fixed and readded support for permessage-deflate WebSocket compression

mojolicious · Feb 2, 2014 · aeee9e8 · aeee9e8
1 parent e98254f
commit aeee9e8
Show file tree

Hide file tree

Showing 6 changed files with 104 additions and 6 deletions.
diff --git a/Changes b/Changes
@@ -1,4 +1,8 @@
 
+4.75  2014-02-02
+  - Fixed and readded support for permessage-deflate WebSocket compression.
+    (Mikey, sri)
+
 4.74  2014-02-02
   - Added all_contents method to Mojo::DOM.
   - Removed support for permessage-deflate WebSocket compression, since there

diff --git a/lib/Mojo/Transaction/WebSocket.pm b/lib/Mojo/Transaction/WebSocket.pm
@@ -1,6 +1,7 @@
 package Mojo::Transaction::WebSocket;
 use Mojo::Base 'Mojo::Transaction';
 
+use Compress::Raw::Zlib 'Z_SYNC_FLUSH';
 use Config;
 use Mojo::JSON;
 use Mojo::Transaction::HTTP;
@@ -25,7 +26,7 @@ use constant {
   PONG         => 10
 };
 
-has 'masked';
+has [qw(compressed masked)];
 has handshake => sub { Mojo::Transaction::HTTP->new };
 has max_websocket_size => sub { $ENV{MOJO_MAX_WEBSOCKET_SIZE} || 262144 };
 
@@ -88,13 +89,29 @@ sub build_message {
   if (exists $frame->{text}) { $frame = [1, 0, 0, 0, TEXT, $frame->{text}] }
   else { $frame = [1, 0, 0, 0, BINARY, $frame->{binary}] }
 
+  # "permessage-deflate" extension
+  return $self->build_frame(@$frame) unless $self->compressed;
+  my $deflate = $self->{deflate} ||= Compress::Raw::Zlib::Deflate->new(
+    AppendOutput => 1,
+    MemLevel     => 8,
+    WindowBits   => -15
+  );
+  $deflate->deflate($frame->[5], my $out);
+  $deflate->flush($out, Z_SYNC_FLUSH);
+  @$frame[1, 5] = (1, substr($out, 0, length($out) - 4));
   return $self->build_frame(@$frame);
 }
 
 sub client_challenge {
   my $self = shift;
+
+  # "permessage-deflate" extension
+  my $headers = $self->res->headers;
+  $self->compressed(1)
+    if ($headers->sec_websocket_extensions // '') =~ /permessage-deflate/;
+
   return _challenge($self->req->headers->sec_websocket_key) eq
-    $self->res->headers->sec_websocket_accept;
+    $headers->sec_websocket_accept;
 }
 
 sub client_handshake {
@@ -239,6 +256,11 @@ sub server_handshake {
     and $res_headers->sec_websocket_protocol($1);
   $res_headers->sec_websocket_accept(
     _challenge($req_headers->sec_websocket_key));
+
+  # "permessage-deflate" extension
+  $self->compressed(1)
+    and $res_headers->sec_websocket_extensions('permessage-deflate')
+    if ($req_headers->sec_websocket_extensions // '') =~ /permessage-deflate/;
 }
 
 sub server_read {
@@ -291,7 +313,19 @@ sub _message {
   # No FIN bit (Continuation)
   return unless $frame->[0];
 
+  # "permessage-deflate" extension (handshake and RSV1)
   my $msg = delete $self->{message};
+  if ($self->compressed && $frame->[1]) {
+    my $inflate = $self->{inflate} ||= Compress::Raw::Zlib::Inflate->new(
+      Bufsize     => $max,
+      LimitOutput => 1,
+      WindowBits  => -15
+    );
+    $inflate->inflate(($msg .= "\x00\x00\xff\xff"), my $out);
+    return $self->finish(1009) if length $msg;
+    $msg = $out;
+  }
+
   $self->emit(json => Mojo::JSON->new->decode($msg))
     if $self->has_subscribers('json');
   $op = delete $self->{op};
@@ -443,6 +477,13 @@ Emitted when a complete WebSocket text message has been received.
 L<Mojo::Transaction::WebSocket> inherits all attributes from
 L<Mojo::Transaction> and implements the following new ones.
 
+=head2 compressed
+
+  my $bool = $ws->compressed;
+  $ws      = $ws->compressed($bool);
+
+Compress messages with C<permessage-deflate> extension.
+
 =head2 handshake
 
   my $handshake = $ws->handshake;

diff --git a/lib/Mojolicious.pm b/lib/Mojolicious.pm
@@ -43,7 +43,7 @@ has types     => sub { Mojolicious::Types->new };
 has validator => sub { Mojolicious::Validator->new };
 
 our $CODENAME = 'Top Hat';
-our $VERSION  = '4.74';
+our $VERSION  = '4.75';
 
 sub AUTOLOAD {
   my $self = shift;
@@ -877,6 +877,8 @@ Maksym Komar
 
 Maxim Vuets
 
+Michael Gregorowicz
+
 Michael Harris
 
 Mike Magowan

diff --git a/lib/Test/Mojo.pm b/lib/Test/Mojo.pm
@@ -903,6 +903,13 @@ Opposite of L</"text_like">.
 Open a WebSocket connection with transparent handshake, takes the same
 arguments as L<Mojo::UserAgent/"websocket">, except for the callback.
 
+  # WebSocket with permessage-deflate compression
+  $t->websocket('/x' => {'Sec-WebSocket-Extensions' => 'permessage-deflate'})
+    ->send_ok('y')
+    ->message_ok
+    ->message_is('z')
+    ->finish_ok;
+
 =head1 SEE ALSO
 
 L<Mojolicious>, L<Mojolicious::Guides>, L<http://mojolicio.us>.

diff --git a/t/mojo/websocket_frames.t b/t/mojo/websocket_frames.t
@@ -206,15 +206,18 @@ is $frame->[5], '', 'no payload';
 isnt(Mojo::Transaction::WebSocket->new->build_frame(1, 0, 0, 0, 2, ''),
   $bytes, 'frames are not equal');
 
-# Binary message roundtrip
-$ws    = Mojo::Transaction::WebSocket->new;
+# Compressed binary message roundtrip
+$ws = Mojo::Transaction::WebSocket->new(compressed => 1);
 $bytes = $ws->build_message({binary => 'just works'});
 $frame = $ws->parse_frame(\($dummy = $bytes));
 is $frame->[0], 1, 'fin flag is set';
-is $frame->[1], 0, 'rsv1 flag is not set';
+is $frame->[1], 1, 'rsv1 flag is set';
 is $frame->[2], 0, 'rsv2 flag is not set';
 is $frame->[3], 0, 'rsv3 flag is not set';
 is $frame->[4], 2, 'binary frame';
 ok $frame->[5], 'has payload';
+isnt(
+  Mojo::Transaction::WebSocket->new->build_message({binary => 'just works'}),
+  $bytes, 'messages are not equal');
 
 done_testing();
diff --git a/t/mojolicious/websocket_lite_app.t b/t/mojolicious/websocket_lite_app.t
@@ -25,6 +25,13 @@ websocket '/echo' => sub {
 
 get '/echo' => {text => 'plain echo!'};
 
+websocket '/no_compression' => sub {
+  my $self = shift;
+  $self->tx->compressed(0);
+  $self->res->headers->remove('Sec-WebSocket-Extensions');
+  $self->on(binary => sub { shift->send({binary => shift}) });
+};
+
 websocket '/json' => sub {
   my $self = shift;
   $self->on(
@@ -150,6 +157,40 @@ $t->websocket_ok('/echo')->send_ok([0, 0, 0, 0, 2, 'c' x 100000])
 # Plain alternative
 $t->get_ok('/echo')->status_is(200)->content_is('plain echo!');
 
+# Compression denied by the server
+$t->websocket_ok(
+  '/no_compression' => {'Sec-WebSocket-Extensions' => 'permessage-deflate'});
+is $t->tx->req->headers->sec_websocket_extensions, 'permessage-deflate',
+  'right "Sec-WebSocket-Extensions" value';
+ok !$t->tx->compressed, 'WebSocket has no compression';
+$t->send_ok({binary => 'a' x 500})
+  ->message_ok->message_is({binary => 'a' x 500})->finish_ok;
+
+# Compressed message ("permessage-deflate")
+$t->websocket_ok(
+  '/echo' => {'Sec-WebSocket-Extensions' => 'permessage-deflate'});
+ok $t->tx->compressed, 'WebSocket has compression';
+$t->send_ok({binary => 'a' x 50000})
+  ->header_is('Sec-WebSocket-Extensions' => 'permessage-deflate');
+is $t->tx->req->headers->sec_websocket_extensions, 'permessage-deflate',
+  'right "Sec-WebSocket-Extensions" value';
+my $payload;
+$t->tx->once(
+  frame => sub {
+    my ($tx, $frame) = @_;
+    $payload = $frame->[5];
+  }
+);
+$t->message_ok->message_is({binary => 'a' x 50000});
+ok length $payload < 262145, 'message has been compressed';
+$t->finish_ok->finished_ok(1005);
+
+# Compressed message exceeding the limit when uncompressed
+$t->websocket_ok(
+  '/echo' => {'Sec-WebSocket-Extensions' => 'permessage-deflate'})
+  ->header_is('Sec-WebSocket-Extensions' => 'permessage-deflate')
+  ->send_ok({binary => 'a' x 1000000})->finished_ok(1009);
+
 # JSON roundtrips
 $t->websocket_ok('/json')->send_ok({json => {test => 23, snowman => '☃'}})
   ->message_ok->json_message_is('' => {test => 24, snowman => '☃'})