mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 02:58:45 +00:00
206 lines
7.5 KiB
Python
206 lines
7.5 KiB
Python
import os
|
|
import asyncio
|
|
import time
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
from proxy._aes import Cipher, algorithms, modes
|
|
from proxy.bridge import MsgSplitter, bridge_ws_reencrypt
|
|
from proxy.network_debug import WsActivity, log_ws_flow
|
|
from proxy.utils import (
|
|
PROTO_ABRIDGED_INT,
|
|
PROTO_INTERMEDIATE_INT,
|
|
PROTO_PADDED_INTERMEDIATE_INT,
|
|
)
|
|
|
|
|
|
def _relay_init() -> bytes:
|
|
return os.urandom(64)
|
|
|
|
|
|
def _encryptor(relay_init: bytes):
|
|
enc = Cipher(
|
|
algorithms.AES(relay_init[8:40]), modes.CTR(relay_init[40:56])
|
|
).encryptor()
|
|
enc.update(b'\x00' * 64)
|
|
return enc
|
|
|
|
|
|
def _abridged(payload: bytes) -> bytes:
|
|
words = len(payload) // 4
|
|
if words < 0x7F:
|
|
return bytes([words]) + payload
|
|
return b'\x7f' + words.to_bytes(3, 'little') + payload
|
|
|
|
|
|
def _intermediate(payload: bytes) -> bytes:
|
|
return len(payload).to_bytes(4, 'little') + payload
|
|
|
|
|
|
class MsgSplitterTest(unittest.TestCase):
|
|
def _split(self, proto_int, packets, chunk_sizes=None):
|
|
relay_init = _relay_init()
|
|
splitter = MsgSplitter(relay_init, proto_int)
|
|
enc = _encryptor(relay_init)
|
|
stream = enc.update(b''.join(packets))
|
|
|
|
chunks = []
|
|
if chunk_sizes is None:
|
|
chunks = [stream]
|
|
else:
|
|
offset = 0
|
|
for size in chunk_sizes:
|
|
chunks.append(stream[offset:offset + size])
|
|
offset += size
|
|
if offset < len(stream):
|
|
chunks.append(stream[offset:])
|
|
|
|
parts = []
|
|
for chunk in chunks:
|
|
parts.extend(splitter.split(chunk))
|
|
return splitter, stream, parts
|
|
|
|
def test_abridged_stream_splits_into_packets(self):
|
|
packets = [_abridged(b'a' * 4), _abridged(b'b' * 16), _abridged(b'c' * 40)]
|
|
_, stream, parts = self._split(PROTO_ABRIDGED_INT, packets)
|
|
self.assertEqual(len(parts), 3)
|
|
self.assertEqual(b''.join(parts), stream)
|
|
self.assertEqual([len(p) for p in parts], [5, 17, 41])
|
|
|
|
def test_intermediate_stream_splits_into_packets(self):
|
|
packets = [_intermediate(b'a' * 8), _intermediate(b'b' * 12)]
|
|
_, stream, parts = self._split(PROTO_INTERMEDIATE_INT, packets)
|
|
self.assertEqual(len(parts), 2)
|
|
self.assertEqual(b''.join(parts), stream)
|
|
|
|
def test_padded_intermediate_uses_intermediate_framing(self):
|
|
packets = [_intermediate(b'z' * 20)]
|
|
_, stream, parts = self._split(PROTO_PADDED_INTERMEDIATE_INT, packets)
|
|
self.assertEqual(parts, [stream])
|
|
|
|
def test_partial_packet_is_buffered_until_complete(self):
|
|
packets = [_abridged(b'a' * 20)]
|
|
_, stream, parts = self._split(
|
|
PROTO_ABRIDGED_INT, packets, chunk_sizes=[1] * (len(packets[0]) - 1)
|
|
)
|
|
self.assertEqual(parts, [stream])
|
|
|
|
def test_split_preserves_stream_across_arbitrary_chunking(self):
|
|
packets = [_intermediate(bytes([i]) * 16) for i in range(8)]
|
|
_, stream, parts = self._split(
|
|
PROTO_INTERMEDIATE_INT, packets, chunk_sizes=[7, 3, 50, 11]
|
|
)
|
|
self.assertEqual(b''.join(parts), stream)
|
|
self.assertEqual(len(parts), 8)
|
|
|
|
def test_empty_chunk_yields_nothing(self):
|
|
splitter = MsgSplitter(_relay_init(), PROTO_INTERMEDIATE_INT)
|
|
self.assertEqual(splitter.split(b''), [])
|
|
|
|
def test_zero_length_packet_disables_splitting(self):
|
|
relay_init = _relay_init()
|
|
splitter = MsgSplitter(relay_init, PROTO_INTERMEDIATE_INT)
|
|
enc = _encryptor(relay_init)
|
|
stream = enc.update((0).to_bytes(4, 'little') + b'tail')
|
|
parts = splitter.split(stream)
|
|
self.assertEqual(parts, [stream])
|
|
self.assertEqual(splitter.split(b'raw'), [b'raw'])
|
|
|
|
def test_flush_returns_buffered_tail_once(self):
|
|
relay_init = _relay_init()
|
|
splitter = MsgSplitter(relay_init, PROTO_INTERMEDIATE_INT)
|
|
enc = _encryptor(relay_init)
|
|
partial = enc.update(_intermediate(b'x' * 32)[:10])
|
|
self.assertEqual(splitter.split(partial), [])
|
|
self.assertEqual(splitter.flush(), [partial])
|
|
self.assertEqual(splitter.flush(), [])
|
|
|
|
|
|
class WsDiagnosticsTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_flow_distinguishes_upstream_send_from_native_backpressure(self):
|
|
class Cipher:
|
|
def update(self, data):
|
|
return data
|
|
|
|
class Socket:
|
|
def __init__(self):
|
|
self.queue = asyncio.Queue()
|
|
self.sending = asyncio.Event()
|
|
self.release = asyncio.Event()
|
|
self.sent = []
|
|
|
|
async def send(self, data):
|
|
self.sent.append(data)
|
|
self.sending.set()
|
|
await self.release.wait()
|
|
|
|
async def recv(self):
|
|
return await self.queue.get()
|
|
|
|
async def close(self):
|
|
pass
|
|
|
|
class Writer:
|
|
def __init__(self):
|
|
self.data = bytearray()
|
|
self.written = asyncio.Event()
|
|
self.release = asyncio.Event()
|
|
|
|
def write(self, data):
|
|
self.data.extend(data)
|
|
self.written.set()
|
|
|
|
async def drain(self):
|
|
await self.release.wait()
|
|
|
|
def close(self):
|
|
pass
|
|
|
|
async def wait_closed(self):
|
|
pass
|
|
|
|
ws, writer = Socket(), Writer()
|
|
reader = asyncio.StreamReader()
|
|
ctx = SimpleNamespace(clt_dec=Cipher(), clt_enc=Cipher(), tg_dec=Cipher(), tg_enc=Cipher())
|
|
with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
|
|
bridge = asyncio.create_task(bridge_ws_reencrypt(reader, writer, ws, 'owned-test', ctx, dc=203))
|
|
try:
|
|
reader.feed_data(b'private-native-payload')
|
|
await asyncio.wait_for(ws.sending.wait(), 1)
|
|
activity = next(item for item in WsActivity.active if item.label == 'owned-test')
|
|
self.assertIsNotNone(activity.sending)
|
|
self.assertEqual(activity.ws_up, 0)
|
|
log_ws_flow(time.monotonic())
|
|
# A reply arriving before send drain finishes must not create
|
|
# a spurious wait for another response when drain resumes.
|
|
await ws.queue.put(b'private-upstream-payload')
|
|
await asyncio.wait_for(writer.written.wait(), 1)
|
|
self.assertEqual(activity.ws_down, len(writer.data))
|
|
self.assertEqual(activity.native_down, 0)
|
|
self.assertIsNotNone(activity.writing_native)
|
|
log_ws_flow(time.monotonic())
|
|
ws.release.set()
|
|
writer.release.set()
|
|
for _ in range(100):
|
|
if activity.native_down and activity.ws_up:
|
|
break
|
|
await asyncio.sleep(.001)
|
|
self.assertEqual(activity.native_down, len(writer.data))
|
|
self.assertIsNone(activity.awaiting_rx)
|
|
self.assertEqual(ws.sent, [b'private-native-payload'])
|
|
self.assertEqual(bytes(writer.data), b'private-upstream-payload')
|
|
reader.feed_eof()
|
|
await asyncio.wait_for(bridge, 1)
|
|
finally:
|
|
bridge.cancel()
|
|
await asyncio.gather(bridge, return_exceptions=True)
|
|
self.assertFalse(WsActivity.active)
|
|
output = '\n'.join(captured.output)
|
|
self.assertIn('WS FLOW DC203', output)
|
|
self.assertIn('WS END DC203', output)
|
|
self.assertNotIn('private-', output)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|