added support for permessage-deflate WebSocket compression (closes #566)

mojolicious · Nov 28, 2013 · 83ebad4 · 83ebad4
1 parent 819a6c8
commit 83ebad4
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 19 deletions.
diff --git a/Changes b/Changes
@@ -1,5 +1,7 @@
 
-4.59  2013-11-27
+4.59  2013-11-28
+  - Added support for permessage-deflate WebSocket compression.
+  - Added compressed attribute to Mojo::Transaction::WebSocket.
   - Relicensed all artwork to CC-SA version 4.0.
 
 4.58  2013-11-19

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;
@@ -26,7 +27,7 @@ use constant {
 };
 
 has handshake => sub { Mojo::Transaction::HTTP->new };
-has 'masked';
+has [qw(compressed masked)];
 has max_websocket_size => sub { $ENV{MOJO_MAX_WEBSOCKET_SIZE} || 262144 };
 
 sub new {
@@ -83,8 +84,13 @@ sub build_frame {
 
 sub client_challenge {
   my $self = shift;
+
+  # "permessage-deflate" extension
+  my $headers = $self->res->headers;
+  $self->_deflate($headers->sec_websocket_extensions);
+
   return _challenge($self->req->headers->sec_websocket_key) eq
-    $self->res->headers->sec_websocket_accept;
+    $headers->sec_websocket_accept;
 }
 
 sub client_handshake {
@@ -94,6 +100,8 @@ sub client_handshake {
   $headers->upgrade('websocket')      unless $headers->upgrade;
   $headers->connection('Upgrade')     unless $headers->connection;
   $headers->sec_websocket_version(13) unless $headers->sec_websocket_version;
+  $headers->sec_websocket_extensions('permessage-deflate')
+    unless $headers->sec_websocket_extensions;
 
   # Generate 16 byte WebSocket challenge
   my $challenge = b64_encode sprintf('%16u', int(rand 9 x 16)), '';
@@ -207,12 +215,12 @@ sub send {
     # Binary or raw text
     $frame
       = exists $frame->{text}
-      ? [1, 0, 0, 0, TEXT, $frame->{text}]
-      : [1, 0, 0, 0, BINARY, $frame->{binary}];
+      ? $self->_compress([1, 0, 0, 0, TEXT,   $frame->{text}])
+      : $self->_compress([1, 0, 0, 0, BINARY, $frame->{binary}]);
   }
 
   # Text
-  $frame = [1, 0, 0, 0, TEXT, encode('UTF-8', $frame)]
+  $frame = $self->_compress([1, 0, 0, 0, TEXT, encode('UTF-8', $frame)])
     if ref $frame ne 'ARRAY';
 
   $self->once(drain => $cb) if $cb;
@@ -238,6 +246,10 @@ sub server_handshake {
     and $res_headers->sec_websocket_protocol($1);
   $res_headers->sec_websocket_accept(
     _challenge($req_headers->sec_websocket_key));
+
+  # "permessage-deflate" extension
+  $res_headers->sec_websocket_extensions('permessage-deflate')
+    if $self->_deflate($req_headers->sec_websocket_extensions);
 }
 
 sub server_read {
@@ -264,6 +276,25 @@ sub server_write {
 
 sub _challenge { b64_encode(sha1_bytes(($_[0] || '') . GUID), '') }
 
+sub _compress {
+  my ($self, $frame) = @_;
+
+  # No compression negotiated
+  return $frame unless $self->compressed;
+
+  # "permessage-deflate" extension
+  $frame->[1] = 1;
+  my $deflate = $self->{deflate}
+    ||= Compress::Raw::Zlib::Deflate->new(WindowBits => -15, MemLevel => 8);
+  $deflate->deflate(\$frame->[5], my $out);
+  $deflate->flush($out, Z_SYNC_FLUSH);
+  $frame->[5] = substr $out, 0, length($out) - 4;
+
+  return $frame;
+}
+
+sub _deflate { ($_[1] // '') =~ /permessage-deflate/i && $_[0]->compressed(1) }
+
 sub _message {
   my ($self, $frame) = @_;
 
@@ -290,8 +321,15 @@ sub _message {
   # No FIN bit (Continuation)
   return unless $frame->[0];
 
-  # Whole message
+  # "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(WindowBits => -15);
+    $inflate->inflate(\($msg .= "\x00\x00\xff\xff"), my $out);
+    $msg = $out;
+  }
+
   $self->emit(json => Mojo::JSON->new->decode($msg))
     if $self->has_subscribers('json');
   $op = delete $self->{op};
@@ -443,6 +481,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(1);
+
+Compress messages with C<permessage-deflate> extension.
+
 =head2 handshake
 
   my $handshake = $ws->handshake;

diff --git a/t/mojo/websocket.t b/t/mojo/websocket.t
@@ -455,15 +455,16 @@ is $stash->{finished}, 1, 'finish event has been emitted once';
 like $log, qr/Inactivity timeout\./, 'right log message';
 app->log->unsubscribe(message => $msg);
 
-# Ping/pong
-my $pong;
+# Ping/pong (with negotiated compression)
+my ($pong, $extensions);
 $ua->websocket(
   '/echo' => sub {
     my ($ua, $tx) = @_;
     $tx->on(
       frame => sub {
         my ($tx, $frame) = @_;
         $pong = $frame->[5] if $frame->[4] == 10;
+        $extensions = $tx->res->headers->sec_websocket_extensions;
         Mojo::IOLoop->stop;
       }
     );
@@ -472,5 +473,6 @@ $ua->websocket(
 );
 Mojo::IOLoop->start;
 is $pong, 'test', 'received pong with payload';
+is $extensions, 'permessage-deflate', 'right "Sec-WebSocket-Extensions" value';
 
 done_testing();
diff --git a/t/mojolicious/websocket_lite_app.t b/t/mojolicious/websocket_lite_app.t
@@ -110,7 +110,7 @@ $t->websocket_ok('/echo')->send_ok('hello again')
   ->message_ok->message_is('echo: hello again')->send_ok('and one more time')
   ->message_ok->message_is('echo: and one more time')->finish_ok;
 
-# Custom headers and protocols
+# Custom headers and protocols with compression
 my $headers = {DNT => 1, 'Sec-WebSocket-Key' => 'NTA2MDAyMDU1NjMzNjkwMg=='};
 $t->websocket_ok('/echo' => $headers => ['foo', 'bar', 'baz'])
   ->header_is('Sec-WebSocket-Accept'   => 'I+x5C3/LJxrmDrWw42nMP4pCSes=')
@@ -119,6 +119,10 @@ $t->websocket_ok('/echo' => $headers => ['foo', 'bar', 'baz'])
 is $t->tx->req->headers->dnt, 1, 'right "DNT" value';
 is $t->tx->req->headers->sec_websocket_protocol, 'foo, bar, baz',
   'right "Sec-WebSocket-Protocol" value';
+is $t->tx->req->headers->sec_websocket_extensions, 'permessage-deflate',
+  'right "Sec-WebSocket-Extensions" value';
+is $t->tx->res->headers->sec_websocket_extensions, 'permessage-deflate',
+  'right "Sec-WebSocket-Extensions" value';
 
 # Bytes
 $t->websocket_ok('/echo')->send_ok({binary => 'bytes!'})
@@ -139,9 +143,9 @@ $t->send_ok({binary => 'a' x 262145})
   ->message_ok->message_is({binary => 'a' x 262145})
   ->finish_ok->finished_ok(1005);
 
-# 64bit binary message (too large)
-$t->websocket_ok('/echo')->send_ok({binary => 'b' x 262145})
-  ->finished_ok(1009);
+# 64bit binary message (too large and no compression)
+$t->websocket_ok('/echo' => {'Sec-WebSocket-Extensions' => 'nothing'})
+  ->send_ok({binary => 'b' x 262145})->finished_ok(1009);
 
 # Binary message in two 64bit frames without FIN bit (too large)
 $t->websocket_ok('/echo')->send_ok([0, 0, 0, 0, 2, 'c' x 100000])
@@ -194,9 +198,11 @@ $t->websocket_ok('/unicode')->send_ok('hello again')
   ->send_ok('and one ☃ more time')
   ->message_ok->message_is('♥: and one ☃ more time')->finish_ok;
 
-# Binary frame and events
+# Binary frame and events (no compression)
 my $bytes = b("I ♥ Mojolicious")->encode('UTF-16LE')->to_string;
-$t->websocket_ok('/bytes');
+$t->websocket_ok('/bytes' => {'Sec-WebSocket-Extensions' => 'nothing'});
+ok !$t->tx->res->headers->sec_websocket_extensions,
+  'no "Sec-WebSocket-Extensions" value';
 my $binary;
 $t->tx->on(
   frame => sub {
@@ -214,10 +220,10 @@ ok !$binary, 'received text frame';
 $t->finish_ok(1000 => 'Have a nice day!');
 is_deeply $close, [1000, 'Have a nice day!'], 'right status and message';
 
-# Binary roundtrips
-$t->websocket_ok('/bytes')->send_ok({binary => $bytes})
-  ->message_ok->message_is($bytes)->send_ok({binary => $bytes})
-  ->message_ok->message_is($bytes)->finish_ok;
+# Binary roundtrips (no compression)
+$t->websocket_ok('/bytes' => {'Sec-WebSocket-Extensions' => 'nothing'})
+  ->send_ok({binary => $bytes})->message_ok->message_is($bytes)
+  ->send_ok({binary => $bytes})->message_ok->message_is($bytes)->finish_ok;
 
 # Two responses
 $t->websocket_ok('/once')->send_ok('hello')