added support for negotiating WebSocket subprotocols

mojolicious · Dec 13, 2015 · 1dfaf12 · 1dfaf12
1 parent 62d0d71
commit 1dfaf12
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 1 deletion.
diff --git a/Changes b/Changes
@@ -1,5 +1,6 @@
 
 6.37  2015-12-13
+  - Added protocol and with_protocols methods to Mojo::Transaction::WebSocket.
   - Fixed a few reference encoding bugs in Mojo::JSON.
   - Fixed a few form generation bugs in Mojo::UserAgent::Transactor.
 

diff --git a/lib/Mojo/Transaction/WebSocket.pm b/lib/Mojo/Transaction/WebSocket.pm
@@ -3,9 +3,10 @@ use Mojo::Base 'Mojo::Transaction';
 
 use Compress::Raw::Zlib 'Z_SYNC_FLUSH';
 use Config;
+use List::Util 'first';
 use Mojo::JSON qw(encode_json j);
 use Mojo::Transaction::HTTP;
-use Mojo::Util qw(b64_encode decode dumper encode sha1_bytes xor_encode);
+use Mojo::Util qw(b64_encode decode dumper encode sha1_bytes trim xor_encode);
 
 use constant DEBUG => $ENV{MOJO_WEBSOCKET_DEBUG} || 0;
 
@@ -209,6 +210,8 @@ sub parse_frame {
   return [$fin, $rsv1, $rsv2, $rsv3, $op, $payload];
 }
 
+sub protocol { shift->res->headers->sec_websocket_protocol }
+
 sub remote_address { shift->handshake->remote_address }
 sub remote_port    { shift->handshake->remote_port }
 sub req            { shift->handshake->req }
@@ -276,6 +279,17 @@ sub with_compression {
     =~ /permessage-deflate/;
 }
 
+sub with_protocols {
+  my $self = shift;
+
+  my %protos = map { trim($_) => 1 } split ',',
+    $self->req->headers->sec_websocket_protocol;
+  return undef unless my $proto = first { $protos{$_} } @_;
+
+  $self->res->headers->sec_websocket_protocol($proto);
+  return $proto;
+}
+
 sub _challenge { b64_encode(sha1_bytes(($_[0] || '') . GUID), '') }
 
 sub _message {
@@ -631,6 +645,12 @@ Parse WebSocket frame.
   say "Opcode: $frame->[4]";
   say "Payload: $frame->[5]";
 
+=head2 protocol
+
+  my $proto = $ws->protocol;
+
+Return negotiated subprotocol or C<undef>.
+
 =head2 remote_address
 
   my $address = $ws->remote_address;
@@ -717,6 +737,12 @@ L<Mojo::Server::Daemon>.
 
 Negotiate C<permessage-deflate> extension for this WebSocket connection.
 
+=head2 with_protocols
+
+  my $proto = $ws->with_protocols('v2.proto', 'v1.proto');
+
+Negotiate subprotocol for this WebSocket connection.
+
 =head1 DEBUGGING
 
 You can set the C<MOJO_WEBSOCKET_DEBUG> environment variable to get some

diff --git a/t/mojolicious/websocket_lite_app.t b/t/mojolicious/websocket_lite_app.t
@@ -27,6 +27,12 @@ websocket '/no_compression' => sub {
   $c->on(binary => sub { shift->send({binary => shift}) });
 };
 
+websocket '/protocols' => sub {
+  my $c = shift;
+  $c->send($c->tx->with_protocols('foo', 'bar', 'baz') // 'none');
+  $c->send($c->tx->protocol // 'none');
+};
+
 websocket '/json' => sub {
   my $c = shift;
   $c->on(
@@ -208,6 +214,24 @@ $t->websocket_ok(
   ->send_ok({binary => $huge})->message_ok->message_is({binary => $huge})
   ->finish_ok;
 
+# Protocol negotiation
+$t->websocket_ok('/protocols' => ['bar'])->message_ok->message_is('bar')
+  ->message_ok->message_is('bar')->finish_ok;
+is $t->tx->protocol, 'bar', 'right protocol';
+is $t->tx->res->headers->sec_websocket_protocol, 'bar',
+  'right "Sec-WebSocket-Protocol" value';
+$t->websocket_ok('/protocols' => ['baz', 'bar', 'foo'])
+  ->message_ok->message_is('foo')->message_ok->message_is('foo')->finish_ok;
+is $t->tx->protocol, 'foo', 'right protocol';
+is $t->tx->res->headers->sec_websocket_protocol, 'foo',
+  'right "Sec-WebSocket-Protocol" value';
+$t->websocket_ok('/protocols' => [''])->message_ok->message_is('none')
+  ->message_ok->message_is('none')->finish_ok;
+is $t->tx->protocol, undef, 'no protocol';
+$t->websocket_ok('/protocols' => ['', '', ''])->message_ok->message_is('none')
+  ->message_ok->message_is('none')->finish_ok;
+is $t->tx->protocol, undef, 'no protocol';
+
 # JSON roundtrips (with a lot of different tests)
 $t->websocket_ok('/json')->send_ok({json => {test => 23, snowman => '☃'}})
   ->message_ok->json_message_is('' => {test => 24, snowman => '☃'})