astelm/tg-ws-proxy
1
0
Fork 0
mirror of https://github.com/Flowseal/tg-ws-proxy.git synced 2026-07-10 05:50:25 +03:00
Code Issues Packages Projects Releases Wiki Activity

perf(bridge): split MTProto packets in O(N) instead of O(N^2) (#913)

Browse Source
This commit is contained in:
Konukhov Yaroslav
2026-05-30 09:25:56 -07:00
committed by GitHub
parent 5915a0e1f3
commit 49e62ca142
1 changed files with 30 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
proxy/bridge.py
+30 -21
View File
@@ -64,19 +64,27 @@ class MsgSplitter:
self._plain_buf.extend(self._dec.update(chunk)) self._plain_buf.extend(self._dec.update(chunk))
parts = [] parts = []
while self._cipher_buf: offset = 0
packet_len = self._next_packet_len() 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: if packet_len is None:
break break
if packet_len <= 0: if packet_len <= 0:
parts.append(bytes(self._cipher_buf)) parts.append(bytes(self._cipher_buf[offset:]))
self._cipher_buf.clear() offset = buf_len
self._plain_buf.clear()
self._disabled = True self._disabled = True
break break
parts.append(bytes(self._cipher_buf[:packet_len])) parts.append(bytes(self._cipher_buf[offset:offset + packet_len]))
del self._cipher_buf[:packet_len] offset += packet_len
del self._plain_buf[:packet_len]
if offset:
del self._cipher_buf[:offset]
del self._plain_buf[:offset]
return parts return parts
def flush(self) -> List[bytes]: def flush(self) -> List[bytes]:
@@ -87,22 +95,23 @@ class MsgSplitter:
self._plain_buf.clear() self._plain_buf.clear()
return [tail] return [tail]
def _next_packet_len(self) -> Optional[int]: def _next_packet_len(self, offset: int, avail: int) -> Optional[int]:
if not self._plain_buf: if avail <= 0:
return None return None
if self._proto == PROTO_ABRIDGED_INT: 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, if self._proto in (PROTO_INTERMEDIATE_INT,
PROTO_PADDED_INTERMEDIATE_INT): PROTO_PADDED_INTERMEDIATE_INT):
return self._next_intermediate_len() return self._next_intermediate_len(offset, avail)
return 0 return 0
def _next_abridged_len(self) -> Optional[int]: def _next_abridged_len(self, offset: int, avail: int) -> Optional[int]:
first = self._plain_buf[0] first = self._plain_buf[offset]
if first in (0x7F, 0xFF): if first in (0x7F, 0xFF):
if len(self._plain_buf) < 4: if avail < 4:
return None 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 header_len = 4
else: else:
payload_len = (first & 0x7F) * 4 payload_len = (first & 0x7F) * 4
@@ -110,18 +119,18 @@ class MsgSplitter:
if payload_len <= 0: if payload_len <= 0:
return 0 return 0
packet_len = header_len + payload_len packet_len = header_len + payload_len
if len(self._plain_buf) < packet_len: if avail < packet_len:
return None return None
return packet_len return packet_len
def _next_intermediate_len(self) -> Optional[int]: def _next_intermediate_len(self, offset: int, avail: int) -> Optional[int]:
if len(self._plain_buf) < 4: if avail < 4:
return None 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: if payload_len <= 0:
return 0 return 0
packet_len = 4 + payload_len packet_len = 4 + payload_len
if len(self._plain_buf) < packet_len: if avail < packet_len:
return None return None
return packet_len return packet_len