mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 11:08:44 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
18175fb4fe | ||
|
|
f742a8957d | ||
|
|
3686c36462 |
No files matched your search
+1
-1
@@ -15,7 +15,7 @@ RUN apt-get update \
|
||||
&& rm -rf /var/lib/apt/lists/*
|
||||
|
||||
WORKDIR /app
|
||||
RUN "$VIRTUAL_ENV/bin/pip" install cryptography==46.0.5 certifi
|
||||
RUN "$VIRTUAL_ENV/bin/pip" install cryptography==46.0.5 certifi "httpx[http2]==0.28.1"
|
||||
|
||||
FROM python:3.12-slim AS runtime
|
||||
|
||||
|
||||
@@ -4,8 +4,8 @@
|
||||
# http://msdn.microsoft.com/en-us/library/ms646997.aspx
|
||||
VSVersionInfo(
|
||||
ffi=FixedFileInfo(
|
||||
filevers=(1, 11, 0, 0),
|
||||
prodvers=(1, 11, 0, 0),
|
||||
filevers=(1, 11, 1, 0),
|
||||
prodvers=(1, 11, 1, 0),
|
||||
mask=0x3f,
|
||||
flags=0x0,
|
||||
OS=0x40004,
|
||||
@@ -21,12 +21,12 @@ VSVersionInfo(
|
||||
[
|
||||
StringStruct(u'CompanyName', u'Flowseal'),
|
||||
StringStruct(u'FileDescription', u'Telegram Desktop WebSocket Bridge Proxy'),
|
||||
StringStruct(u'FileVersion', u'1.11.0.0'),
|
||||
StringStruct(u'FileVersion', u'1.11.1.0'),
|
||||
StringStruct(u'InternalName', u'TgWsProxy'),
|
||||
StringStruct(u'LegalCopyright', u'Copyright (c) Flowseal. MIT License.'),
|
||||
StringStruct(u'OriginalFilename', u'TgWsProxy.exe'),
|
||||
StringStruct(u'ProductName', u'TG WS Proxy'),
|
||||
StringStruct(u'ProductVersion', u'1.11.0.0'),
|
||||
StringStruct(u'ProductVersion', u'1.11.1.0'),
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
from .config import parse_dc_ip_list, proxy_config, coerce_domain_list
|
||||
from .utils import get_link_host, build_github_opener
|
||||
|
||||
__version__ = "1.11.0"
|
||||
__version__ = "1.11.1"
|
||||
|
||||
__all__ = ["__version__", "get_link_host", "proxy_config", "parse_dc_ip_list", "build_github_opener", "coerce_domain_list"]
|
||||
+49
-12
@@ -79,6 +79,14 @@ class _ReplyBufferFull(BufferError):
|
||||
"""A local memory limit"""
|
||||
|
||||
|
||||
class _MTProtoTransportError(Exception):
|
||||
"""An endpoint error that the native client must receive, not just EOF."""
|
||||
|
||||
def __init__(self, code: int, source: str):
|
||||
super().__init__('%d (%s)' % (code, source))
|
||||
self.code = code
|
||||
|
||||
|
||||
@dataclass
|
||||
class _ReplayPacket:
|
||||
sent_at: float
|
||||
@@ -139,6 +147,7 @@ class _HttpChannel:
|
||||
def __init__(self, lane, channel_id: int, label: str):
|
||||
self.lane, self.channel_id, self.label = lane, channel_id, label
|
||||
self.closed = False
|
||||
self.transport_error: Optional[int] = None
|
||||
self.pending: Dict[asyncio.Task, int] = {}
|
||||
self.pending_since: Dict[asyncio.Task, float] = {}
|
||||
self.sent_since: Dict[asyncio.Task, float] = {}
|
||||
@@ -336,6 +345,8 @@ class _HttpChannel:
|
||||
def _fail(self, error: Exception) -> None:
|
||||
if self.closed:
|
||||
return
|
||||
if isinstance(error, _MTProtoTransportError):
|
||||
self.transport_error = error.code
|
||||
self.closed = True
|
||||
self.recovery_wakeup.set()
|
||||
while not self.queue.empty():
|
||||
@@ -460,14 +471,27 @@ class _HttpChannel:
|
||||
if trace is not None:
|
||||
trace.phase = 'cancelled-' + trace.phase
|
||||
raise
|
||||
except _MTProtoTransportError as exc:
|
||||
if trace is not None:
|
||||
trace.phase = 'transport-error'
|
||||
stats.h2_errors += 1
|
||||
if not self.closed:
|
||||
log.warning('[%s] H2 lane=%d channel=%d host=%s MTProto transport error=%s '
|
||||
'request=%d encrypted=%d bytes=%d replay=%d; reporting to client',
|
||||
self.label, self.lane.lane_id, self.channel_id, self.lane.host, exc,
|
||||
trace.request if trace is not None else 0,
|
||||
body[:8] != b'\x00' * 8, len(body), replay)
|
||||
self._fail(exc)
|
||||
except Exception as exc:
|
||||
if trace is not None:
|
||||
trace.phase = 'error-' + type(exc).__name__
|
||||
stats.h2_errors += 1
|
||||
if isinstance(exc, (httpx.HTTPError, OSError, ValueError, BufferError)):
|
||||
log.warning('[%s] H2 lane=%d channel=%d request failed: %s: %s; '
|
||||
'detail=[%s]; '
|
||||
'closing native channel, reconnect may select another domain',
|
||||
self.label, self.lane.lane_id, self.channel_id, type(exc).__name__, exc)
|
||||
self.label, self.lane.lane_id, self.channel_id, type(exc).__name__, exc,
|
||||
trace.summary(time.monotonic()) if trace is not None else '-')
|
||||
else:
|
||||
log.exception('[%s] H2 lane=%d channel=%d unexpected request failure',
|
||||
self.label, self.lane.lane_id, self.channel_id)
|
||||
@@ -610,8 +634,7 @@ class _HttpLane:
|
||||
extensions={'trace': self._trace}) as response:
|
||||
if response.http_version != 'HTTP/2':
|
||||
raise ValueError('CF /api preflight negotiated ' + response.http_version)
|
||||
if response.status_code in (403, 429) or (response.status_code >= 500
|
||||
and response.status_code != 501):
|
||||
if not (200 <= response.status_code < 300 or response.status_code in (405, 501)):
|
||||
raise ConnectionError('CF /api preflight HTTP %d cf_ray=%s' % (
|
||||
response.status_code, response.headers.get('cf-ray', '-')))
|
||||
log.debug('H2 lane=%d ready host=%s status=%d ms=%.0f',
|
||||
@@ -696,6 +719,9 @@ class _HttpLane:
|
||||
self.max_headers_seconds = max(self.max_headers_seconds, headers_at - started)
|
||||
if response.http_version != 'HTTP/2':
|
||||
raise ValueError('CF /api downgraded to ' + response.http_version)
|
||||
if response.status_code in (403, 404, 429, 444):
|
||||
raise _MTProtoTransportError(-response.status_code, 'HTTP %d cf_ray=%s' % (
|
||||
response.status_code, response.headers.get('cf-ray', '-')))
|
||||
if response.status_code != 200:
|
||||
raise ConnectionError('CF /api HTTP %d cf_ray=%s' % (
|
||||
response.status_code, response.headers.get('cf-ray', '-')))
|
||||
@@ -727,8 +753,9 @@ class _HttpLane:
|
||||
if content and len(content) % 4:
|
||||
raise ValueError('CF /api response is not an aligned MTProto packet')
|
||||
if len(content) == 4:
|
||||
log.warning('H2 lane=%d channel=%d request=%d MTProto transport error=%d',
|
||||
self.lane_id, channel_id, request_id, struct.unpack('<i', content)[0])
|
||||
code, = struct.unpack('<i', content)
|
||||
if code < 0:
|
||||
raise _MTProtoTransportError(code, 'HTTP 200 payload')
|
||||
self.responses += 1
|
||||
if channel is not None:
|
||||
channel.http_bytes += received
|
||||
@@ -750,6 +777,9 @@ class _HttpLane:
|
||||
(headers_at - started) * 1000, body_seconds * 1000,
|
||||
trace.summary(finished) if trace is not None else '-')
|
||||
return bytes(content)
|
||||
except _MTProtoTransportError:
|
||||
self.errors += 1
|
||||
raise
|
||||
except _ReplyBufferFull:
|
||||
raise
|
||||
except (httpx.HTTPError, OSError, ValueError, BufferError):
|
||||
@@ -930,19 +960,26 @@ async def bridge_h2(reader, writer, channel: _HttpChannel, ctx, tag: bytes) -> N
|
||||
tasks = [asyncio.create_task(upload()), asyncio.create_task(download()),
|
||||
asyncio.create_task(channel._recover())]
|
||||
try:
|
||||
done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
|
||||
for task in done:
|
||||
task.result()
|
||||
try:
|
||||
done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
|
||||
finally:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
if channel.transport_error is not None:
|
||||
await write_native(_encode_reply(struct.pack('<i', channel.transport_error), tag))
|
||||
channel.down += 4
|
||||
stats.bytes_down += 4
|
||||
else:
|
||||
for task in done:
|
||||
task.result()
|
||||
except asyncio.IncompleteReadError:
|
||||
log.debug('[%s] H2 channel=%d native EOF', channel.label, channel.channel_id)
|
||||
except _NativeDeliveryError as exc:
|
||||
log.debug('[%s] H2 channel=%d native closed: %s', channel.label, channel.channel_id, exc)
|
||||
except (OSError, ValueError, BufferError, asyncio.TimeoutError) as exc:
|
||||
except (httpx.HTTPError, OSError, ValueError, BufferError, asyncio.TimeoutError) as exc:
|
||||
if not channel.closed:
|
||||
log.warning('[%s] H2 channel=%d bridge ended: %s: %s',
|
||||
channel.label, channel.channel_id, type(exc).__name__, exc)
|
||||
finally:
|
||||
for task in tasks:
|
||||
task.cancel()
|
||||
await asyncio.gather(*tasks, return_exceptions=True)
|
||||
await channel.close()
|
||||
+67
-2
@@ -155,16 +155,75 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase):
|
||||
calls = []
|
||||
async def handler(request):
|
||||
calls.append(request.content)
|
||||
return httpx.Response(403, extensions={'http_version': b'HTTP/2'})
|
||||
return httpx.Response(503, extensions={'http_version': b'HTTP/2'})
|
||||
self._transport(handler)
|
||||
channel = self._channel()
|
||||
await channel.send(b'x' * 40, False)
|
||||
with self.assertRaisesRegex(ConnectionError, '403'):
|
||||
with self.assertRaisesRegex(ConnectionError, '503'):
|
||||
await asyncio.wait_for(channel.receive(), 1)
|
||||
self.assertEqual(calls, [b'x' * 40])
|
||||
self.assertTrue(channel.closed)
|
||||
self.assertFalse(self.lane.closed)
|
||||
|
||||
async def test_transport_error_reaches_client_before_close_without_cooling_domain(self):
|
||||
class ErrorBody(httpx.AsyncByteStream):
|
||||
async def __aiter__(self):
|
||||
raise AssertionError('HTTP transport error bodies must be ignored')
|
||||
yield b''
|
||||
|
||||
for tag in (PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE):
|
||||
for status in (403, 404, 429, 444):
|
||||
with self.subTest(tag=tag, status=status):
|
||||
calls = []
|
||||
|
||||
def handler(request):
|
||||
calls.append(request.content)
|
||||
if request.content == b'h' * 40:
|
||||
return httpx.Response(200, content=b'r' * 40,
|
||||
extensions={'http_version': b'HTTP/2'})
|
||||
return httpx.Response(status, stream=ErrorBody(),
|
||||
headers={'content-type': 'text/html'},
|
||||
extensions={'http_version': b'HTTP/2'})
|
||||
|
||||
self._transport(handler)
|
||||
channel = self._channel()
|
||||
output = bytearray()
|
||||
draining, release = asyncio.Event(), asyncio.Event()
|
||||
|
||||
async def drain():
|
||||
draining.set()
|
||||
await release.wait()
|
||||
|
||||
writer = SimpleNamespace(write=output.extend, drain=drain)
|
||||
reader = asyncio.StreamReader()
|
||||
packet = b'k' * 40
|
||||
reader.feed_data(_encode_reply(packet, tag))
|
||||
ctx = SimpleNamespace(clt_enc=_IdentityCipher(), clt_dec=_IdentityCipher())
|
||||
task = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, tag))
|
||||
try:
|
||||
await asyncio.wait_for(draining.wait(), 1)
|
||||
self.assertFalse(task.done())
|
||||
self.assertFalse(channel.close_done)
|
||||
if tag == PROTO_TAG_ABRIDGED:
|
||||
self.assertEqual(output, b'\x01' + struct.pack('<i', -status))
|
||||
else:
|
||||
length, = struct.unpack('<I', output[:4])
|
||||
self.assertEqual(length, len(output) - 4)
|
||||
self.assertEqual(output[4:8], struct.pack('<i', -status))
|
||||
self.assertLessEqual(length, 7)
|
||||
self.assertEqual(self.lane.failed_until, 0)
|
||||
healthy = self._channel()
|
||||
await healthy.send(b'h' * 40, False)
|
||||
self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'r' * 40)
|
||||
self.assertEqual(calls, [packet, b'h' * 40])
|
||||
self.assertFalse(self.lane.closed)
|
||||
finally:
|
||||
release.set()
|
||||
await asyncio.wait_for(task, 1)
|
||||
self.assertTrue(channel.close_done)
|
||||
self.assertFalse(channel.pending)
|
||||
await healthy.close()
|
||||
|
||||
async def test_long_polls_cannot_block_fresh_request_and_retired_packet_recovers(self):
|
||||
bodies = [b'samekey!' + bytes([index]) * 32 for index in range(MAX_CHANNEL_REQUESTS)]
|
||||
fresh = b'samekey!' + b'n' * 32
|
||||
@@ -352,6 +411,12 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase):
|
||||
self._transport(lambda request: httpx.Response(501, extensions={'http_version': b'HTTP/2'}))
|
||||
await self.lane._preflight()
|
||||
|
||||
async def test_missing_http_endpoint_is_rejected_before_native_requests(self):
|
||||
self._transport(lambda request: httpx.Response(
|
||||
404, extensions={'http_version': b'HTTP/2'}))
|
||||
with self.assertRaisesRegex(ConnectionError, 'preflight HTTP 404'):
|
||||
await self.lane._preflight()
|
||||
|
||||
async def _packet_bridge(self, *, content_length=True, truncate=False, concurrent_small=False):
|
||||
first_chunk_read, release_tail = asyncio.Event(), asyncio.Event()
|
||||
body = b'x' * (64 * 1024 + 40)
|
||||
|
||||
+226
-2
@@ -1,7 +1,9 @@
|
||||
"""Wire-level regressions: real TLS sockets, ALPN and HTTP/2 flow control."""
|
||||
import asyncio
|
||||
import datetime
|
||||
import hashlib
|
||||
import ipaddress
|
||||
import os
|
||||
import ssl
|
||||
import struct
|
||||
import tempfile
|
||||
@@ -9,7 +11,7 @@ import time
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
from h2.config import H2Configuration
|
||||
@@ -19,7 +21,10 @@ from h2.settings import SettingCodes
|
||||
|
||||
from proxy.h2_transport import H2Transport, STREAM_RECEIVE_WINDOW, _Connection
|
||||
from proxy.cf_h2 import _HttpChannel, _HttpLane, bridge_h2
|
||||
from proxy.utils import PROTO_TAG_INTERMEDIATE
|
||||
from proxy import tg_ws_proxy
|
||||
from proxy._aes import Cipher, algorithms, modes
|
||||
from proxy.config import proxy_config
|
||||
from proxy.utils import PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE
|
||||
|
||||
|
||||
class H2WireTest(unittest.IsolatedAsyncioTestCase):
|
||||
@@ -424,6 +429,184 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
|
||||
self.assertEqual(await first, b'a' * 65532)
|
||||
self.assertEqual(await second, b'b' * 65532)
|
||||
|
||||
async def test_http_404_is_forwarded_while_other_native_channel_keeps_working(self):
|
||||
class Cipher:
|
||||
def update(self, data):
|
||||
return data
|
||||
|
||||
class Writer:
|
||||
def __init__(self):
|
||||
self.data = bytearray()
|
||||
self.changed = asyncio.Event()
|
||||
|
||||
def write(self, data):
|
||||
self.data.extend(data)
|
||||
|
||||
async def drain(self):
|
||||
self.changed.set()
|
||||
|
||||
channel = _HttpChannel(self.lane, 1, 'bad-key-test')
|
||||
reader, writer = asyncio.StreamReader(), Writer()
|
||||
ctx = SimpleNamespace(clt_dec=Cipher(), clt_enc=Cipher())
|
||||
bridge = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE))
|
||||
self.tasks.append(bridge)
|
||||
reader.feed_data(struct.pack('<I', 40) + b'\x00' * 40)
|
||||
self.respond(await asyncio.wait_for(self.requests.get(), 1), b'p' * 84)
|
||||
await asyncio.wait_for(writer.changed.wait(), 1)
|
||||
reader.feed_data((struct.pack('<I', 40) + b'k' * 40) * 2)
|
||||
bad = await asyncio.wait_for(self.requests.get(), 1)
|
||||
pending = await asyncio.wait_for(self.requests.get(), 1)
|
||||
healthy = _HttpChannel(self.lane, 2, 'healthy-key-test')
|
||||
await healthy.send(b'h' * 40, False)
|
||||
good = await asyncio.wait_for(self.requests.get(), 1)
|
||||
self.respond(bad, b'<html>not found</html>', status=404, end=False)
|
||||
await asyncio.wait_for(bridge, 1)
|
||||
self.assertEqual(writer.data, struct.pack('<I', 84) + b'p' * 84
|
||||
+ struct.pack('<Ii', 4, -404))
|
||||
self.assertTrue(channel.close_done)
|
||||
self.assertFalse(channel.pending)
|
||||
self.assertNotIn(pending[1], self.transport.connection.streams)
|
||||
self.assertFalse(healthy.closed)
|
||||
self.respond(good, b'g' * 40)
|
||||
self.assertEqual(await asyncio.wait_for(healthy.receive(), 1), b'g' * 40)
|
||||
recovered = _HttpChannel(self.lane, 3, 'replacement-key-test')
|
||||
await recovered.send(b'n' * 40, False)
|
||||
self.respond(await asyncio.wait_for(self.requests.get(), 1), b'r' * 40)
|
||||
self.assertEqual(await asyncio.wait_for(recovered.receive(), 1), b'r' * 40)
|
||||
self.assertEqual(len(self.connections), 1)
|
||||
self.assertEqual(self.lane.failed_until, 0)
|
||||
|
||||
async def test_binary_transport_error_is_terminal_without_domain_cooldown(self):
|
||||
channel = _HttpChannel(self.lane, 1, 'binary-error-test')
|
||||
await channel.send(b'k' * 40, False)
|
||||
self.respond(await asyncio.wait_for(self.requests.get(), 1), struct.pack('<i', -404))
|
||||
await asyncio.gather(*list(channel.pending))
|
||||
self.assertTrue(channel.closed)
|
||||
self.assertEqual(channel.transport_error, -404)
|
||||
self.assertEqual(self.lane.failed_until, 0)
|
||||
self.assertEqual(self.lane.reply_buffer_bytes, 0)
|
||||
await channel.close()
|
||||
|
||||
async def test_obfuscated_clients_preserve_packets_through_full_h2_fallback(self):
|
||||
secret = os.urandom(16)
|
||||
channels = []
|
||||
clients = []
|
||||
handlers = []
|
||||
|
||||
async def open_channel(dc, label):
|
||||
self.assertEqual(dc, 2)
|
||||
channel = _HttpChannel(self.lane, len(channels) + 1, label)
|
||||
channels.append(channel)
|
||||
return channel
|
||||
|
||||
def accept(reader, writer):
|
||||
task = asyncio.create_task(tg_ws_proxy._handle_client(reader, writer, secret))
|
||||
handlers.append(task)
|
||||
self.tasks.append(task)
|
||||
|
||||
def client_init(tag):
|
||||
raw = bytearray(os.urandom(64))
|
||||
raw[0] = 0x42
|
||||
raw[56:62] = tag + struct.pack('<h', -2)
|
||||
keys = bytes(raw[8:56])
|
||||
up = Cipher(algorithms.AES(hashlib.sha256(keys[:32] + secret).digest()),
|
||||
modes.CTR(keys[32:])).encryptor()
|
||||
reverse = keys[::-1]
|
||||
down = Cipher(algorithms.AES(hashlib.sha256(reverse[:32] + secret).digest()),
|
||||
modes.CTR(reverse[32:])).encryptor()
|
||||
encrypted = up.update(bytes(raw))
|
||||
return bytes(raw[:56]) + encrypted[56:], up, down
|
||||
|
||||
def frame(body, tag, padding):
|
||||
if tag == PROTO_TAG_ABRIDGED:
|
||||
words = len(body) // 4
|
||||
prefix = (bytes([words | 0x80]) if words < 127
|
||||
else b'\xff' + words.to_bytes(3, 'little'))
|
||||
else:
|
||||
if tag == PROTO_TAG_SECURE:
|
||||
body += os.urandom(padding)
|
||||
prefix = struct.pack('<I', len(body) | 0x80000000)
|
||||
return prefix + body
|
||||
|
||||
async def check_reply(peer, expected, tag):
|
||||
reader, _, _, down = peer
|
||||
|
||||
async def read(count):
|
||||
return down.update(await asyncio.wait_for(reader.readexactly(count), 2))
|
||||
|
||||
if tag == PROTO_TAG_ABRIDGED:
|
||||
words = (await read(1))[0]
|
||||
if words == 127:
|
||||
words = int.from_bytes(await read(3), 'little')
|
||||
length = words * 4
|
||||
else:
|
||||
length, = struct.unpack('<I', await read(4))
|
||||
body = await read(length)
|
||||
self.assertEqual(body[:len(expected)], expected)
|
||||
self.assertIn(len(body) - len(expected), range(4) if tag == PROTO_TAG_SECURE else (0,))
|
||||
|
||||
server = await asyncio.start_server(accept, '127.0.0.1', 0)
|
||||
pool = SimpleNamespace(open=AsyncMock(side_effect=open_channel))
|
||||
try:
|
||||
with patch.object(tg_ws_proxy, 'cf_h2_pool', pool), \
|
||||
patch.object(tg_ws_proxy.ws_pool, 'get', AsyncMock(return_value=None)) as direct, \
|
||||
patch.multiple(proxy_config, fallback_cfproxy=True, cfproxy_worker_domains=[],
|
||||
force_test_dc=False, fake_tls_domain='', proxy_protocol=False):
|
||||
for tag in (PROTO_TAG_ABRIDGED, PROTO_TAG_INTERMEDIATE, PROTO_TAG_SECURE):
|
||||
with self.subTest(tag=tag.hex()):
|
||||
peers = []
|
||||
requests = {}
|
||||
for index in range(2):
|
||||
reader, writer = await asyncio.open_connection(
|
||||
'127.0.0.1', server.sockets[0].getsockname()[1])
|
||||
clients.append(writer)
|
||||
init, up, down = client_init(tag)
|
||||
peer = (reader, writer, up, down)
|
||||
peers.append(peer)
|
||||
bodies = [b'\x00' * 8 + os.urandom(8) + struct.pack('<I', 20) + os.urandom(20)]
|
||||
bodies += [b'samekey!' + os.urandom(size - 8) for size in (40, 131112, 56)]
|
||||
wire = init + up.update(b''.join(
|
||||
frame(body, tag, padding) for body, padding in zip(bodies, (15, 1, 7, 0))))
|
||||
for start, end in ((0, 13), (13, 57), (57, 65), (65, 78), (78, len(wire))):
|
||||
writer.write(wire[start:end])
|
||||
await writer.drain()
|
||||
await asyncio.sleep(0)
|
||||
expected = {body: packet for packet, body in enumerate(bodies)}
|
||||
for _ in bodies:
|
||||
request = await asyncio.wait_for(self.requests.get(), 2)
|
||||
self.assertTrue(request[2] in expected, 'HTTP body differs from native input')
|
||||
packet = expected.pop(request[2])
|
||||
requests[index, packet] = request
|
||||
|
||||
for index, packet in ((1, 2), (0, 2), (0, 0), (1, 0), (0, 1)):
|
||||
reply = bytes([0x30 + index, packet]) * 36
|
||||
self.respond(requests[index, packet], reply)
|
||||
await check_reply(peers[index], reply, tag)
|
||||
self.respond(requests[0, 3], b'not found', status=404)
|
||||
await check_reply(peers[0], struct.pack('<i', -404), tag)
|
||||
self.assertEqual(await asyncio.wait_for(peers[0][0].read(), 2), b'')
|
||||
self.respond(requests[1, 1], b'healthy!' * 9)
|
||||
await check_reply(peers[1], b'healthy!' * 9, tag)
|
||||
self.respond(requests[1, 3], struct.pack('<i', -429))
|
||||
await check_reply(peers[1], struct.pack('<i', -429), tag)
|
||||
self.assertEqual(await asyncio.wait_for(peers[1][0].read(), 2), b'')
|
||||
await asyncio.wait_for(asyncio.gather(*handlers), 2)
|
||||
self.assertEqual(direct.await_count, 6)
|
||||
for call in direct.await_args_list:
|
||||
self.assertEqual(call.args, (2, True))
|
||||
self.assertEqual(call.kwargs, {'is_test_dc': False})
|
||||
self.assertEqual(pool.open.await_count, 6)
|
||||
self.assertTrue(all(channel.close_done and not channel.pending for channel in channels))
|
||||
self.assertTrue(all(channel.quick_requests == 4 for channel in channels))
|
||||
self.assertEqual(len(self.connections), 1)
|
||||
self.assertEqual(self.lane.failed_until, 0)
|
||||
finally:
|
||||
for writer in clients:
|
||||
writer.close()
|
||||
await writer.wait_closed()
|
||||
server.close()
|
||||
await server.wait_closed()
|
||||
|
||||
async def test_file_chunk_fits_first_window_without_extra_network_round_trip(self):
|
||||
# A 128 KiB file part plus its MTProto envelope must fit before any
|
||||
# WINDOW_UPDATE arrives. This matters even when the consumer is fast.
|
||||
@@ -820,6 +1003,47 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
|
||||
self.respond(good_request)
|
||||
self.assertEqual(await good, b'r' * 40)
|
||||
|
||||
async def test_reset_during_response_closes_native_bridge_without_escaping(self):
|
||||
cipher = SimpleNamespace(update=lambda data: data)
|
||||
ctx = SimpleNamespace(clt_dec=cipher, clt_enc=cipher)
|
||||
output = bytearray()
|
||||
writer = SimpleNamespace(write=output.extend, drain=AsyncMock())
|
||||
reader = asyncio.StreamReader()
|
||||
reader.feed_data(struct.pack('<I', 40) + b'k' * 40)
|
||||
channel = _HttpChannel(self.lane, 1, 'reset-during-body')
|
||||
|
||||
with patch.object(channel, '_recover', side_effect=asyncio.Event().wait), \
|
||||
self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
|
||||
bridge = asyncio.create_task(bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE))
|
||||
self.tasks.append(bridge)
|
||||
request = await asyncio.wait_for(self.requests.get(), 1)
|
||||
healthy, good_request = await self.post(b'h' * 40)
|
||||
self.respond(request, b'partial-response!!!!', end=False, length=40)
|
||||
|
||||
async def body_received():
|
||||
while self.lane.reply_buffer_bytes != 20:
|
||||
await asyncio.sleep(.001)
|
||||
|
||||
await asyncio.wait_for(body_received(), 1)
|
||||
conn, sid, _ = request
|
||||
conn.h2.reset_stream(sid, error_code=2)
|
||||
self.flush(conn)
|
||||
await asyncio.wait_for(bridge, 1)
|
||||
self.assertTrue(channel.close_done)
|
||||
self.assertFalse(channel.pending)
|
||||
self.assertIsNone(channel.transport_error)
|
||||
self.assertFalse(output, 'A partial HTTP response must never reach the native client')
|
||||
self.assertEqual(self.lane.reply_buffer_bytes, 0)
|
||||
self.assertIsNone(self.transport.connection.error)
|
||||
self.respond(good_request)
|
||||
self.assertEqual(await asyncio.wait_for(healthy, 1), b'r' * 40)
|
||||
self.assertEqual(len(self.connections), 1)
|
||||
failures = [record for record in captured.records if record.levelname in ('WARNING', 'ERROR')]
|
||||
self.assertEqual(len(failures), 1)
|
||||
self.assertIn('reset: 2', failures[0].getMessage())
|
||||
self.assertIn('down=20', failures[0].getMessage())
|
||||
self.assertIn('sid=1', failures[0].getMessage())
|
||||
|
||||
async def test_connection_loss_fails_waiters_and_next_request_reconnects(self):
|
||||
tasks = [await self.post() for _ in range(3)]
|
||||
self.connections[0].writer.close()
|
||||
|
||||
Reference in new issue
Block a user