Repository: abhinavsingh/proxy.py
Pull request #1055 (2022-01-24 10:46:26)

[Core] Invoke `selector.close` on shutdown
This was done in order to debug the `N semaphores leaked` warnings.  Turns out this problem is in the Python core.  Sending out PR for the `close` related changes.

codecov[bot]: # [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) Report
> Merging [#1055](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (ed0fc7e) into [develop](https://codecov.io/gh/abhinavsingh/proxy.py/commit/3858f3ae76584c33614661a1ac40154a42ccf060?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) (3858f3a) will **decrease** coverage by `0.52%`.
> The diff coverage is `50.00%`.

[![Impacted file tree graph](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/graphs/tree.svg?width=650&height=150&src=pr&token=Zh9J7b4la2&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)

```diff
@@             Coverage Diff             @@
##           develop    #1055      +/-   ##
===========================================
- Coverage    87.32%   86.79%   -0.53%     
===========================================
  Files          160      160              
  Lines         6683     6688       +5     
  Branches       675      620      -55     
===========================================
- Hits          5836     5805      -31     
- Misses         741      760      +19     
- Partials       106      123      +17     
```

| Flag | Coverage Δ | |
|---|---|---|
| pytest | `86.79% <50.00%> (-0.35%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Impacted Files](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | Coverage Δ | |
|---|---|---|
| [proxy/http/websocket/client.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC93ZWJzb2NrZXQvY2xpZW50LnB5) | `52.77% <0.00%> (-0.75%)` | :arrow_down: |
| [proxy/core/acceptor/acceptor.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS9hY2NlcHRvci9hY2NlcHRvci5weQ==) | `95.57% <50.00%> (-1.75%)` | :arrow_down: |
| [proxy/http/handler.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvaHR0cC9oYW5kbGVyLnB5) | `75.22% <50.00%> (-2.46%)` | :arrow_down: |
| [proxy/core/work/threadless.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-cHJveHkvY29yZS93b3JrL3RocmVhZGxlc3MucHk=) | `81.28% <100.00%> (-1.89%)` | :arrow_down: |
| [tests/integration/test\_integration.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaW50ZWdyYXRpb24vdGVzdF9pbnRlZ3JhdGlvbi5weQ==) | `94.31% <0.00%> (-5.69%)` | :arrow_down: |
| [tests/test\_main.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvdGVzdF9tYWluLnB5) | `95.73% <0.00%> (-4.27%)` | :arrow_down: |
| [tests/http/parser/test\_proxy\_protocol.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvaHR0cC9wYXJzZXIvdGVzdF9wcm94eV9wcm90b2NvbC5weQ==) | `96.07% <0.00%> (-3.93%)` | :arrow_down: |
| [tests/test\_set\_open\_file\_limit.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvdGVzdF9zZXRfb3Blbl9maWxlX2xpbWl0LnB5) | `96.42% <0.00%> (-3.58%)` | :arrow_down: |
| [tests/testing/test\_embed.py](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh#diff-dGVzdHMvdGVzdGluZy90ZXN0X2VtYmVkLnB5) | `90.32% <0.00%> (-3.23%)` | :arrow_down: |
| ... and [6 more](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Last update [3858f3a...ed0fc7e](https://codecov.io/gh/abhinavsingh/proxy.py/pull/1055?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Abhinav+Singh).


Patch:
diff --git a/proxy/core/acceptor/acceptor.py b/proxy/core/acceptor/acceptor.py
index 0df4c2a..299c0ef 100644
--- a/proxy/core/acceptor/acceptor.py
+++ b/proxy/core/acceptor/acceptor.py
@@ -178,6 +178,7 @@ class Acceptor(multiprocessing.Process):
             for fileno in self.socks:
                 self.socks[fileno].close()
             self.socks.clear()
+            self.selector.close()
             logger.debug('Acceptor#%d shutdown', self.idd)
 
     def _recv_and_setup_socks(self) -> None:
@@ -207,7 +208,8 @@ class Acceptor(multiprocessing.Process):
         self._lthread.start()
 
     def _stop_local(self) -> None:
-        if self._lthread is not None and self._local_work_queue is not None:
+        if self._lthread is not None and \
+                self._local_work_queue is not None:
             self._local_work_queue.put(False)
             self._lthread.join()
 
diff --git a/proxy/core/work/threadless.py b/proxy/core/work/threadless.py
index 1ad84ac..f43c0a4 100644
--- a/proxy/core/work/threadless.py
+++ b/proxy/core/work/threadless.py
@@ -419,6 +419,7 @@ class Threadless(ABC, Generic[T]):
             if wqfileno is not None:
                 self.selector.unregister(wqfileno)
                 self.close_work_queue()
+            self.selector.close()
             assert self.loop is not None
             self.loop.run_until_complete(self.loop.shutdown_asyncgens())
             self.loop.close()
diff --git a/proxy/http/handler.py b/proxy/http/handler.py
index dc2d7cf..bea158a 100644
--- a/proxy/http/handler.py
+++ b/proxy/http/handler.py
@@ -351,6 +351,8 @@ class HttpProtocolHandler(BaseTcpServerHandler[HttpClientConnection]):
             )
         finally:
             self.shutdown()
+            if self.selector:
+                self.selector.close()
             loop.close()
 
     async def _run_once(self) -> bool:
diff --git a/proxy/http/parser/parser.py b/proxy/http/parser/parser.py
index 6573cab..00790c3 100644
--- a/proxy/http/parser/parser.py
+++ b/proxy/http/parser/parser.py
@@ -114,7 +114,7 @@ class HttpParser:
     def header(self, key: bytes) -> bytes:
         """Convenient method to return original header value from internal data structure."""
         if self.headers is None or key.lower() not in self.headers:
-            raise KeyError('%s not found in headers', text_(key))
+            raise KeyError('%s not found in headers' % text_(key))
         return self.headers[key.lower()][1]
 
     def has_header(self, key: bytes) -> bool:
diff --git a/proxy/http/websocket/client.py b/proxy/http/websocket/client.py
index c8c53eb..13e70e7 100644
--- a/proxy/http/websocket/client.py
+++ b/proxy/http/websocket/client.py
@@ -27,6 +27,9 @@ from ...common.constants import (
 
 
 class WebsocketClient(TcpConnection):
+    """Websocket client connection.
+
+    TODO: Make me compatible with the work framework."""
 
     def __init__(
         self,
@@ -57,10 +60,14 @@ class WebsocketClient(TcpConnection):
         return self.sock
 
     def handshake(self) -> None:
+        """Start websocket upgrade & handshake protocol"""
         self.upgrade()
         self.sock.setblocking(False)
 
     def upgrade(self) -> None:
+        """Creates a key and sends websocket handshake packet to upstream.
+        Receives response from the server and asserts that websocket
+        accept header is valid in the response."""
         key = base64.b64encode(secrets.token_bytes(16))
         self.sock.send(
             build_websocket_handshake_request(
@@ -74,12 +81,6 @@ class WebsocketClient(TcpConnection):
         accept = response.header(b'Sec-Websocket-Accept')
         assert WebsocketFrame.key_to_accept(key) == accept
 
-    def ping(self, data: Optional[bytes] = None) -> None:
-        pass    # pragma: no cover
-
-    def pong(self, data: Optional[bytes] = None) -> None:
-        pass    # pragma: no cover
-
     def shutdown(self, _data: Optional[bytes] = None) -> None:
         """Closes connection with the server."""
         super().close()
@@ -121,3 +122,4 @@ class WebsocketClient(TcpConnection):
                 except OSError:
                     pass
             self.sock.close()
+            self.selector.close()


Test patch:
diff --git a/tests/http/websocket/test_websocket_client.py b/tests/http/websocket/test_websocket_client.py
index 4e31e48..92df52e 100644
--- a/tests/http/websocket/test_websocket_client.py
+++ b/tests/http/websocket/test_websocket_client.py
@@ -8,6 +8,8 @@
     :copyright: (c) 2013-present by Abhinav Singh and contributors.
     :license: BSD, see LICENSE for more details.
 """
+import selectors
+
 import unittest
 from unittest import mock
 
@@ -15,18 +17,19 @@ from proxy.common.utils import (
     build_websocket_handshake_request, build_websocket_handshake_response,
 )
 from proxy.http.websocket import WebsocketFrame, WebsocketClient
-from proxy.common.constants import DEFAULT_PORT
+from proxy.common.constants import DEFAULT_PORT, DEFAULT_BUFFER_SIZE
 
 
 class TestWebsocketClient(unittest.TestCase):
 
-    @mock.patch('proxy.http.websocket.client.socket.gethostbyname')
     @mock.patch('base64.b64encode')
+    @mock.patch('proxy.http.websocket.client.socket.gethostbyname')
     @mock.patch('proxy.http.websocket.client.new_socket_connection')
-    def test_handshake(
-        self, mock_connect: mock.Mock,
-        mock_b64encode: mock.Mock,
-        mock_gethostbyname: mock.Mock,
+    def test_handshake_success(
+            self,
+            mock_connect: mock.Mock,
+            mock_gethostbyname: mock.Mock,
+            mock_b64encode: mock.Mock,
     ) -> None:
         key = b'MySecretKey'
         mock_b64encode.return_value = key
@@ -35,9 +38,71 @@ class TestWebsocketClient(unittest.TestCase):
             build_websocket_handshake_response(
                 WebsocketFrame.key_to_accept(key),
             )
+        mock_connect.assert_not_called()
         client = WebsocketClient(b'localhost', DEFAULT_PORT)
+        mock_connect.assert_called_once()
         mock_connect.return_value.send.assert_not_called()
         client.handshake()
         mock_connect.return_value.send.assert_called_with(
             build_websocket_handshake_request(key),
         )
+        mock_connect.return_value.recv.assert_called_once_with(
+            DEFAULT_BUFFER_SIZE,
+        )
+
+    @mock.patch('base64.b64encode')
+    @mock.patch('selectors.DefaultSelector')
+    @mock.patch('proxy.http.websocket.client.new_socket_connection')
+    def test_send_recv_frames_success(
+            self,
+            mock_connect: mock.Mock,
+            mock_selector: mock.Mock,
+            mock_b64encode: mock.Mock,
+    ) -> None:
+        key = b'MySecretKey'
+        mock_b64encode.return_value = key
+        mock_connect.return_value.recv.side_effect = [
+            build_websocket_handshake_response(
+                WebsocketFrame.key_to_accept(key),
+            ),
+            WebsocketFrame.text(b'world'),
+        ]
+
+        def on_message(frame: WebsocketFrame) -> None:
+            assert frame.build() == WebsocketFrame.text(b'world')
+
+        client = WebsocketClient(
+            b'localhost', DEFAULT_PORT, on_message=on_message,
+        )
+        mock_selector.assert_called_once()
+        client.handshake()
+        client.queue(memoryview(WebsocketFrame.text(b'hello')))
+        mock_connect.return_value.send.assert_called_once()
+        mock_selector.return_value.select.side_effect = [
+            [
+                (mock.Mock(), selectors.EVENT_WRITE),
+            ],
+        ]
+        client.run_once()
+        self.assertEqual(mock_connect.return_value.send.call_count, 2)
+        mock_selector.return_value.select.side_effect = [
+            [
+                (mock.Mock(), selectors.EVENT_READ),
+            ],
+        ]
+        client.run_once()
+
+    @mock.patch('selectors.DefaultSelector')
+    @mock.patch('proxy.http.websocket.client.new_socket_connection')
+    def test_run(
+            self,
+            mock_connect: mock.Mock,
+            mock_selector: mock.Mock,
+    ) -> None:
+        mock_selector.return_value.select.side_effect = KeyboardInterrupt
+        client = WebsocketClient(b'localhost', DEFAULT_PORT)
+        client.run()
+        mock_connect.return_value.shutdown.assert_called_once()
+        mock_connect.return_value.close.assert_called_once()
+        mock_selector.return_value.unregister.assert_called_once_with(mock_connect.return_value)
+        mock_selector.return_value.close.assert_called_once()