mirror of
https://github.com/Flowseal/tg-ws-proxy.git
synced 2026-10-21 11:08:44 +00:00
859 lines
42 KiB
Python
859 lines
42 KiB
Python
"""Wire-level regressions: real TLS sockets, ALPN and HTTP/2 flow control."""
|
|
import asyncio
|
|
import datetime
|
|
import ipaddress
|
|
import ssl
|
|
import struct
|
|
import tempfile
|
|
import time
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import 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.utils import PROTO_TAG_INTERMEDIATE
|
|
|
|
|
|
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_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_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' * (128 * 1024))
|
|
self.assertEqual(await asyncio.wait_for(fast, .5), b'f' * (128 * 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_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()
|