Fixes #356 WebSockets not working with TLS interception (#365)

Author: abhinavsingh

helper/chrome_with_proxy.sh

@@ -20,6 +20,6 @@ fi
   --no-first-run \
   --no-default-browser-check \
   --user-data-dir="$(mktemp -d -t 'chrome-remote_data_dir')" \
-  --proxy-server=${PROXY_PY_ADDR} \
-  --ignore-urlfetcher-cert-requests \
-  --ignore-certificate-errors
+  --proxy-server=${PROXY_PY_ADDR} #\
+#  --ignore-urlfetcher-cert-requests \
+#  --ignore-certificate-errors

proxy/http/handler.py

@@ -288,13 +288,13 @@ def handle_writables(self, writables: List[Union[int, HasFileno]]) -> bool:
 
             try:
                 self.client.flush()
-            except OSError:
-                logger.error('OSError when flushing buffer to client')
-                return True
             except BrokenPipeError:
                 logger.error(
                     'BrokenPipeError when flushing buffer for client')
                 return True
+            except OSError:
+                logger.error('OSError when flushing buffer to client')
+                return True
         return False
 
     def handle_readables(self, readables: List[Union[int, HasFileno]]) -> bool:
@@ -359,7 +359,7 @@ def handle_readables(self, readables: List[Union[int, HasFileno]]) -> bool:
             except HttpProtocolException as e:
                 logger.debug(
                     'HttpProtocolException type raised')
-                response = e.response(self.request)
+                response: Optional[memoryview] = e.response(self.request)
                 if response:
                     self.client.queue(response)
                 return True

proxy/http/parser.py

@@ -258,3 +258,8 @@ def is_http_1_1_keep_alive(self) -> bool:
         return self.version == HTTP_1_1 and \
             (not self.has_header(b'Connection') or
              self.header(b'Connection').lower() == b'keep-alive')
+
+    def is_connection_upgrade(self) -> bool:
+        return self.version == HTTP_1_1 and \
+            self.has_header(b'Connection') and \
+            self.has_header(b'Upgrade')

proxy/http/proxy/server.py

@@ -207,6 +207,15 @@ def on_client_data(self, raw: memoryview) -> Optional[memoryview]:
             if self.request.state == httpParserStates.COMPLETE and (
                     self.request.method != httpMethods.CONNECT or
                     self.flags.tls_interception_enabled()):
+
+                if self.pipeline_request is not None and \
+                        self.pipeline_request.is_connection_upgrade():
+                    # Previous pipelined request was a WebSocket
+                    # upgrade request. Incoming client data now
+                    # must be treated as WebSocket protocol packets.
+                    self.server.queue(raw)
+                    return None
+
                 if self.pipeline_request is None:
                     self.pipeline_request = HttpParser(
                         httpParserTypes.REQUEST_PARSER)
@@ -226,7 +235,8 @@ def on_client_data(self, raw: memoryview) -> Optional[memoryview]:
                     self.server.queue(
                         memoryview(
                             self.pipeline_request.build()))
-                    self.pipeline_request = None
+                    if not self.pipeline_request.is_connection_upgrade():
+                        self.pipeline_request = None
             else:
                 self.server.queue(raw)
             return None

proxy/http/websocket.py

@@ -231,7 +231,7 @@ def run_once(self) -> bool:
         self.selector.register(self.sock.fileno(), ev)
         events = self.selector.select(timeout=1)
         self.selector.unregister(self.sock)
-        for key, mask in events:
+        for _, mask in events:
             if mask & selectors.EVENT_READ and self.on_message:
                 raw = self.recv()
                 if raw is None or raw.tobytes() == b'':