diff --git a/proxy/bridge.py b/proxy/bridge.py index 0fa2e15..eff59eb 100644 --- a/proxy/bridge.py +++ b/proxy/bridge.py @@ -64,19 +64,27 @@ class MsgSplitter: self._plain_buf.extend(self._dec.update(chunk)) parts = [] - while self._cipher_buf: - packet_len = self._next_packet_len() + offset = 0 + buf_len = len(self._cipher_buf) + # Walk the buffer with an offset instead of deleting each packet from + # the front. Front-deletion on a bytearray shifts the remaining bytes, + # so a chunk holding many small packets degrades to O(N^2); a single + # trailing del keeps splitting O(N). + while offset < buf_len: + packet_len = self._next_packet_len(offset, buf_len - offset) if packet_len is None: break if packet_len <= 0: - parts.append(bytes(self._cipher_buf)) - self._cipher_buf.clear() - self._plain_buf.clear() + parts.append(bytes(self._cipher_buf[offset:])) + offset = buf_len self._disabled = True break - parts.append(bytes(self._cipher_buf[:packet_len])) - del self._cipher_buf[:packet_len] - del self._plain_buf[:packet_len] + parts.append(bytes(self._cipher_buf[offset:offset + packet_len])) + offset += packet_len + + if offset: + del self._cipher_buf[:offset] + del self._plain_buf[:offset] return parts def flush(self) -> List[bytes]: @@ -87,22 +95,23 @@ class MsgSplitter: self._plain_buf.clear() return [tail] - def _next_packet_len(self) -> Optional[int]: - if not self._plain_buf: + def _next_packet_len(self, offset: int, avail: int) -> Optional[int]: + if avail <= 0: return None if self._proto == PROTO_ABRIDGED_INT: - return self._next_abridged_len() + return self._next_abridged_len(offset, avail) if self._proto in (PROTO_INTERMEDIATE_INT, PROTO_PADDED_INTERMEDIATE_INT): - return self._next_intermediate_len() + return self._next_intermediate_len(offset, avail) return 0 - def _next_abridged_len(self) -> Optional[int]: - first = self._plain_buf[0] + def _next_abridged_len(self, offset: int, avail: int) -> Optional[int]: + first = self._plain_buf[offset] if first in (0x7F, 0xFF): - if len(self._plain_buf) < 4: + if avail < 4: return None - payload_len = int.from_bytes(self._plain_buf[1:4], 'little') * 4 + payload_len = int.from_bytes( + self._plain_buf[offset + 1:offset + 4], 'little') * 4 header_len = 4 else: payload_len = (first & 0x7F) * 4 @@ -110,18 +119,18 @@ class MsgSplitter: if payload_len <= 0: return 0 packet_len = header_len + payload_len - if len(self._plain_buf) < packet_len: + if avail < packet_len: return None return packet_len - def _next_intermediate_len(self) -> Optional[int]: - if len(self._plain_buf) < 4: + def _next_intermediate_len(self, offset: int, avail: int) -> Optional[int]: + if avail < 4: return None - payload_len = _st_I_le.unpack_from(self._plain_buf, 0)[0] & 0x7FFFFFFF + payload_len = _st_I_le.unpack_from(self._plain_buf, offset)[0] & 0x7FFFFFFF if payload_len <= 0: return 0 packet_len = 4 + payload_len - if len(self._plain_buf) < packet_len: + if avail < packet_len: return None return packet_len