added basic WebSocket subprotocol negotiation support

mojolicious · May 23, 2013 · af89e25 · af89e25
1 parent 7718f7b
commit af89e25
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 1 deletion.
diff --git a/lib/Mojo/Transaction/WebSocket.pm b/lib/Mojo/Transaction/WebSocket.pm
@@ -2,6 +2,7 @@ package Mojo::Transaction::WebSocket;
 use Mojo::Base 'Mojo::Transaction';
 
 use Config;
+use List::Util 'first';
 use Mojo::JSON;
 use Mojo::Transaction::HTTP;
 use Mojo::Util qw(b64_encode decode encode sha1_bytes xor_encode);
@@ -262,6 +263,20 @@ sub server_write {
   return delete $self->{write} // '';
 }
 
+sub subprotocol {
+  my $self = shift;
+
+  return $self->res->headers->sec_websocket_protocol if $self->masked;
+
+  return undef unless my $sub = $self->req->headers->sec_websocket_protocol;
+  for my $proto (split /,\s*/, $sub) {
+    next unless my $match = first { $proto eq $_ } @_;
+    $self->res->headers->sec_websocket_protocol($match);
+    return $match;
+  }
+  return undef;
+}
+
 sub _challenge { b64_encode(sha1_bytes(($_[0] || '') . GUID), '') }
 
 sub _message {
@@ -649,6 +664,13 @@ Read data server-side, used to implement web servers.
 
 Write data server-side, used to implement web servers.
 
+=head2 subprotocol
+
+  my $proto = $ws->subprotocol;
+  my $proto = $ws->subprotocol('v1.foo', 'v2.foo');
+
+Negotiate WebSocket subprotocol and return C<undef> if negotiation failed.
+
 =head1 DEBUGGING
 
 You can set the MOJO_WEBSOCKET_DEBUG environment variable to get some advanced

diff --git a/lib/Test/Mojo.pm b/lib/Test/Mojo.pm
@@ -255,6 +255,12 @@ sub status_isnt {
   return $self->_test('isnt', $self->tx->res->code, $status, $desc);
 }
 
+sub subprotocol_is {
+  my ($self, $proto, $desc) = @_;
+  return $self->_test('is', $self->tx->subprotocol,
+    $proto, $desc || 'exact match for subprotocol');
+}
+
 sub text_is {
   my ($self, $selector, $value, $desc) = @_;
   $desc ||= encode 'UTF-8', qq{exact match for selector "$selector"};
@@ -843,6 +849,13 @@ Check response status for exact match.
 
 Opposite of C<status_is>.
 
+=head2 subprotocol_is
+
+  $t = $t->subprotocol_is('v1.proto');
+  $t = $t->subprotocol_is('v1.proto', 'right protocol');
+
+Check WebSocket subprotocol for exact match.
+
 =head2 text_is
 
   $t = $t->text_is('div.foo[x=y]' => 'Hello!');

diff --git a/t/mojolicious/websocket_lite_app.t b/t/mojolicious/websocket_lite_app.t
@@ -81,6 +81,13 @@ websocket '/once' => sub {
   );
 };
 
+websocket '/negotiate' => sub {
+  my $self = shift;
+  return $self->rendered(426)
+    unless my $proto = $self->tx->subprotocol('v2.echo', 'echo');
+  $self->on(message => sub { shift->send($proto . ': ' . shift) });
+};
+
 under '/nested';
 
 websocket sub {
@@ -112,7 +119,7 @@ $t->websocket_ok('/echo')->send_ok('hello again')
 
 # Custom header and protocols
 $t->websocket_ok('/echo' => {DNT => 1} => ['foo', 'bar', 'baz'])
-  ->header_is('Sec-WebSocket-Protocol' => 'foo')->send_ok('hello')
+  ->subprotocol_is('foo')->send_ok('hello')
   ->message_ok->message_is('echo: hello')->finish_ok;
 is $t->tx->req->headers->dnt, 1, 'right "DNT" value';
 is $t->tx->req->headers->sec_websocket_protocol, 'foo, bar, baz',
@@ -223,6 +230,20 @@ $t->websocket_ok('/once')->send_ok('hello')
   ->send_ok('hello')->message_ok->message_is('ONE: hello')->send_ok('hello')
   ->message_ok->message_is('ONE: hello')->finish_ok;
 
+# Failed subprotocol negotiation
+my $tx = $t->ua->build_websocket_tx('/negotiate');
+$t->request_ok($tx)->status_is(426);
+$tx = $t->ua->build_websocket_tx('/negotiate' => ['v3.echo', 'v4.echo']);
+$t->request_ok($tx)->status_is(426);
+
+# Successful subprotocol negotiation
+$t->websocket_ok('/negotiate' => ['v3.echo', 'v2.echo'])
+  ->subprotocol_is('v2.echo')->send_ok('hello')
+  ->message_ok->message_is('v2.echo: hello')->finish_ok;
+$t->websocket_ok('/negotiate' => ['echo', 'v2.echo'])->subprotocol_is('echo')
+  ->send_ok('hello again')->message_ok->message_is('echo: hello again')
+  ->finish_ok;
+
 # Nested WebSocket
 $t->websocket_ok('/nested')->send_ok('hello')
   ->message_ok->message_is('nested echo: hello')->finished_ok(1000);