File Coverage

blib/lib/Plack/App/Proxy/WebSocket.pm

Criterion	Covered	Total	%
statement	30	93	32.2
branch	0	32	0.0
condition	0	9	0.0
subroutine	10	13	76.9
pod	2	2	100.0
total	42	149	28.1

package Plack::App::Proxy::WebSocket;
{
  $Plack::App::Proxy::WebSocket::VERSION = '0.03';
}
# ABSTRACT: proxy HTTP and WebSocket connections

use warnings FATAL => 'all';
use strict;

use AnyEvent::Handle;
use AnyEvent::Socket;
use HTTP::Headers;
use HTTP::Request;
use HTTP::Parser::XS qw/parse_http_response HEADERS_AS_HASHREF/;
use Plack::Request;
use URI;

use parent 'Plack::App::Proxy';


sub call {
    my ($self, $env) = @_;
    my $req = Plack::Request->new($env);

    # detect a protocol upgrade handshake or just proxy as usual
    my $upgrade = $req->header('Upgrade') or return $self->SUPER::call($env);

    $env->{'psgi.streaming'} or die "Plack server support for psgi.streaming is required";
    my $client_fh = $env->{'psgix.io'} or die "Plack server support for the psgix.io extension is required";

    my $url = $self->build_url_from_env($env) or return [502, [], ["Bad Gateway"]];
    my $uri = URI->new($url);

    sub {
        my $res = shift;

        # set up an event loop if the server is blocking
        my $cv;
        unless ($env->{'psgi.nonblocking'}) {
            $env->{'psgi.errors'}->print("Plack server support for psgi.nonblocking is highly recommended.\n");
            $cv = AE::cv;
        }

        tcp_connect $uri->host, $uri->port, sub {
            my $server_fh = shift;

            # return 502 if connection to server fails
            unless ($server_fh) {
                $res->([502, [], ["Bad Gateway"]]);
                $cv->send if $cv;
                return;
            }

            my $client = AnyEvent::Handle->new(fh => $client_fh);
            my $server = AnyEvent::Handle->new(fh => $server_fh);

            # forward request from the client
            my $headers = $self->build_headers_from_env($env, $req, $uri);
            $headers->{Upgrade} = $upgrade;
            $headers->{Connection} = 'Upgrade';
            my $hs = HTTP::Request->new('GET', $uri->path, HTTP::Headers->new(%$headers));
            $hs->protocol($req->protocol);
            $server->push_write($hs->as_string);

            my $buffer = "";
            my $writer;

            # buffer the exchange between the client and server
            $client->on_read(sub {
                my $hdl = shift;
                my $buf = delete $hdl->{rbuf};
                $server->push_write($buf);
            });
            $server->on_read(sub {
                my $hdl = shift;
                my $buf = delete $hdl->{rbuf};

                return $writer->write($buf) if $writer;
                $buffer .= $buf;

                my ($ret, $http_version, $status, $message, $headers) =
                    parse_http_response($buffer, HEADERS_AS_HASHREF);
                $server->push_shutdown if $ret == -2;
                return if $ret < 0;

                $headers = [$self->response_headers(HTTP::Headers->new(%$headers))] unless $status == 101;
                $writer = $res->([$status, $headers]);
                $writer->write(substr($buffer, $ret));
                $buffer = undef;
            });

            # shut down the sockets and exit the loop if an error occurs
            $client->on_error(sub {
                $client->destroy;
                $server->push_shutdown;
                $cv->send if $cv;
                $writer->close if $writer;
            });
            $server->on_error(sub {
                $server->destroy;
                # get the client handle's attention
                $client->push_shutdown;
            });
        };

        $cv->recv if $cv;
    };
}


sub build_headers_from_env {
    my ($self, $env, $req, $uri) = @_;

    my $headers = $self->SUPER::build_headers_from_env($env, $req);

    # if x-forwarded-for already existed, append the remote address; the super
    # method fails to maintain a list of mutiple proxies
    if (my $forwarded_for = $env->{HTTP_X_FORWARDED_FOR}) {
        $headers->{'X-Forwarded-For'} = "$forwarded_for, $env->{REMOTE_ADDR}";
    }

    # the super method depends on the user agent to add the host header if it
    # is missing, so set the host if it needs to be set
    if ($uri && !$headers->{'Host'}) {
        $headers->{'Host'} = $uri->host_port;
    }

    $headers->{'X-Forwarded-Proto'} ||= $env->{'psgi.url_scheme'};
    $headers->{'X-Real-IP'} ||= $env->{REMOTE_ADDR};

    $headers;
}

1;

__END__