wscat (WebSocket netcat) implemented in Perl

wscat.pl

#!/usr/bin/env perl

use strict;
use warnings;
use autodie;

use Digest::SHA ();
use IO::Select ();
use IO::Socket::INET ();
use MIME::Base64 ();
use Socket ();
use URI ();

use Protocol::WebSocket::Frame ();

use constant WS_GUID => '258EAFA5-E914-47DA-95CA-C5AB0DC85B11';

use constant MAX_CHUNK_SIZE => 64000;

use constant CRLF => "\x0d\x0a";

use constant ERROR_SIGS => qw( HUP QUIT ABRT USR1 USR2 SEGV PIPE ALRM TERM );

run( @ARGV ) if !caller;

sub run {
    my ($uri) = @_;

    my $uri_obj = URI->new($uri);

    my $uri_scheme = $uri_obj->scheme();

    if (!$uri_scheme) {
        die "Need a URI!\n";
    }

    if ($uri_obj->scheme() !~ m<\Awss?\z>) {
        die sprintf "Invalid schema: “%s” ($uri)\n", $uri_obj->scheme();
    }

    #Reparse this as an HTTP URI since URI.pm doesn’t historically
    #parse WebSocket URIs.
    $uri =~ s<\Aws><http>;
    $uri_obj = URI->new($uri);

    my $inet;

    if ($uri_scheme eq 'ws') {
        my $iaddr = Socket::inet_aton($uri_obj->host());

        my $port = $uri_obj->port() || 80;
        my $paddr = Socket::pack_sockaddr_in( $port, $iaddr );

        socket( $inet, Socket::PF_INET(), Socket::SOCK_STREAM(), 0 );
        connect( $inet, $paddr );
    }
    else {
        require IO::Socket::SSL;

        $inet = IO::Socket::SSL->new(
            PeerHost => $uri_obj->host(),
            PeerPort => $uri_obj->port() || 443,
            SSL_hostname => $uri_obj->host(),
        );

        die "IO::Socket::SSL: [$!][$@]\n" if !$inet;
    }

    my $get_arg = $uri_obj->path();
    if (!length $get_arg) {
        $get_arg = '/';
    }
    if (length $uri_obj->query()) {
        $get_arg .= '.' . $get_arg->query();
    }

    #TODO: Rewrite without IO::Select?
    my $s = IO::Select->new();

    $inet->blocking(0);

    $s->add($inet);

    my $buf_sr = _handshake_as_client( $s, $inet, $uri_obj->host(), $get_arg );

    _mux_after_handshake( \*STDIN, \*STDOUT, $s, $inet, $$buf_sr );

    exit 1;
}

sub _handshake_as_client {
    my ($from_remote_s, $to_remote, $host, $get_arg) = @_;

    my $key = join q<>, map { _two_random_bytes() } 1 .. 8;

    my $key_b64 = MIME::Base64::encode_base64($key);
    chomp $key_b64;

    my $hdr = join(
        CRLF,
        "GET $get_arg HTTP/1.1",
        "Host: $host",
        'Upgrade: websocket',
        'Connection: Upgrade',
        'Sec-WebSocket-Version: 13',
        "Sec-WebSocket-Key: $key_b64",
        q<>,
        q<>,
    );

    #Write out the client handshake.
    syswrite( $to_remote, $hdr );

    my $handshake_ok;

    my $buf;

    my $got_first_line;
    my $got_upgrade;
    my $got_connection;
    my $got_accept;

    #Read the server handshake.
  HANDSHAKE:
    while (!$handshake_ok) {
        for my $fh ($from_remote_s->can_read(1)) {
            sysread $fh, my $buf, MAX_CHUNK_SIZE;

            my $lf1;

            while ( -1 != ($lf1 = index($buf, "\x0a") ) ) {
                my $line = substr( $buf, 0, 1 + $lf1, q<> );

                if ( $line eq "\x0d\x0a" || $line eq "\x0a" ) {
                    my @missing = (
                        ( $got_first_line ? () : 'status line' ),
                        ( $got_upgrade ? () : 'Upgrade' ),
                        ( $got_connection ? () : 'Connection' ),
                        ( $got_accept ? () : 'Sec-WebSocket-Accept' ),
                    );

                    if (@missing) {
                        die( "Headers missing: " . join ', ', @missing );
                    }

                    last HANDSHAKE;
                }

                if (!$got_first_line) {
                    $got_first_line = 1;

                    if ($line !~ m<\AHTTP/1.1 101 >) {
                        _chomp_crlf($line);
                        die "Unfamiliar first line: “$line”\n";
                    }
                }

                my ($key, $val) = split m<\s*:\s*>, $line, 2;

                if ($key =~ m<\Aupgrade\z>i) {
                    _chomp_crlf($val);
                    my @vals = split m<\s*,\s*>, $val;
                    if (!grep { m<\Awebsocket\z>i } @vals) {
                        die "“$key” ($val) must contain “websocket”!\n";
                    }
                    $got_upgrade = 1;
                }
                elsif ($key =~ m<\Aconnection\z>i) {
                    _chomp_crlf($val);
                    my @vals = split m<\s*,\s*>, $val;
                    if (!grep { m<\Aupgrade\z>i } @vals) {
                        die "“$key” ($val) must contain “upgrade”!\n";
                    }
                    $got_connection = 1;
                }
                elsif ($key =~ m<\Asec-websocket-accept\z>i) {
                    my $accept_val = Digest::SHA::sha1_base64( $key_b64 . WS_GUID() );
                    _pad_b64($accept_val);

                    _chomp_crlf($val);
                    if ( $val ne $accept_val ) {
                        die "“$key” must be “$accept_val”, not “$val”! (Key was “$key_b64”)\n";
                    }
                    $got_accept = 1;
                }
            }
        }
    }

    return \$buf;
}

sub _mux_after_handshake {
    my ($from_caller, $to_caller, $from_remote_s, $to_remote, $buf) = @_;

    #Funnel from us to them
    my $pid = fork;
    if (!$pid) {
        die "fork(): $!" if !defined $pid;

        eval {
            for my $sig (ERROR_SIGS()) {
                $SIG{$sig} = sub {
                    my $frame = Protocol::WebSocket::Frame->new(
                        buffer => pack('n', 1011),
                        type => 'close',
                    );

                    syswrite( $to_remote, $frame->to_bytes() );

                    $SIG{$_[0]} = 'DEFAULT';
                    kill $_[0], $$;
                };
            }

            $SIG{'INT'} = sub {
                my $frame = Protocol::WebSocket::Frame->new(
                    buffer => pack('n', 1000),
                    type => 'close',
                );

                syswrite( $to_remote, $frame->to_bytes() );

                $SIG{$_[0]} = 'DEFAULT';
                kill $_[0], $$;
            };

            while (my $line = readline $from_caller) {
                my $frame = Protocol::WebSocket::Frame->new(
                    buffer => $line,
                    type => 'binary',
                );

                syswrite( $to_remote, $frame->to_bytes() );
            }
        };

        exit 1; #we get here from … ??
    };

    my $frame = Protocol::WebSocket::Frame->new(
        buffer => $buf,
        type => 'binary',
    );

    eval {

        #Funnel from them to us
        while (1) {
            while ( my $next = $frame->next_bytes() ) {
                print {$to_caller} $next;
            }

            my @ready = $from_remote_s->can_read(1);

            for my $fh (@ready) {
                while ( CORE::sysread $fh, my $buf, MAX_CHUNK_SIZE ) {
                    $frame->append($buf);
                }
                die "sysread(): $!" if !$!{'EAGAIN'};
            }
        }

        die "Never get here!";
    };

    kill 'TERM', $pid;
}

sub _two_random_bytes {
    my $rnd = int rand 65536;
    return pack 's', $rnd;
}

sub _pad_b64 {
    if (length($_[0]) % 4) {
        $_[0] .= '=' x (4 - (length($_[0]) % 4));
    }
}

sub _chomp_crlf { $_[0] =~ s<\x0d?\x0a\z><> }