From 380e0cc3ce58f8e0b65f94908e707426fed9f54b Mon Sep 17 00:00:00 2001
From: Abhinav Singh <126065+abhinavsingh@users.noreply.github.com>
Date: Tue, 23 Apr 2024 18:51:53 +0530
Subject: [PATCH] Catch `KeyError` within Threadless executors (#1396)

---
 .gitignore                    |  1 -
 .vscode/settings.json         | 16 +++++++++++++++-
 proxy/core/work/threadless.py | 22 +++++++++++++++-------
 proxy/http/parser/parser.py   |  7 +++++++
 4 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/.gitignore b/.gitignore
index 088334e2..dcc36e2b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -36,6 +36,5 @@ build
 
 pyreverse.png
 profile.svg
-
 *-pre-push
 jaxl-api-credentials*.json
diff --git a/.vscode/settings.json b/.vscode/settings.json
index e2fed886..f8894c3d 100644
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@@ -41,5 +41,19 @@
   "python.linting.flake8Args": ["--config", ".flake8"],
   "python.linting.mypyEnabled": true,
   "python.formatting.provider": "autopep8",
-  "autoDocstring.docstringFormat": "sphinx"
+  "autoDocstring.docstringFormat": "sphinx",
+  "emeraldwalk.runonsave": {
+    "commands": [
+      {
+        "match": "\\.py$",
+        "isAsync": false,
+        "cmd": "./.venv/bin/autoflake --in-place --remove-all-unused-imports \"${file}\""
+      },
+      {
+        "match": "\\.py$",
+        "isAsync": false,
+        "cmd": "./.venv/bin/isort \"${file}\""
+      }
+    ]
+  }
 }
diff --git a/proxy/core/work/threadless.py b/proxy/core/work/threadless.py
index f43c0a47..e638940f 100644
--- a/proxy/core/work/threadless.py
+++ b/proxy/core/work/threadless.py
@@ -191,13 +191,21 @@ class Threadless(ABC, Generic[T]):
             #
             # TODO: Also remove offending work from pool to avoid spin loop.
             elif fileno != -1:
-                self.selector.register(fileno, events=mask, data=work_id)
-                self.registered_events_by_work_ids[work_id][fileno] = mask
-                logger.debug(
-                    'fd#{0} registered for mask#{1} by work#{2}'.format(
-                        fileno, mask, work_id,
-                    ),
-                )
+                try:
+                    self.selector.register(fileno, events=mask, data=work_id)
+                    self.registered_events_by_work_ids[work_id][fileno] = mask
+                    logger.debug(
+                        'fd#{0} registered for mask#{1} by work#{2}'.format(
+                            fileno,
+                            mask,
+                            work_id,
+                        ),
+                    )
+                except KeyError as exc:
+                    logger.warning(
+                        'KeyError when trying to register fd#{0}'.format(fileno),
+                        exc_info=exc,
+                    )
 
     async def _update_conn_pool_events(self) -> None:
         if not self._upstream_conn_pool:
diff --git a/proxy/http/parser/parser.py b/proxy/http/parser/parser.py
index 6b882f8f..77846671 100644
--- a/proxy/http/parser/parser.py
+++ b/proxy/http/parser/parser.py
@@ -216,6 +216,13 @@ class HttpParser:
             self.has_header(b'Connection') and \
             self.has_header(b'Upgrade')
 
+    @property
+    def is_websocket_upgrade(self) -> bool:
+        return (
+            self.is_connection_upgrade
+            and self.header(b'upgrade').lower() == b'websocket'
+        )
+
     @property
     def is_https_tunnel(self) -> bool:
         """Returns true for HTTPS CONNECT tunnel request."""