From fbf7250b2424cded3894c8ae7b58e29a0d395761 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Wed, 26 Jan 2022 17:19:14 +0530 Subject: [PATCH 1/3] Avoid using `tobytes` where possible --- proxy/core/connection/connection.py | 6 +++--- proxy/http/handler.py | 5 ++--- proxy/http/server/web.py | 6 ++++-- proxy/http/websocket/client.py | 2 +- tests/http/test_protocol_handler.py | 2 +- 5 files changed, 11 insertions(+), 10 deletions(-) diff --git a/proxy/core/connection/connection.py b/proxy/core/connection/connection.py index b4d387e0c6..5f0a39cd66 100644 --- a/proxy/core/connection/connection.py +++ b/proxy/core/connection/connection.py @@ -83,16 +83,16 @@ def flush(self, max_send_size: Optional[int] = None) -> int: """Users must handle BrokenPipeError exceptions""" if not self.has_buffer(): return 0 - mv = self.buffer[0].tobytes() - max_send_size = max_send_size or DEFAULT_MAX_SEND_SIZE + mv = self.buffer[0] # TODO: Assemble multiple packets if total # size remains below max send size. + max_send_size = max_send_size or DEFAULT_MAX_SEND_SIZE sent: int = self.send(mv[:max_send_size]) if sent == len(mv): self.buffer.pop(0) self._num_buffer -= 1 else: - self.buffer[0] = memoryview(mv[sent:]) + self.buffer[0] = mv[sent:] del mv logger.debug('flushed %d bytes to %s' % (sent, self.tag)) return sent diff --git a/proxy/http/handler.py b/proxy/http/handler.py index 6d6cf01c07..e1aff89ee6 100644 --- a/proxy/http/handler.py +++ b/proxy/http/handler.py @@ -267,10 +267,9 @@ def _discover_plugin_klass(self, protocol: int) -> Optional[Type['HttpProtocolHa def _parse_first_request(self, data: memoryview) -> bool: # Parse http request - # - # TODO(abhinavsingh): Remove .tobytes after parser is - # memoryview compliant try: + # TODO(abhinavsingh): Remove .tobytes after parser is + # memoryview compliant self.request.parse(data.tobytes()) except HttpProtocolException as e: # noqa: WPS329 self.work.queue(BAD_REQUEST_RESPONSE_PKT) diff --git a/proxy/http/server/web.py b/proxy/http/server/web.py index 04b97b4f50..c3376b3765 100644 --- a/proxy/http/server/web.py +++ b/proxy/http/server/web.py @@ -174,8 +174,10 @@ async def read_from_descriptors(self, r: Readables) -> bool: def on_client_data(self, raw: memoryview) -> None: if self.switched_protocol == httpProtocolTypes.WEBSOCKET: - # TODO(abhinavsingh): Remove .tobytes after websocket frame parser - # is memoryview compliant + # TODO(abhinavsingh): Do we really tobytes() here? + # Websocket parser currently doesn't depend on internal + # buffers, due to which it can directly parse out of + # memory views. But how about large payloads scenarios? remaining = raw.tobytes() frame = WebsocketFrame() while remaining != b'': diff --git a/proxy/http/websocket/client.py b/proxy/http/websocket/client.py index dd05e2f1d0..223ff50482 100644 --- a/proxy/http/websocket/client.py +++ b/proxy/http/websocket/client.py @@ -96,7 +96,7 @@ def run_once(self) -> bool: if mask & selectors.EVENT_READ and self.on_message: # TODO: client recvbuf size flag currently not used here raw = self.recv() - if raw is None or raw.tobytes() == b'': + if raw is None or raw == b'': self.closed = True return True frame = WebsocketFrame() diff --git a/tests/http/test_protocol_handler.py b/tests/http/test_protocol_handler.py index a2425fd723..a491e286f8 100644 --- a/tests/http/test_protocol_handler.py +++ b/tests/http/test_protocol_handler.py @@ -456,7 +456,7 @@ async def assert_data_queued( CRLF, ]) server.queue.assert_called_once() - self.assertEqual(server.queue.call_args_list[0][0][0].tobytes(), pkt) + self.assertEqual(server.queue.call_args_list[0][0][0], pkt) server.buffer_size.return_value = len(pkt) async def assert_data_queued_to_server(self, server: mock.Mock) -> None: From ed0db3efde4f44340ff48e573908ae8d7f49b7c1 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Wed, 26 Jan 2022 17:27:29 +0530 Subject: [PATCH 2/3] `send` accepts `Union[memoryview, bytes]` now --- proxy/core/connection/connection.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/proxy/core/connection/connection.py b/proxy/core/connection/connection.py index 5f0a39cd66..4d4901ed4f 100644 --- a/proxy/core/connection/connection.py +++ b/proxy/core/connection/connection.py @@ -10,7 +10,7 @@ """ import logging from abc import ABC, abstractmethod -from typing import List, Optional +from typing import List, Optional, Union from .types import tcpConnectionTypes from ...common.types import TcpOrTlsSocket @@ -47,7 +47,7 @@ def connection(self) -> TcpOrTlsSocket: """Must return the socket connection to use in this class.""" raise TcpConnectionUninitializedException() # pragma: no cover - def send(self, data: bytes) -> int: + def send(self, data: Union[memoryview, bytes]) -> int: """Users must handle BrokenPipeError exceptions""" # logger.info(data) return self.connection.send(data) From 04f85a973e51b2903c006e0b1e0e97a4e8be76c9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 26 Jan 2022 11:58:38 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- proxy/core/connection/connection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proxy/core/connection/connection.py b/proxy/core/connection/connection.py index 4d4901ed4f..a8c090d137 100644 --- a/proxy/core/connection/connection.py +++ b/proxy/core/connection/connection.py @@ -10,7 +10,7 @@ """ import logging from abc import ABC, abstractmethod -from typing import List, Optional, Union +from typing import List, Union, Optional from .types import tcpConnectionTypes from ...common.types import TcpOrTlsSocket