Add Client websocket deflate init support.(not tested)

Add no takeover extension to server and client.

aiohttp/client.py

@@ -349,7 +349,8 @@ def ws_connect(self, url, *,
                    origin=None,
                    headers=None,
                    proxy=None,
-                   proxy_auth=None):
+                   proxy_auth=None,
+                   compress=False):
         """Initiate websocket connection."""
         return _WSRequestContextManager(
             self._ws_connect(url,
@@ -363,7 +364,8 @@ def ws_connect(self, url, *,
                              origin=origin,
                              headers=headers,
                              proxy=proxy,
-                             proxy_auth=proxy_auth))
+                             proxy_auth=proxy_auth,
+                             compress=compress))
 
     @asyncio.coroutine
     def _ws_connect(self, url, *,
@@ -377,7 +379,8 @@ def _ws_connect(self, url, *,
                     origin=None,
                     headers=None,
                     proxy=None,
-                    proxy_auth=None):
+                    proxy_auth=None,
+                    compress=False):
 
         if headers is None:
             headers = CIMultiDict()
@@ -399,6 +402,8 @@ def _ws_connect(self, url, *,
             headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = ','.join(protocols)
         if origin is not None:
             headers[hdrs.ORIGIN] = origin
+        if compress:
+            headers[hdrs.SEC_WEBSOCKET_EXTENSIONS] = 'permessage-deflate'
 
         # send request
         resp = yield from self.get(url, headers=headers,
@@ -457,12 +462,32 @@ def _ws_connect(self, url, *,
                         protocol = proto
                         break
 
+            # websocket compress
+            compress_notakeover = False
+            if compress:
+                if hdrs.SEC_WEBSOCKET_EXTENSIONS not in resp.headers:
+                    compress = False
+                else:
+                    exts = resp.headers[
+                        hdrs.SEC_WEBSOCKET_EXTENSIONS].split(',')
+                    for ext in exts:
+                        params = [x.strip() for x in ext.split(';')]
+                        if params[0] == 'permessage-deflate':
+                            for param in params:
+                                if param == 'client_no_context_takeover':
+                                    compress_notakeover = True
+                                    break
+                            if compress_notakeover:
+                                break
+
             proto = resp.connection.protocol
             reader = FlowControlDataQueue(
                 proto, limit=2 ** 16, loop=self._loop)
             proto.set_parser(WebSocketReader(reader), reader)
             resp.connection.writer.set_tcp_nodelay(True)
-            writer = WebSocketWriter(resp.connection.writer, use_mask=True)
+            writer = WebSocketWriter(
+                resp.connection.writer, use_mask=True,
+                compress=compress, notakeover=compress_notakeover)
         except Exception:
             resp.close()
             raise

aiohttp/http_websocket.py

@@ -7,11 +7,10 @@
 import json
 import random
 import sys
+import zlib
 from enum import IntEnum
 from struct import Struct
 
-import zlib
-
 from . import hdrs
 from .helpers import NO_EXTENSIONS, noop
 from .http_exceptions import HttpBadRequest, HttpProcessingError
@@ -261,7 +260,8 @@ def _feed_data(self, data):
                     payload_merged = b''.join(self._partial)
 
                     if compressed:
-                        payload_merged = self._decompressobj.decompress(payload_merged + _WS_DEFLATE_TRAILING)
+                        payload_merged = self._decompressobj.decompress(
+                            payload_merged + _WS_DEFLATE_TRAILING)
 
                     self._partial.clear()
 
@@ -428,12 +428,13 @@ class WebSocketWriter:
 
     def __init__(self, stream, *,
                  use_mask=False, limit=DEFAULT_LIMIT, random=random.Random(),
-                 compress=False):
+                 compress=False, notakeover=False):
         self.stream = stream
         self.writer = stream.transport
         self.use_mask = use_mask
         self.randrange = random.randrange
         self.compress = compress
+        self.notakeover = notakeover
         self._closing = False
         self._limit = limit
         self._output_size = 0
@@ -445,8 +446,10 @@ def _send_frame(self, message, opcode):
             ws_logger.warning('websocket connection is closing.')
 
         rsv = 0
-        if self.compress and opcode < 8:
-            if not self._compressobj:
+
+        # Only compress larger packets
+        if self.compress and opcode < 8 and len(message) > 124:
+            if not self._compressobj or self.notakeover:
                 self._compressobj = zlib.compressobj(wbits=-self.compress)
 
             message = self._compressobj.compress(message)
@@ -592,30 +595,41 @@ def do_handshake(method, headers, stream,
             hashlib.sha1(key.encode() + WS_KEY).digest()).decode())]
 
     compress = 0
+    compress_notakeover = False
 
     extensions = headers.get(hdrs.SEC_WEBSOCKET_EXTENSIONS)
     if extensions:
-        extensions = [ [s.strip() for s in s1.split(';')]
-            for s1 in extensions.split(',')]
+        extensions = [[s.strip() for s in s1.split(';')]
+                      for s1 in extensions.split(',')]
 
         for ext in extensions:
             if ext[0] == 'permessage-deflate':
+                enabledext = ['permessage-deflate']
                 compress = 15
                 for param in ext[1:]:
                     if param.startswith('server_max_window_bits'):
                         compress = int(param.split('=')[1])
-                        break
+                        enabledext.append((
+                            'server_max_window_bits=' + str(compress)))
+                    elif param == 'server_no_context_takeover':
+                        compress_notakeover = True
+                        enabledext.append(('server_no_context_takeover'))
+                    # Ignore Client Takeover
+                    # elif param == 'client_no_context_takeover':
+                    #     compress_notakeover |= WSCompressNoTakeover.NT_CLIENT
+
+                response_headers.append((
+                    hdrs.SEC_WEBSOCKET_EXTENSIONS, '; '.join(enabledext)))
                 break
 
-        if compress:
-            response_headers.append((hdrs.SEC_WEBSOCKET_EXTENSIONS, 'permessage-deflate'))
-
     if protocol:
         response_headers.append((hdrs.SEC_WEBSOCKET_PROTOCOL, protocol))
 
     # response code, headers, None, writer, protocol
     return (101,
             response_headers,
             None,
-            WebSocketWriter(stream, limit=write_buffer_size, compress=compress),
+            WebSocketWriter(
+                stream, limit=write_buffer_size,
+                compress=compress, notakeover=compress_notakeover),
             protocol)