fix a few bugs in proxy tests

mojolicious · Jun 16, 2016 · 6d5cf25 · 6d5cf25
1 parent ddae356
commit 6d5cf25
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/lib/Mojo/UserAgent.pm b/lib/Mojo/UserAgent.pm
@@ -126,6 +126,7 @@ sub _connect_proxy {
       my ($self, $tx) = @_;
 
       # CONNECT failed
+      $old->previous($tx)->req->via_proxy(0);
       my $id = $tx->connection;
       if ($tx->error || !$tx->res->is_status_class(200) || !$tx->keep_alive) {
         $old->res->error({message => 'Proxy connection failed'});
@@ -134,7 +135,6 @@ sub _connect_proxy {
       }
 
       # Start real transaction without TLS upgrade
-      $old->req->via_proxy(0);
       return $self->_start($loop, $old->connection($id), $cb)
         unless $tx->req->url->protocol eq 'https';
 

diff --git a/t/mojo/websocket_proxy.t b/t/mojo/websocket_proxy.t
@@ -4,6 +4,7 @@ BEGIN { $ENV{MOJO_REACTOR} = 'Mojo::Reactor::Poll' }
 
 use Test::More;
 use Mojo::IOLoop;
+use Mojo::IOLoop::Server;
 use Mojo::Server::Daemon;
 use Mojo::UserAgent;
 use Mojolicious::Lite;
@@ -35,6 +36,7 @@ $daemon->listen(['http://127.0.0.1'])->start;
 my $port = Mojo::IOLoop->acceptor($daemon->acceptors->[0])->port;
 
 # CONNECT proxy server for testing
+my $dummy = Mojo::IOLoop::Server->generate_port;
 my (%buffer, $connected, $read, $sent);
 my $nf
   = "HTTP/1.1 404 NOT FOUND\x0d\x0a"
@@ -60,7 +62,7 @@ my $id = Mojo::IOLoop->server(
           $buffer{$id}{client} = '';
           if ($buffer =~ /CONNECT (\S+):(\d+)?/) {
             $connected = "$1:$2";
-            my $fail = $2 == $port + 1;
+            my $fail = $2 == $dummy;
 
             # Connection to server
             $buffer{$id}{connection} = Mojo::IOLoop->client(
@@ -193,10 +195,9 @@ ok $sent > 25, 'sent enough';
 
 # Proxy WebSocket with bad target
 $ua->proxy->http("http://127.0.0.1:$proxy");
-my $port2 = $port + 1;
 my ($success, $err);
 $ua->websocket(
-  "ws://127.0.0.1:$port2/test" => sub {
+  "ws://127.0.0.1:$dummy/test" => sub {
     my ($ua, $tx) = @_;
     $success = $tx->success;
     $err     = $tx->error;

diff --git a/t/mojo/websocket_proxy_tls.t b/t/mojo/websocket_proxy_tls.t
@@ -54,6 +54,7 @@ $daemon->listen([$listen])->start;
 my $port = Mojo::IOLoop->acceptor($daemon->acceptors->[0])->port;
 
 # Connect proxy server for testing
+my $dummy = Mojo::IOLoop::Server->generate_port;
 my (%buffer, $connected, $read, $sent);
 my $nf
   = "HTTP/1.1 501 FOO\x0d\x0a"
@@ -79,7 +80,7 @@ my $id = Mojo::IOLoop->server(
           $buffer{$id}{client} = '';
           if ($buffer =~ /CONNECT (\S+):(\d+)?/) {
             $connected = "$1:$2";
-            my $fail = $2 == $port + 1;
+            my $fail = $2 == $dummy;
 
             # Connection to server
             $buffer{$id}{connection} = Mojo::IOLoop->client(
@@ -260,30 +261,31 @@ is $tx->res->body, "https://127.0.0.1:$port/proxy", 'right content';
 
 # Proxy WebSocket with bad target
 $ua->proxy->https("http://127.0.0.1:$proxy");
-my $port2 = $port + 1;
-my ($success, $err);
+my ($success, $err, $leak);
 $ua->websocket(
-  "wss://127.0.0.1:$port2/test" => sub {
+  "wss://127.0.0.1:$dummy/test" => sub {
     my ($ua, $tx) = @_;
     $success = $tx->success;
+    $leak    = !!Mojo::IOLoop->stream($tx->previous->connection);
     $err     = $tx->res->error;
     Mojo::IOLoop->stop;
   }
 );
 Mojo::IOLoop->start;
 ok !$success, 'no success';
+ok !$leak,    'connection has been removed';
 is $err->{message}, 'Proxy connection failed', 'right error';
 
 # Failed TLS handshake through proxy
 $tx = $ua->get("https://127.0.0.1:$close");
-is $err->{message}, 'Proxy connection failed', 'right error';
+like $tx->error->{message}, qr/handshake problems/, 'right error';
 
 # Idle connection through proxy
 $ua->on(
   start => sub { shift->connect_timeout(0.25) if pop->req->method eq 'CONNECT' }
 );
 $tx = $ua->get("https://127.0.0.1:$idle");
-is $err->{message}, 'Proxy connection failed', 'right error';
+is $tx->error->{message}, 'Connect timeout', 'right error';
 $ua->connect_timeout(10);
 
 # Blocking proxy request again