astelm/tg-ws-proxy
1
0
Fork 0
mirror of https://github.com/Flowseal/tg-ws-proxy.git synced 2026-06-12 09:41:43 +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

51
proxy/bridge.py
View File

@@ -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