mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-09-05 18:16:11 +03:00
delewer
131 lines
4.1 KiB
Python
131 lines
4.1 KiB
Python
import asyncio
|
|
import unittest
|
|
|
|
from proxy.raw_websocket import RawWebSocket, WsHandshakeError, _xor_mask
|
|
|
|
|
|
def _raw_frame(opcode, data, fin=True):
|
|
b0 = (0x80 if fin else 0x00) | opcode
|
|
n = len(data)
|
|
if n < 126:
|
|
return bytes([b0, n]) + data
|
|
if n < 65536:
|
|
return bytes([b0, 126]) + n.to_bytes(2, 'big') + data
|
|
return bytes([b0, 127]) + n.to_bytes(8, 'big') + data
|
|
|
|
|
|
class _NullWriter:
|
|
def write(self, data):
|
|
pass
|
|
|
|
async def drain(self):
|
|
pass
|
|
|
|
|
|
def _recv(chunks, cls=RawWebSocket):
|
|
async def _run():
|
|
reader = asyncio.StreamReader()
|
|
for chunk in chunks:
|
|
reader.feed_data(chunk)
|
|
reader.feed_eof()
|
|
ws = cls(reader, _NullWriter())
|
|
return ws, await ws.recv()
|
|
|
|
return asyncio.run(_run())
|
|
|
|
|
|
class XorMaskTest(unittest.TestCase):
|
|
def test_roundtrip(self):
|
|
data = bytes(range(256)) * 3
|
|
mask = b'\x01\x02\x03\x04'
|
|
self.assertEqual(_xor_mask(_xor_mask(data, mask), mask), data)
|
|
|
|
def test_empty_payload(self):
|
|
self.assertEqual(_xor_mask(b'', b'\x01\x02\x03\x04'), b'')
|
|
|
|
|
|
class BuildFrameTest(unittest.TestCase):
|
|
def test_short_unmasked_frame(self):
|
|
self.assertEqual(
|
|
RawWebSocket._build_frame(RawWebSocket.OP_BINARY, b'abc'),
|
|
b'\x82\x03abc',
|
|
)
|
|
|
|
def test_extended_length_selects_16bit_header(self):
|
|
frame = RawWebSocket._build_frame(RawWebSocket.OP_BINARY, b'x' * 200)
|
|
self.assertEqual(frame[:2], b'\x82\x7e')
|
|
self.assertEqual(int.from_bytes(frame[2:4], 'big'), 200)
|
|
|
|
def test_masked_frame_sets_mask_bit_and_is_reversible(self):
|
|
payload = b'payload'
|
|
frame = RawWebSocket._build_frame(
|
|
RawWebSocket.OP_BINARY, payload, mask=True)
|
|
self.assertTrue(frame[1] & 0x80)
|
|
self.assertEqual(_xor_mask(frame[6:], frame[2:6]), payload)
|
|
|
|
|
|
class RecvTest(unittest.TestCase):
|
|
def test_returns_unfragmented_message(self):
|
|
_, msg = _recv([_raw_frame(RawWebSocket.OP_BINARY, b'hello')])
|
|
self.assertEqual(msg, b'hello')
|
|
|
|
def test_reassembles_fragmented_message(self):
|
|
_, msg = _recv([
|
|
_raw_frame(RawWebSocket.OP_BINARY, b'AAA', False),
|
|
_raw_frame(RawWebSocket.OP_CONT, b'BBB', False),
|
|
_raw_frame(RawWebSocket.OP_CONT, b'CCC', True),
|
|
])
|
|
self.assertEqual(msg, b'AAABBBCCC')
|
|
|
|
def test_control_frame_between_fragments_is_skipped(self):
|
|
_, msg = _recv([
|
|
_raw_frame(RawWebSocket.OP_BINARY, b'AAA', False),
|
|
_raw_frame(RawWebSocket.OP_PONG, b''),
|
|
_raw_frame(RawWebSocket.OP_CONT, b'BBB', True),
|
|
])
|
|
self.assertEqual(msg, b'AAABBB')
|
|
|
|
def test_close_frame_returns_none(self):
|
|
ws, msg = _recv([_raw_frame(RawWebSocket.OP_CLOSE, b'\x03\xe8')])
|
|
self.assertIsNone(msg)
|
|
self.assertTrue(ws._closed)
|
|
|
|
def test_oversized_frame_is_rejected_before_reading_payload(self):
|
|
header = bytes([0x82, 127]) + (1 << 40).to_bytes(8, 'big')
|
|
with self.assertRaises(ConnectionError):
|
|
_recv([header])
|
|
|
|
def test_reassembled_message_exceeding_limit_is_rejected(self):
|
|
class _Capped(RawWebSocket):
|
|
__slots__ = ()
|
|
MAX_MESSAGE_LEN = 1500
|
|
|
|
chunk = b'x' * 1024
|
|
with self.assertRaises(ConnectionError):
|
|
_recv([
|
|
_raw_frame(RawWebSocket.OP_BINARY, chunk, False),
|
|
_raw_frame(RawWebSocket.OP_CONT, chunk, False),
|
|
], cls=_Capped)
|
|
|
|
|
|
class ParseCloseTest(unittest.TestCase):
|
|
def test_known_code_gets_name(self):
|
|
code, reason = RawWebSocket._parse_close(b'\x03\xe8bye')
|
|
self.assertEqual(code, 1000)
|
|
self.assertIn('normal', reason)
|
|
|
|
def test_empty_payload(self):
|
|
self.assertEqual(RawWebSocket._parse_close(b''), (None, ''))
|
|
|
|
|
|
class HandshakeErrorTest(unittest.TestCase):
|
|
def test_redirect_status_codes(self):
|
|
for code in (301, 302, 303, 307, 308):
|
|
self.assertTrue(WsHandshakeError(code, '').is_redirect)
|
|
for code in (0, 200, 429, 502):
|
|
self.assertFalse(WsHandshakeError(code, '').is_redirect)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|