added support for reason strings to Mojo::Transaction::WebSocket

mojolicious · Apr 11, 2013 · 0772343 · 0772343
1 parent fbe03ac
commit 0772343
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 21 deletions.
diff --git a/lib/Mojo/Transaction/WebSocket.pm b/lib/Mojo/Transaction/WebSocket.pm
@@ -107,10 +107,14 @@ sub client_write { shift->server_write(@_) }
 sub connection { shift->handshake->connection }
 
 sub finish {
-  my ($self, $code) = @_;
-  my $payload = $code ? pack('n', $code) : '';
-  $self->{status} //= $code ? $code : 1005;
+  my $self = shift;
+
+  my $close = $self->{close} = [@_];
+  my $payload = $close->[0] ? pack('n', $close->[0]) : '';
+  $payload .= encode 'UTF-8', $close->[1] if defined $close->[1];
+  $close->[0] //= 1005;
   $self->send([1, 0, 0, 0, CLOSE, $payload])->{finished} = 1;
+
   return $self;
 }
 
@@ -217,7 +221,7 @@ sub send {
 sub server_close {
   my $self = shift;
   $self->{state} = 'finished';
-  return $self->emit(finish => $self->{status} // 1006);
+  return $self->emit(finish => $self->{close} ? (@{$self->{close}}) : 1006);
 }
 
 sub server_handshake {
@@ -269,10 +273,11 @@ sub _message {
 
   # Close
   if ($op == CLOSE) {
-    my $code;
-    $code = unpack 'n', substr($frame->[5], 0, 2) if length $frame->[5] >= 2;
-    $self->{status} = $code // 1005;
-    return $self->finish($code);
+    return $self->finish(1005) unless length $frame->[5] >= 2;
+    my $close = $self->{close} = [];
+    $close->[0] = unpack 'n', substr($frame->[5], 0, 2, '');
+    $close->[1] = decode 'UTF-8', $frame->[5];
+    return $self->finish(@$close);
   }
 
   # Append chunk and check message size
@@ -358,7 +363,7 @@ Emitted once all data has been sent.
 =head2 finish
 
   $ws->on(finish => sub {
-    my ($ws, $code) = @_;
+    my ($ws, $code, $reason) = @_;
     ...
   });
 
@@ -513,6 +518,7 @@ Connection identifier or socket.
 
   $ws = $ws->finish;
   $ws = $ws->finish(1000);
+  $ws = $ws->finish(1003, 'What was that?');
 
 Finish the WebSocket connection gracefully.
 

diff --git a/lib/Mojolicious/Controller.pm b/lib/Mojolicious/Controller.pm
@@ -64,20 +64,20 @@ sub cookie {
 }
 
 sub finish {
-  my ($self, $chunk) = @_;
+  my $self = shift;
 
   # WebSocket
   my $tx = $self->tx;
-  $tx->finish($chunk) and return $self if $tx->is_websocket;
+  $tx->finish(@_) and return $self if $tx->is_websocket;
 
   # Chunked stream
   if ($tx->res->is_chunked) {
-    $self->write_chunk($chunk) if defined $chunk;
+    $self->write_chunk(@_) if @_;
     return $self->write_chunk('');
   }
 
   # Normal stream
-  $self->write($chunk) if defined $chunk;
+  $self->write(@_) if @_;
   return $self->write('');
 }
 
@@ -553,6 +553,7 @@ Access request cookie values and create new response cookies.
 
   $c = $c->finish;
   $c = $c->finish(1000);
+  $c = $c->finish(1003, 'What was that?');
   $c = $c->finish('Bye!');
 
 Gracefully end WebSocket connection or long poll stream.

diff --git a/lib/Test/Mojo.pm b/lib/Test/Mojo.pm
@@ -111,7 +111,7 @@ sub finish_ok {
 sub finished_ok {
   my ($self, $code) = @_;
   Mojo::IOLoop->one_tick while !$self->{finished};
-  return $self->_test('is', $self->{finished}, $code, 'finished websocket');
+  return $self->_test('is', $self->{finished}[0], $code, 'finished websocket');
 }
 
 sub get_ok  { shift->_request_ok(get  => @_) }
@@ -305,12 +305,12 @@ sub websocket_ok {
 
   # Establish WebSocket connection
   $self->{messages} = [];
-  $self->{finished} = 0;
+  $self->{finished} = undef;
   $self->ua->websocket(
     $url => @_ => sub {
       my ($ua, $tx) = @_;
       $self->tx($tx);
-      $tx->on(finish => sub { $self->{finished} = pop });
+      $tx->on(finish => sub { shift; $self->{finished} = [@_] });
       $tx->on(binary => sub { push @{$self->{messages}}, [binary => pop] });
       $tx->on(text   => sub { push @{$self->{messages}}, [text   => pop] });
       Mojo::IOLoop->stop;
@@ -606,6 +606,7 @@ Opposite of C<element_exists>.
 
   $t = $t->finish_ok;
   $t = $t->finish_ok(1000);
+  $t = $t->finish_ok(1003, 'What was that?');
 
 Finish WebSocket connection gracefully.
 

diff --git a/t/mojo/websocket.t b/t/mojo/websocket.t
@@ -59,7 +59,7 @@ websocket '/socket' => sub {
     $self->req->headers->host => sub {
       my $self = shift;
       $self->send(Mojo::IOLoop->stream($self->tx->connection)->timeout);
-      $self->finish(1000);
+      $self->finish(1000, 'I ♥ Mojolicious!');
     }
   );
 };
@@ -207,15 +207,16 @@ my $sock = IO::Socket::INET->new(PeerAddr => '127.0.0.1', PeerPort => $port);
 $sock->blocking(0);
 $tx->connection($sock);
 $result = '';
-my ($local, $early, $status);
+my ($local, $early, $status, $msg);
 $ua->start(
   $tx => sub {
     my ($ua, $tx) = @_;
     $early = $finished;
     $tx->on(
       finish => sub {
-        my ($tx, $code) = @_;
+        my ($tx, $code, $reason) = @_;
         $status = $code;
+        $msg    = $reason;
         Mojo::IOLoop->stop;
       }
     );
@@ -233,6 +234,7 @@ Mojo::IOLoop->start;
 is $finished, 1,    'finish event has been emitted';
 is $early,    1,    'finish event has been emitted at the right time';
 is $status,   1000, 'right status';
+is $msg, 'I ♥ Mojolicious!', 'right message';
 ok $result =~ /^lalala(\d+)$/, 'right result';
 is $1, 15, 'right timeout';
 ok $local, 'local port';
@@ -467,8 +469,8 @@ Mojo::IOLoop->start;
 is $result, 'foo bar', 'right result';
 
 # Dies
-($finished, $code) = ();
-my ($websocket, $msg);
+($finished, $code, $msg) = ();
+my $websocket;
 $ua->websocket(
   '/dead' => sub {
     my ($ua, $tx) = @_;