mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 02:58:45 +00:00
1125 lines
55 KiB
Python
1125 lines
55 KiB
Python
"""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
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
import httpx
|
|
from h2.config import H2Configuration
|
|
from h2.connection import H2Connection
|
|
from h2.events import DataReceived, RequestReceived, StreamEnded, StreamReset
|
|
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 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 H2ContextTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_lanes_reuse_verified_context_without_changing_websocket_contexts(self):
|
|
from proxy.h2_transport import _default_ssl_context
|
|
from proxy.raw_websocket import _ssl_ctx, _ssl_ctx_fronting
|
|
from proxy.utils import create_ssl_context
|
|
|
|
_default_ssl_context.cache_clear()
|
|
self.addCleanup(_default_ssl_context.cache_clear)
|
|
lanes = []
|
|
try:
|
|
with patch('proxy.h2_transport.create_ssl_context', wraps=create_ssl_context) as create:
|
|
lanes = [_HttpLane('kws4.example.org', 1), _HttpLane('kws2.example.net', 2)]
|
|
create.assert_called_once_with()
|
|
first, second = (lane.client._transport.ssl_context for lane in lanes)
|
|
self.assertIs(first, second)
|
|
self.assertIsNot(first, _ssl_ctx)
|
|
self.assertIsNot(first, _ssl_ctx_fronting)
|
|
self.assertTrue(first.check_hostname)
|
|
self.assertEqual(first.verify_mode, ssl.CERT_REQUIRED)
|
|
self.assertTrue(_ssl_ctx.check_hostname)
|
|
self.assertFalse(_ssl_ctx_fronting.check_hostname)
|
|
self.assertEqual(set(first.get_ca_certs(binary_form=True)),
|
|
set(_ssl_ctx.get_ca_certs(binary_form=True)))
|
|
|
|
custom = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
|
|
custom.check_hostname = False
|
|
transport = H2Transport(custom)
|
|
self.assertIs(transport.ssl_context, custom)
|
|
self.assertTrue(first.check_hostname)
|
|
await transport.aclose()
|
|
finally:
|
|
await asyncio.gather(*(lane.close() for lane in lanes))
|
|
|
|
|
|
class H2WireTest(unittest.IsolatedAsyncioTestCase):
|
|
@classmethod
|
|
def setUpClass(cls):
|
|
from cryptography import x509
|
|
from cryptography.hazmat.primitives import hashes, serialization
|
|
from cryptography.hazmat.primitives.asymmetric import rsa
|
|
from cryptography.x509.oid import NameOID
|
|
|
|
cls.folder = tempfile.TemporaryDirectory(prefix='tg-h2-tests-')
|
|
cls.addClassCleanup(cls.folder.cleanup)
|
|
key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
|
name = x509.Name([x509.NameAttribute(NameOID.COMMON_NAME, 'localhost')])
|
|
now = datetime.datetime.now(datetime.timezone.utc)
|
|
cert = (x509.CertificateBuilder().subject_name(name).issuer_name(name)
|
|
.public_key(key.public_key()).serial_number(x509.random_serial_number())
|
|
.not_valid_before(now - datetime.timedelta(minutes=1))
|
|
.not_valid_after(now + datetime.timedelta(hours=1))
|
|
.add_extension(x509.SubjectAlternativeName([
|
|
x509.IPAddress(ipaddress.ip_address('127.0.0.1'))]), critical=False)
|
|
.sign(key, hashes.SHA256()))
|
|
cls.cert_path = Path(cls.folder.name) / 'cert.pem'
|
|
key_path = Path(cls.folder.name) / 'key.pem'
|
|
cls.cert_path.write_bytes(cert.public_bytes(serialization.Encoding.PEM))
|
|
key_path.write_bytes(key.private_bytes(serialization.Encoding.PEM,
|
|
serialization.PrivateFormat.PKCS8,
|
|
serialization.NoEncryption()))
|
|
cls.server_ssl = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
|
|
cls.server_ssl.load_cert_chain(str(cls.cert_path), str(key_path))
|
|
cls.server_ssl.set_alpn_protocols(['h2'])
|
|
|
|
async def asyncSetUp(self):
|
|
self.requests = asyncio.Queue()
|
|
self.resets = asyncio.Queue()
|
|
self.connections = []
|
|
self.handlers = []
|
|
self.tasks = []
|
|
self.server_max_streams = 64
|
|
self.server = await asyncio.start_server(self.serve, '127.0.0.1', 0, ssl=self.server_ssl)
|
|
port = self.server.sockets[0].getsockname()[1]
|
|
self.lane = _HttpLane('127.0.0.1:%d' % port, 1)
|
|
await self.lane.client.aclose()
|
|
self.transport = H2Transport(ssl.create_default_context(cafile=str(self.cert_path)))
|
|
self.lane.client = httpx.AsyncClient(transport=self.transport, trust_env=False,
|
|
timeout=httpx.Timeout(2, read=1))
|
|
|
|
async def asyncTearDown(self):
|
|
for task in self.tasks:
|
|
task.cancel()
|
|
await asyncio.gather(*self.tasks, return_exceptions=True)
|
|
await self.lane.close()
|
|
self.server.close()
|
|
await self.server.wait_closed()
|
|
for connection in self.connections:
|
|
connection.writer.close()
|
|
await asyncio.wait_for(asyncio.gather(*self.handlers), 3)
|
|
|
|
async def serve(self, reader, writer):
|
|
self.handlers.append(asyncio.current_task())
|
|
conn = SimpleNamespace(h2=H2Connection(config=H2Configuration(client_side=False)),
|
|
writer=writer, bodies={}, pending={})
|
|
self.connections.append(conn)
|
|
conn.h2.initiate_connection()
|
|
conn.h2.update_settings({SettingCodes.MAX_CONCURRENT_STREAMS: self.server_max_streams})
|
|
writer.write(conn.h2.data_to_send())
|
|
try:
|
|
while True:
|
|
data = await reader.read(65536)
|
|
if not data:
|
|
break
|
|
for event in conn.h2.receive_data(data):
|
|
if isinstance(event, RequestReceived):
|
|
conn.bodies[event.stream_id] = bytearray()
|
|
elif isinstance(event, DataReceived):
|
|
conn.bodies[event.stream_id].extend(event.data)
|
|
conn.h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id)
|
|
elif isinstance(event, StreamEnded):
|
|
await self.requests.put((conn, event.stream_id, bytes(conn.bodies[event.stream_id])))
|
|
elif isinstance(event, StreamReset):
|
|
conn.pending.pop(event.stream_id, None)
|
|
await self.resets.put((event.stream_id, event.error_code))
|
|
self.flush(conn)
|
|
await writer.drain()
|
|
except (ConnectionError, ssl.SSLError):
|
|
pass
|
|
finally:
|
|
writer.close()
|
|
try:
|
|
await writer.wait_closed()
|
|
except (ConnectionError, ssl.SSLError):
|
|
pass
|
|
|
|
def flush(self, conn):
|
|
for sid, (body, end) in list(conn.pending.items()):
|
|
while body:
|
|
count = min(len(body), conn.h2.max_outbound_frame_size,
|
|
conn.h2.local_flow_control_window(sid))
|
|
if count <= 0:
|
|
break
|
|
chunk, body = body[:count], body[count:]
|
|
conn.h2.send_data(sid, chunk, end_stream=end and not body)
|
|
if body:
|
|
conn.pending[sid] = (body, end)
|
|
else:
|
|
del conn.pending[sid]
|
|
conn.writer.write(conn.h2.data_to_send())
|
|
|
|
def respond(self, request, body=b'r' * 40, *, end=True, length=None, status=200):
|
|
conn, sid, _ = request
|
|
conn.h2.send_headers(sid, [(b':status', str(status).encode()),
|
|
(b'content-type', b'application/octet-stream'),
|
|
(b'content-length', str(len(body) if length is None else length).encode())],
|
|
end_stream=end and not body)
|
|
if body:
|
|
conn.pending[sid] = (body, end)
|
|
self.flush(conn)
|
|
|
|
async def post(self, body=b'x' * 40):
|
|
task = asyncio.create_task(self.lane._post(body, len(self.tasks) + 1))
|
|
self.tasks.append(task)
|
|
request = await asyncio.wait_for(self.requests.get(), 2)
|
|
self.assertEqual(request[2], body)
|
|
return task, request
|
|
|
|
async def test_fast_response_does_not_wait_for_unrelated_long_polls(self):
|
|
slow1, request1 = await self.post()
|
|
slow2, request2 = await self.post()
|
|
fast, request3 = await self.post()
|
|
self.respond(request3)
|
|
self.assertEqual(await asyncio.wait_for(fast, .5), b'r' * 40)
|
|
self.assertFalse(slow1.done())
|
|
self.assertFalse(slow2.done())
|
|
self.respond(request1)
|
|
self.respond(request2)
|
|
await asyncio.gather(slow1, slow2)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_shared_default_context_rejects_untrusted_certificate(self):
|
|
async with httpx.AsyncClient(transport=H2Transport(), trust_env=False) as client:
|
|
with self.assertRaises(httpx.ConnectError):
|
|
await client.get('https://' + self.lane.host + '/api')
|
|
task, request = await self.post()
|
|
self.respond(request)
|
|
self.assertEqual(await asyncio.wait_for(task, 1), b'r' * 40)
|
|
|
|
async def test_debug_trace_follows_real_stream_without_logging_payload(self):
|
|
channel = _HttpChannel(self.lane, 1, 'trace-test')
|
|
body = b'private-payload-' + b'x' * 24
|
|
with self.assertLogs('tg-mtproto-proxy', level='DEBUG') as captured:
|
|
await channel.send(body, False)
|
|
request = await asyncio.wait_for(self.requests.get(), 1)
|
|
trace = next(iter(channel.request_traces.values()))
|
|
|
|
async def phase(expected):
|
|
while trace.phase != expected:
|
|
await asyncio.sleep(.001)
|
|
|
|
await asyncio.wait_for(phase('headers'), .5)
|
|
self.assertEqual(trace.stream, request[1])
|
|
self.assertEqual(trace.packet, 1)
|
|
self.assertEqual(trace.request, 1)
|
|
self.assertIsNotNone(trace.sent_at)
|
|
self.assertIsNone(trace.headers_at)
|
|
self.respond(request, b'r' * 20, end=False, length=40)
|
|
await asyncio.wait_for(phase('body'), .5)
|
|
self.assertIsNotNone(trace.headers_at)
|
|
conn, sid, _ = request
|
|
conn.h2.send_data(sid, b'r' * 20, end_stream=True)
|
|
self.flush(conn)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
|
|
await asyncio.gather(*list(channel.pending))
|
|
self.assertEqual(trace.received, 40)
|
|
self.assertEqual(trace.phase, 'complete')
|
|
self.assertFalse(channel.request_traces)
|
|
self.assertIs(channel.recent_traces[-1], trace)
|
|
self.lane.log_flow(trace.finished_at)
|
|
await channel.close()
|
|
output = '\n'.join(captured.output)
|
|
self.assertIn('http_down=40', output)
|
|
self.assertIn('small=1 large=0', output)
|
|
self.assertNotIn('private-payload', output)
|
|
self.assertFalse(channel.recent_traces)
|
|
|
|
async def test_stale_recovery_slot_replaced_without_resetting_original_request(self):
|
|
channel = _HttpChannel(self.lane, 1, 'recovery-wire-test')
|
|
bodies = [b'samekey!' + bytes([index]) * 32 for index in range(3)]
|
|
originals = []
|
|
for body in bodies[:2]:
|
|
await channel.send(body, False)
|
|
originals.append(await asyncio.wait_for(self.requests.get(), 1))
|
|
|
|
with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_SLOT_SECONDS', .1):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
old_poll = await asyncio.wait_for(self.requests.get(), 1)
|
|
other_poll = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.assertEqual([old_poll[2], other_poll[2]], bodies[:2])
|
|
|
|
# The two duplicate requests remain silent. A new native request
|
|
# also gets no reply until its own duplicate solicits one.
|
|
await channel.send(bodies[2], False)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
probe = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.assertEqual(probe[2], bodies[2])
|
|
reset_sid, reset_code = await asyncio.wait_for(self.resets.get(), 1)
|
|
self.assertEqual((reset_sid, reset_code), (old_poll[1], 8))
|
|
self.assertIn(original[1], self.transport.connection.streams)
|
|
for request in originals:
|
|
self.assertIn(request[1], self.transport.connection.streams)
|
|
self.assertIn(other_poll[1], self.transport.connection.streams)
|
|
self.assertEqual(len(channel.replay_pending), 2)
|
|
self.respond(probe, b'r' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
|
|
# The original stream remains usable, even after recovery.
|
|
self.respond(original, b'o' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'o' * 40)
|
|
self.assertEqual(channel.replay_rotations, 1)
|
|
self.assertEqual(self.lane.failed_until, 0)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_normal_long_poll_survives_six_seconds_and_delivers_late_reply(self):
|
|
channel = _HttpChannel(self.lane, 1, 'normal-long-poll-test')
|
|
bodies = [b'samekey!' + bytes([index]) * 32 for index in range(3)]
|
|
for body in bodies[:2]:
|
|
await channel.send(body, False)
|
|
self.respond(await asyncio.wait_for(self.requests.get(), 1))
|
|
await asyncio.wait_for(channel.receive(), 1)
|
|
with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
poll = await asyncio.wait_for(self.requests.get(), 1)
|
|
# Advance request ages, keeping the real TLS/H2 streams open.
|
|
# Six seconds is still inside Telegram's default 25s HTTP wait.
|
|
for task in channel.replay_pending:
|
|
channel.pending_since[task] = time.monotonic() - 6
|
|
channel.sent_since[task] = time.monotonic() - 6
|
|
await asyncio.sleep(.03)
|
|
self.assertTrue(self.resets.empty(), 'normal HTTP wait was reset before its delayed result')
|
|
await channel.send(bodies[2], False)
|
|
fresh = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.assertEqual(fresh[2], bodies[2])
|
|
self.respond(fresh, b'n' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'n' * 40)
|
|
self.respond(poll, b'd' * (128 * 1024))
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'd' * (128 * 1024))
|
|
self.assertEqual(channel.replay_rotations, 0)
|
|
self.assertEqual(channel.replay_requests, 1)
|
|
|
|
async def test_replay_rotation_does_not_interrupt_upload(self):
|
|
channel = _HttpChannel(self.lane, 1, 'replay-upload-test')
|
|
body = b'samekey!' + b'p' * 32
|
|
await channel.send(body, False)
|
|
self.respond(await asyncio.wait_for(self.requests.get(), 1))
|
|
await asyncio.wait_for(channel.receive(), 1)
|
|
upload_started, resume = asyncio.Event(), asyncio.Event()
|
|
send_body = _Connection.send_body
|
|
|
|
async def paused_send(connection, request, stream):
|
|
upload_started.set()
|
|
await resume.wait()
|
|
await send_body(connection, request, stream)
|
|
|
|
with patch.object(_Connection, 'send_body', paused_send), \
|
|
patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
await asyncio.wait_for(upload_started.wait(), 1)
|
|
probe = next(iter(channel.replay_pending))
|
|
channel.pending_since[probe] = time.monotonic() - 40
|
|
await asyncio.sleep(.03)
|
|
self.assertFalse(probe.done())
|
|
self.assertTrue(self.resets.empty())
|
|
self.assertEqual(channel.replay_rotations, 0)
|
|
resume.set()
|
|
request = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.respond(request)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
|
|
|
|
async def test_recovery_deadline_starts_after_original_upload_finishes(self):
|
|
channel = _HttpChannel(self.lane, 1, 'unfinished-upload-test')
|
|
body = b'samekey!' + b'q' * 32
|
|
started, resume = asyncio.Event(), asyncio.Event()
|
|
send_body = _Connection.send_body
|
|
|
|
async def paused_send(connection, request, stream):
|
|
started.set()
|
|
await resume.wait()
|
|
await send_body(connection, request, stream)
|
|
|
|
with patch.object(_Connection, 'send_body', paused_send), \
|
|
patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .02), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
|
|
await channel.send(body, False)
|
|
await asyncio.wait_for(started.wait(), .5)
|
|
self.tasks.append(asyncio.create_task(channel._recover()))
|
|
await asyncio.sleep(.06)
|
|
self.assertEqual(channel.requests, 1, 'unsent original was duplicated')
|
|
self.assertTrue(self.requests.empty())
|
|
resume.set()
|
|
original = await asyncio.wait_for(self.requests.get(), .5)
|
|
probe = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(original[2], body)
|
|
self.assertEqual(probe[2], body)
|
|
self.respond(probe)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
|
|
self.assertTrue(self.resets.empty())
|
|
|
|
async def test_body_tail_does_not_wait_for_unrelated_long_polls(self):
|
|
slow1, _ = await self.post()
|
|
slow2, _ = await self.post()
|
|
fast, request = await self.post()
|
|
self.respond(request, b'a' * 20, end=False, length=40)
|
|
await asyncio.sleep(.03)
|
|
self.assertFalse(fast.done())
|
|
conn, sid, _ = request
|
|
conn.h2.send_data(sid, b'b' * 20, end_stream=True)
|
|
self.flush(conn)
|
|
self.assertEqual(await asyncio.wait_for(fast, .5), b'a' * 20 + b'b' * 20)
|
|
self.assertFalse(slow1.done())
|
|
self.assertFalse(slow2.done())
|
|
|
|
async def test_receiving_body_is_marked_before_request_coroutine_resumes(self):
|
|
channel = _HttpChannel(self.lane, 1, 'early-headers-test')
|
|
resume = asyncio.Event()
|
|
send_body = _Connection.send_body
|
|
|
|
async def paused_send(connection, request, stream):
|
|
await send_body(connection, request, stream)
|
|
await resume.wait()
|
|
|
|
with patch.object(_Connection, 'send_body', paused_send):
|
|
await channel.send(b'k' * 40, False)
|
|
request = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.respond(request, b'a' * 20, end=False, length=40)
|
|
|
|
async def headers_received():
|
|
while channel.receiving_http == 0:
|
|
await asyncio.sleep(.005)
|
|
try:
|
|
await asyncio.wait_for(headers_received(), .5)
|
|
self.assertEqual(channel.receiving_http, 1)
|
|
self.assertTrue(channel.queue.empty())
|
|
finally:
|
|
resume.set()
|
|
conn, sid, _ = request
|
|
conn.h2.send_data(sid, b'b' * 20, end_stream=True)
|
|
self.flush(conn)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'a' * 20 + b'b' * 20)
|
|
self.assertEqual(channel.receiving_http, 0)
|
|
|
|
async def test_native_channel_delivers_ready_reply_while_large_body_is_pending(self):
|
|
class IdentityCipher:
|
|
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)
|
|
self.changed.set()
|
|
|
|
async def drain(self):
|
|
pass
|
|
|
|
reader, writer = asyncio.StreamReader(), Writer()
|
|
channel = _HttpChannel(self.lane, 1, 'wire-test')
|
|
ctx = SimpleNamespace(clt_enc=IdentityCipher(), clt_dec=IdentityCipher())
|
|
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'x' * 40)
|
|
slow = await asyncio.wait_for(self.requests.get(), 2)
|
|
self.respond(slow, b'a' * 65536, end=False, length=65576)
|
|
reader.feed_data(struct.pack('<I', 40) + b'y' * 40)
|
|
fast = await asyncio.wait_for(self.requests.get(), 2)
|
|
self.respond(fast, b'f' * 40)
|
|
await asyncio.wait_for(writer.changed.wait(), .5)
|
|
self.assertEqual(writer.data, struct.pack('<I', 40) + b'f' * 40)
|
|
writer.changed.clear()
|
|
conn, sid, _ = slow
|
|
# Finish after the first 64 KiB have passed through flow control.
|
|
for _ in range(100):
|
|
if sid not in conn.pending:
|
|
break
|
|
await asyncio.sleep(.001)
|
|
self.assertNotIn(sid, conn.pending)
|
|
conn.pending[sid] = (b'b' * 40, True)
|
|
self.flush(conn)
|
|
await asyncio.wait_for(writer.changed.wait(), .5)
|
|
self.assertEqual(writer.data, struct.pack('<I', 40) + b'f' * 40
|
|
+ struct.pack('<I', 65576) + b'a' * 65536 + b'b' * 40)
|
|
|
|
async def test_parallel_first_windows_do_not_wait_for_connection_credit(self):
|
|
first, request1 = await self.post()
|
|
second, request2 = await self.post()
|
|
self.respond(request1, b'a' * 65532)
|
|
self.respond(request2, b'b' * 65532)
|
|
# No event-loop yield between responses: the server has not received
|
|
# any WINDOW_UPDATE credit for either body yet.
|
|
self.assertFalse(request1[0].pending)
|
|
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.
|
|
body = b'f' * (128 * 1024 + 168)
|
|
task, request = await self.post()
|
|
self.respond(request, body)
|
|
self.assertNotIn(request[1], list(request[0].pending))
|
|
self.assertEqual(await asyncio.wait_for(task, 1), body)
|
|
|
|
async def test_stalled_body_does_not_block_recovery_of_another_native_request(self):
|
|
channel = _HttpChannel(self.lane, 1, 'stalled-body-test')
|
|
await channel.send(b'k' * 8 + b'a' * 32, False)
|
|
slow = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.respond(slow, b'a' * 20, end=False, length=40)
|
|
fresh_body = b'k' * 8 + b'b' * 32
|
|
await channel.send(fresh_body, False)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .05):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
probe = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(probe[2], fresh_body)
|
|
self.respond(probe, b'r' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
|
|
self.assertIn(slow[1], self.transport.connection.streams)
|
|
self.assertIn(original[1], self.transport.connection.streams)
|
|
self.assertTrue(self.resets.empty())
|
|
conn, sid, _ = slow
|
|
conn.h2.send_data(sid, b'a' * 20, end_stream=True)
|
|
self.flush(conn)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'a' * 40)
|
|
self.assertEqual(self.lane.failed_until, 0)
|
|
|
|
async def test_first_recovery_without_result_does_not_end_recovery(self):
|
|
channel = _HttpChannel(self.lane, 1, 'second-recovery-test')
|
|
body = b'k' * 40
|
|
await channel.send(body, False)
|
|
self.respond(await asyncio.wait_for(self.requests.get(), 1), b'')
|
|
with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .02):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
first = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.respond(first, b'')
|
|
second = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(first[2], body)
|
|
self.assertEqual(second[2], body)
|
|
self.respond(second, b'r' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
|
|
self.assertEqual(channel.replay_requests, 2)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_completed_http_ack_does_not_add_idle_delay_to_each_chunk(self):
|
|
channel = _HttpChannel(self.lane, 1, 'chunked-download-poll-test')
|
|
# An HTTP reply can be only an MTProto acknowledgement. The queued
|
|
# file chunk needs another HTTP receiver, with no new native upload.
|
|
# Keep the production 1s idle threshold; shorten only scheduler ticks.
|
|
with patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
|
|
self.tasks.append(asyncio.create_task(channel._recover()))
|
|
for index in range(4):
|
|
body = b'samekey!' + bytes([index]) * 32
|
|
await channel.send(body, False)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.assertEqual(original[2], body)
|
|
self.respond(original, b'a' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'a' * 40)
|
|
poll = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(poll[2], body)
|
|
reply = bytes([index]) * (128 * 1024)
|
|
self.respond(poll, reply)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), reply)
|
|
self.assertTrue(self.resets.empty())
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_drained_receiver_wakes_first_poll_without_waiting_for_timer(self):
|
|
channel = _HttpChannel(self.lane, 1, 'receiver-wakeup-test')
|
|
body = b'samekey!' + b'q' * 32
|
|
with patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', 10):
|
|
self.tasks.append(asyncio.create_task(channel._recover()))
|
|
await channel.send(body, False)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.respond(original, b'a' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'a' * 40)
|
|
poll = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(poll[2], body)
|
|
self.respond(poll, b'f' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'f' * 40)
|
|
await asyncio.sleep(.03)
|
|
self.assertTrue(self.requests.empty(), 'a wakeup bypassed the retry backoff')
|
|
self.assertEqual(channel.replay_requests, 1)
|
|
|
|
async def test_native_write_completion_wakes_poll_but_backpressure_does_not(self):
|
|
class IdentityCipher:
|
|
def update(self, data):
|
|
return data
|
|
|
|
class Writer:
|
|
def __init__(self):
|
|
self.started, self.resume = asyncio.Event(), asyncio.Event()
|
|
|
|
def write(self, data):
|
|
self.started.set()
|
|
|
|
async def drain(self):
|
|
await self.resume.wait()
|
|
|
|
channel = _HttpChannel(self.lane, 1, 'native-wakeup-test')
|
|
reader, writer = asyncio.StreamReader(), Writer()
|
|
ctx = SimpleNamespace(clt_enc=IdentityCipher(), clt_dec=IdentityCipher())
|
|
body = b'samekey!' + b'q' * 32
|
|
with patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', 10):
|
|
self.tasks.append(asyncio.create_task(
|
|
bridge_h2(reader, writer, channel, ctx, PROTO_TAG_INTERMEDIATE)))
|
|
reader.feed_data(struct.pack('<I', len(body)) + body)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.respond(original)
|
|
await asyncio.wait_for(writer.started.wait(), .5)
|
|
await asyncio.sleep(.03)
|
|
self.assertTrue(self.requests.empty())
|
|
writer.resume.set()
|
|
poll = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(poll[2], body)
|
|
self.assertIsNone(channel.delivering_since)
|
|
self.assertEqual(channel.down, 40)
|
|
|
|
async def test_idle_poll_preserves_slot_for_new_overdue_request(self):
|
|
channel = _HttpChannel(self.lane, 1, 'idle-poll-capacity-test')
|
|
bodies = [b'samekey!' + bytes([index]) * 32 for index in range(4)]
|
|
for body in bodies[:3]:
|
|
await channel.send(body, False)
|
|
self.respond(await asyncio.wait_for(self.requests.get(), 1))
|
|
await asyncio.wait_for(channel.receive(), 1)
|
|
with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .06), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005):
|
|
self.tasks.append(asyncio.create_task(channel._recover()))
|
|
poll = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(poll[2], bodies[2])
|
|
await asyncio.sleep(.04)
|
|
self.assertTrue(self.requests.empty(), 'completed history occupied the spare recovery slot')
|
|
await channel.send(bodies[3], False)
|
|
original = await asyncio.wait_for(self.requests.get(), .5)
|
|
recovery = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(original[2], bodies[3])
|
|
self.assertEqual(recovery[2], bodies[3])
|
|
self.respond(recovery, b'n' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'n' * 40)
|
|
# Neither the old receiver nor the new original was cancelled.
|
|
self.assertTrue(self.resets.empty())
|
|
self.respond(poll, b'p' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'p' * 40)
|
|
self.respond(original, b'o' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'o' * 40)
|
|
self.assertEqual(channel.replay_rotations, 0)
|
|
|
|
async def test_busy_channel_recovers_old_request_despite_newer_replies(self):
|
|
channel = _HttpChannel(self.lane, 1, 'busy-recovery-test')
|
|
body = b'samekey!' + b't' * 32
|
|
original = None
|
|
recovered = asyncio.Event()
|
|
chatter_reply = asyncio.Event()
|
|
chatter_requests = 0
|
|
|
|
async def server():
|
|
nonlocal original, chatter_requests
|
|
while True:
|
|
request = await self.requests.get()
|
|
if request[2] == body:
|
|
if original is None:
|
|
original = request
|
|
else:
|
|
self.respond(request, b't' * 40)
|
|
else:
|
|
chatter_requests += 1
|
|
self.respond(request, b'c' * 40)
|
|
|
|
async def chatter():
|
|
index = 0
|
|
while True:
|
|
chatter_reply.clear()
|
|
await channel.send(b'samekey!' + struct.pack('<I', index) + b'c' * 28, False)
|
|
index += 1
|
|
await chatter_reply.wait()
|
|
await asyncio.sleep(.005)
|
|
|
|
async def receiver():
|
|
while True:
|
|
if await channel.receive() == b't' * 40:
|
|
recovered.set()
|
|
else:
|
|
chatter_reply.set()
|
|
|
|
with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .05), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .12), \
|
|
patch('proxy.cf_h2.REPLAY_HISTORY_PACKETS', 3):
|
|
await channel.send(body, False)
|
|
self.tasks.extend(asyncio.create_task(coro) for coro in (
|
|
server(), chatter(), receiver(), channel._recover()))
|
|
await asyncio.wait_for(recovered.wait(), .6)
|
|
self.assertGreater(chatter_requests, 3)
|
|
self.assertLessEqual(len(channel.replay_history), 3)
|
|
self.assertIn(original[1], self.transport.connection.streams)
|
|
self.assertTrue(self.resets.empty())
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_receiving_other_body_does_not_mask_overdue_request(self):
|
|
channel = _HttpChannel(self.lane, 1, 'parallel-body-recovery-test')
|
|
body = b'samekey!' + b't' * 32
|
|
await channel.send(body, False)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
await channel.send(b'samekey!' + b'b' * 32, False)
|
|
other = await asyncio.wait_for(self.requests.get(), 1)
|
|
self.respond(other, b'b' * 4, end=False, length=400)
|
|
|
|
async def body_progress():
|
|
for _ in range(90):
|
|
await asyncio.sleep(.01)
|
|
conn, sid, _ = other
|
|
conn.h2.send_data(sid, b'b' * 4)
|
|
self.flush(conn)
|
|
|
|
with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .03), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .07):
|
|
self.tasks.extend(asyncio.create_task(coro) for coro in (body_progress(), channel._recover()))
|
|
probe = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(probe[2], body)
|
|
self.assertEqual(channel.receiving_http, 1)
|
|
self.respond(probe, b't' * 40)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b't' * 40)
|
|
self.assertIn(original[1], self.transport.connection.streams)
|
|
self.assertIn(other[1], self.transport.connection.streams)
|
|
self.assertTrue(self.resets.empty())
|
|
|
|
async def test_silent_recovery_polls_refresh_without_new_native_input(self):
|
|
channel = _HttpChannel(self.lane, 1, 'refresh-recovery-test')
|
|
bodies = [b'samekey!' + bytes([index]) * 32 for index in range(2)]
|
|
for body in bodies:
|
|
await channel.send(body, False)
|
|
await asyncio.wait_for(self.requests.get(), 1)
|
|
with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_SLOT_SECONDS', .2), \
|
|
patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .02):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
latest = await asyncio.wait_for(self.requests.get(), .5)
|
|
older = await asyncio.wait_for(self.requests.get(), .5)
|
|
channel.sent_since[channel.replay_history[0].pending] = time.monotonic() - 1
|
|
refreshed = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(refreshed[2], bodies[0])
|
|
self.assertEqual(latest[2], bodies[0])
|
|
self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (latest[1], 8))
|
|
self.assertIn(older[1], self.transport.connection.streams)
|
|
self.assertEqual(len(channel.replay_pending), 2)
|
|
self.respond(refreshed)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), .5), b'r' * 40)
|
|
self.assertEqual(channel.replay_rotations, 1)
|
|
|
|
async def test_recovery_never_resets_replays_with_partially_received_bodies(self):
|
|
channel = _HttpChannel(self.lane, 1, 'protected-body-test')
|
|
bodies = [b'k' * 8 + bytes([index]) * 32 for index in range(3)]
|
|
for body in bodies[:2]:
|
|
await channel.send(body, False)
|
|
await asyncio.wait_for(self.requests.get(), 1)
|
|
with patch('proxy.cf_h2.REPLAY_REQUEST_SECONDS', .01), \
|
|
patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .005), \
|
|
patch('proxy.cf_h2.REPLAY_SLOT_SECONDS', .04):
|
|
recovery = asyncio.create_task(channel._recover())
|
|
self.tasks.append(recovery)
|
|
probes = [await asyncio.wait_for(self.requests.get(), 1) for _ in range(2)]
|
|
for probe in probes:
|
|
self.respond(probe, b'p' * 20, end=False, length=40)
|
|
await channel.send(bodies[2], False)
|
|
original = await asyncio.wait_for(self.requests.get(), 1)
|
|
await asyncio.sleep(.1)
|
|
self.assertEqual(channel.receiving_http, 2)
|
|
self.assertTrue(self.requests.empty())
|
|
self.assertTrue(self.resets.empty())
|
|
self.assertIn(original[1], self.transport.connection.streams)
|
|
# Finishing a partial response frees a slot without throwing its
|
|
# MTProto result away, and recovery can then use that slot.
|
|
conn, sid, _ = probes[0]
|
|
conn.h2.send_data(sid, b'p' * 20, end_stream=True)
|
|
self.flush(conn)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'p' * 40)
|
|
fresh_probe = await asyncio.wait_for(self.requests.get(), .5)
|
|
self.assertEqual(fresh_probe[2], bodies[2])
|
|
self.respond(fresh_probe)
|
|
self.assertEqual(await asyncio.wait_for(channel.receive(), 1), b'r' * 40)
|
|
self.assertEqual(channel.replay_rotations, 0)
|
|
|
|
async def test_preflight_and_posts_use_same_tls_connection(self):
|
|
preflight = asyncio.create_task(self.lane._preflight())
|
|
self.tasks.append(preflight)
|
|
request = await asyncio.wait_for(self.requests.get(), 2)
|
|
self.respond(request, b'', status=501)
|
|
await preflight
|
|
task, request = await self.post()
|
|
self.respond(request)
|
|
await task
|
|
self.assertEqual(self.lane.tcp_connections, 1)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_concurrent_waiters_respect_remote_stream_limit(self):
|
|
self.server_max_streams = 2
|
|
pending = [asyncio.create_task(self.lane._post(b'x' * 40, i)) for i in range(12)]
|
|
self.tasks.extend(pending)
|
|
for _ in range(6):
|
|
batch = [await asyncio.wait_for(self.requests.get(), 2) for _ in range(2)]
|
|
await asyncio.sleep(.01)
|
|
self.assertTrue(self.requests.empty())
|
|
for request in batch:
|
|
self.respond(request)
|
|
self.assertEqual(await asyncio.gather(*pending), [b'r' * 40] * 12)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_cancellation_resets_stream_and_reuses_server_capacity(self):
|
|
self.server_max_streams = 2
|
|
healthy, healthy_request = await self.post()
|
|
for _ in range(5):
|
|
task, request = await self.post()
|
|
task.cancel()
|
|
await asyncio.gather(task, return_exceptions=True)
|
|
self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
|
|
self.respond(healthy_request)
|
|
self.assertEqual(await asyncio.wait_for(healthy, .5), b'r' * 40)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_cancel_while_waiting_for_stream_slot_does_not_leak_capacity(self):
|
|
self.server_max_streams = 1
|
|
first, request = await self.post()
|
|
waiting = asyncio.create_task(self.lane._post(b'w' * 40, 2))
|
|
self.tasks.append(waiting)
|
|
await asyncio.sleep(.03)
|
|
self.assertTrue(self.requests.empty())
|
|
waiting.cancel()
|
|
await asyncio.gather(waiting, return_exceptions=True)
|
|
self.respond(request)
|
|
await first
|
|
next_task, next_request = await self.post()
|
|
self.respond(next_request)
|
|
await next_task
|
|
self.assertFalse(self.transport.connection.streams)
|
|
|
|
async def test_large_upload_and_download_cross_flow_control_windows(self):
|
|
body = b'x' * (256 * 1024)
|
|
task, request = await self.post(body)
|
|
self.respond(request, body)
|
|
self.assertEqual(await asyncio.wait_for(task, 2), body)
|
|
self.assertEqual(len(self.connections), 1)
|
|
|
|
async def test_stalled_body_consumer_does_not_block_other_streams(self):
|
|
# Keep a streaming response open without reading it. Its window must
|
|
# bound buffering while other responses can exceed a connection window.
|
|
opening = asyncio.create_task(self.lane.client.send(
|
|
self.lane.client.build_request('POST', 'https://' + self.lane.host + '/api', content=b'x' * 40),
|
|
stream=True))
|
|
self.tasks.append(opening)
|
|
request = await asyncio.wait_for(self.requests.get(), 2)
|
|
self.respond(request, b's' * (1024 * 1024))
|
|
response = await opening
|
|
fast, fast_request = await self.post()
|
|
self.respond(fast_request, b'f' * (4 * 1024 * 1024))
|
|
self.assertEqual(await asyncio.wait_for(fast, 2), b'f' * (4 * 1024 * 1024))
|
|
stream = self.transport.connection.streams[request[1]]
|
|
self.assertLessEqual(sum(len(data) for data, _ in stream.chunks), STREAM_RECEIVE_WINDOW)
|
|
await response.aclose()
|
|
self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
|
|
|
|
async def test_read_timeout_resets_only_its_stream(self):
|
|
self.lane.client.timeout = httpx.Timeout(2, read=.15)
|
|
stalled, request = await self.post()
|
|
healthy, healthy_request = await self.post()
|
|
self.respond(healthy_request)
|
|
self.assertEqual(await healthy, b'r' * 40)
|
|
with self.assertRaises(httpx.ReadTimeout):
|
|
await stalled
|
|
self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
|
|
self.assertIsNone(self.transport.connection.error)
|
|
|
|
async def test_peer_reset_does_not_break_other_streams(self):
|
|
bad, request = await self.post()
|
|
good, good_request = await self.post()
|
|
conn, sid, _ = request
|
|
conn.h2.reset_stream(sid, error_code=8)
|
|
self.flush(conn)
|
|
with self.assertRaises(httpx.RemoteProtocolError):
|
|
await bad
|
|
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()
|
|
for task, _ in tasks:
|
|
with self.assertRaises(httpx.HTTPError):
|
|
await asyncio.wait_for(task, .5)
|
|
task, request = await self.post()
|
|
self.respond(request)
|
|
self.assertEqual(await task, b'r' * 40)
|
|
self.assertEqual(len(self.connections), 2)
|
|
|
|
async def test_goaway_fails_pending_without_replaying_and_reconnects(self):
|
|
task, request = await self.post()
|
|
conn, _, _ = request
|
|
conn.h2.close_connection(error_code=0)
|
|
self.flush(conn)
|
|
with self.assertRaisesRegex(httpx.RemoteProtocolError, 'GOAWAY'):
|
|
await task
|
|
self.assertTrue(self.requests.empty())
|
|
next_task, next_request = await self.post()
|
|
self.respond(next_request)
|
|
await next_task
|
|
self.assertEqual(len(self.connections), 2)
|
|
|
|
async def test_close_stops_reader_and_wakes_pending_requests(self):
|
|
task, _ = await self.post()
|
|
connection = self.transport.connection
|
|
await self.lane.close()
|
|
with self.assertRaises(httpx.HTTPError):
|
|
await asyncio.wait_for(task, .5)
|
|
self.assertTrue(connection.reader_task.done())
|
|
self.assertTrue(connection.writer.is_closing())
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|