Code

Code/venv/test.py

@@ -0,0 +1,286 @@
+#!/usr/bin/env python3
+import argparse
+import asyncio
+import time
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+from aioquic.quic.configuration import QuicConfiguration
+from aioquic.quic.connection import QuicConnection
+from aioquic.tls import SessionTicket
+from aioquic.asyncio.client import connect
+from aioquic.quic.events import QuicEvent, HandshakeCompleted, ConnectionTerminated, ProtocolNegotiated, DatagramFrameReceived
+from aioquic.h0.connection import H0_ALPN
+from aioquic.h3.connection import H3_ALPN
+
+# RFC references:
+# - RFC 9000 QUIC transport: anti-amplification (<=3x before address validation), Initial >=1200B[1]
+# - RFC 9001 TLS over QUIC: handshake, tokens, address validation via Retry/token acceptance[2]
+# - RFC 8446 TLS 1.3 (used by QUIC): handshake and tokens at TLS layer carried in QUIC[3]
+
+@dataclass
+class FlowReport:
+    host: str
+    port: int
+    alpn: str
+    used_retry: bool
+    token_present: bool
+    address_validated_ts: Optional[float]
+    client_bytes_sent_before_av: int
+    server_bytes_recv_before_av: int
+    client_bytes_sent_total: int
+    server_bytes_recv_total: int
+    compliant_before_av: Optional[bool]
+    error: Optional[str]
+    handshake_completed: bool
+    protocol: Optional[str]
+
+class ByteMeteringTransport(asyncio.DatagramTransport):
+    def __init__(self, inner: asyncio.DatagramTransport):
+        self._inner = inner
+        self.bytes_sent = 0
+
+    def sendto(self, data: bytes, addr=None):
+        # UDP payload bytes we send
+        self.bytes_sent += len(data)
+        self._inner.sendto(data, addr)
+
+    # Delegate all other methods
+    def __getattr__(self, item):
+        return getattr(self._inner, item)
+
+class MeteringProtocol(asyncio.DatagramProtocol):
+    def __init__(self, quic: QuicConnection, on_datagram_recv):
+        self.quic = quic
+        self.on_datagram_recv = on_datagram_recv
+        self.transport: Optional[ByteMeteringTransport] = None
+
+    def connection_made(self, transport):
+        # Wrap real transport with metering
+        self.transport = ByteMeteringTransport(transport)
+        self.quic._network_path.send_datagram = self.transport.sendto  # ensure quic uses our wrapper
+
+    def datagram_received(self, data, addr):
+        # Count server->client bytes (payload)
+        self.on_datagram_recv(len(data))
+        self.quic.receive_datagram(data, time.time(), addr)
+
+    def error_received(self, exc):
+        pass
+
+    def connection_lost(self, exc):
+        pass
+
+async def quic_attempt(host: str, port: int, alpn: str, sni: Optional[str], timeout: float, token: Optional[bytes]) -> FlowReport:
+    cfg = QuicConfiguration(
+        is_client=True,
+        alpn_protocols=[alpn],
+        verify_mode=None,  # for probing; set to CERT_REQUIRED for strict TLS
+        server_name=sni or host,
+    )
+    # Inject token if provided (address validation token)
+    if token:
+        cfg.retry = True  # signal client can handle retry/token
+        cfg.token = token
+
+    # Bookkeeping
+    server_recv_total = 0
+    server_recv_before_av = 0
+    client_sent_before_av = 0
+    client_sent_total = 0
+    address_validated_ts = None
+    handshake_completed = False
+    proto = None
+    used_retry = False
+    token_present = token is not None
+    error = None
+
+    # Handler to count per-datagram bytes from server
+    def on_dgram_recv(n: int):
+        nonlocal server_recv_total
+        server_recv_total += n
+
+    loop = asyncio.get_event_loop()
+
+    # Low-level connect to intercept transport
+    try:
+        async with connect(
+            host,
+            port,
+            configuration=cfg,
+            create_protocol=None,
+            wait_connected=False,  # we want to observe early phases
+        ) as client:
+            # Access underlying connection and transport
+            quic: QuicConnection = client._quic
+            protocol = MeteringProtocol(quic, on_dgram_recv)
+
+            # Create custom UDP endpoint so we can wrap transport
+            transport, _ = await loop.create_datagram_endpoint(
+                lambda: protocol,
+                remote_addr=(host, port),
+            )
+            client._transport = protocol.transport  # swap client transport to our wrapper
+
+            start = time.time()
+            deadline = start + timeout
+
+            # Kick off handshake by sending Initial
+            quic.connect()
+            # Ensure Initial datagram >=1200B (aioquic does this by spec[1])
+            for datagram, addr in quic.datagrams_to_send(now=time.time()):
+                protocol.transport.sendto(datagram, addr)
+
+            # Event loop
+            while time.time() < deadline:
+                await asyncio.sleep(0.001)
+
+                # Pump outbound
+                for datagram, addr in quic.datagrams_to_send(now=time.time()):
+                    protocol.transport.sendto(datagram, addr)
+
+                # Process events
+                for event in quic.poll_events():
+                    # Negotiate protocol (ALPN)
+                    if isinstance(event, ProtocolNegotiated):
+                        proto = event.alpn_protocol
+                    # Handshake completed -> address validated by definition[2]
+                    if isinstance(event, HandshakeCompleted):
+                        handshake_completed = True
+                        if address_validated_ts is None:
+                            address_validated_ts = time.time()
+                    if isinstance(event, ConnectionTerminated):
+                        # End loop gracefully
+                        deadline = time.time()
+                        break
+
+                # Track bytes sent by client
+                client_sent_total = protocol.transport.bytes_sent
+
+                # Address Validation moment:
+                if address_validated_ts is None:
+                    # Before AV, snapshot budgets
+                    client_sent_before_av = client_sent_total
+                    server_recv_before_av = server_recv_total
+
+                # Give QUIC timers a tick
+                quic.handle_timer(time.time())
+
+                # Heuristic: detect Retry packet
+                # In aioquic, quic._retry_sent/_retry_received are internal; fallback: token becomes set after Retry
+                if quic._retry_received:  # type: ignore[attr-defined]
+                    used_retry = True
+
+                # Exit if handshake done and settled
+                if handshake_completed and (time.time() - address_validated_ts) > 0.1:
+                    break
+
+            # Close connection
+            quic.close(error_code=0x0)
+            for datagram, addr in quic.datagrams_to_send(now=time.time()):
+                protocol.transport.sendto(datagram, addr)
+            await asyncio.sleep(0.05)
+            transport.close()
+
+    except Exception as e:
+        error = str(e)
+
+    compliant = None
+    if address_validated_ts is None:
+        # If AV never happened, we still evaluate budget until timeout
+        if client_sent_before_av > 0:
+            compliant = server_recv_before_av <= 3 * client_sent_before_av  # RFC 9000 3× rule[1]
+    else:
+        # Evaluate until AV moment
+        if client_sent_before_av > 0:
+            compliant = server_recv_before_av <= 3 * client_sent_before_av  # RFC 9000 3× rule[1]
+
+    return FlowReport(
+        host=host,
+        port=port,
+        alpn=alpn,
+        used_retry=used_retry,
+        token_present=token_present,
+        address_validated_ts=address_validated_ts,
+        client_bytes_sent_before_av=client_sent_before_av,
+        server_bytes_recv_before_av=server_recv_before_av,
+        client_bytes_sent_total=client_sent_total,
+        server_bytes_recv_total=server_recv_total,
+        compliant_before_av=compliant,
+        error=error,
+        handshake_completed=handshake_completed,
+        protocol=proto,
+    )
+
+def print_report(title: str, report: FlowReport):
+    print(f"=== {title} ===")
+    print(f"Target: {report.host}:{report.port}, ALPN={report.alpn}")
+    print(f"Protocol negotiated: {report.protocol}")
+    print(f"Retry observed: {report.used_retry}, Token provided: {report.token_present}")
+    print(f"Handshake completed: {report.handshake_completed}, AddressValidated: {report.address_validated_ts is not None}")
+    print(f"Client bytes sent before AV: {report.client_bytes_sent_before_av}")
+    print(f"Server bytes received before AV: {report.server_bytes_recv_before_av}")
+    budget = 3 * report.client_bytes_sent_before_av
+    print(f"Anti-Amplification budget (3× client-before-AV): {budget}")
+    if report.compliant_before_av is None:
+        print("Compliance before AV: Unknown (insufficient data)")
+    else:
+        status = "COMPLIANT" if report.compliant_before_av else "VIOLATION"
+        print(f"Compliance before AV: {status}")
+    print(f"Totals: client_sent={report.client_bytes_sent_total}, server_recv={report.server_bytes_recv_total}")
+    if report.error:
+        print(f"Error: {report.error}")
+    print()
+
+async def main():
+    parser = argparse.ArgumentParser(description="QUIC Anti-Amplification & Retry Probe")
+    parser.add_argument("--host", required=True)
+    parser.add_argument("--port", type=int, default=443)
+    parser.add_argument("--alpn", default="h3")
+    parser.add_argument("--sni", default=None)
+    parser.add_argument("--timeout", type=float, default=5.0)
+    parser.add_argument("--no-retry", action="store_true")
+    args = parser.parse_args()
+
+    # First attempt: no token
+    report1 = await quic_attempt(
+        host=args.host,
+        port=args.port,
+        alpn=args.alpn,
+        sni=args.sni,
+        timeout=args.timeout,
+        token=None,
+    )
+    print_report("Attempt 1 (no token)", report1)
+
+    if report1.error:
+        return
+
+    if args.no_retry or not report1.used_retry:
+        # Either server didn't send Retry or user skipped token attempt
+        return
+
+    # If Retry observed, retrieve token from the connection state
+    # aioquic stores token in connection._token after Retry; we must run a second connection using it
+    # For simplicity, we reuse the token observed in first attempt if accessible; otherwise we try a second connect which should present the stored token via session ticket.
+    # Note: Access to private member is implementation-dependent; guarded with getattr.
+    retry_token: Optional[bytes] = None
+    try:
+        # Not ideal, but for probing we may access private attr
+        retry_token = getattr(report1, "retry_token", None)
+    except Exception:
+        retry_token = None
+
+    # Second attempt: with token if we have one (many servers encode address validation in Retry, token echoed by client)
+    report2 = await quic_attempt(
+        host=args.host,
+        port=args.port,
+        alpn=args.alpn,
+        sni=args.sni,
+        timeout=args.timeout,
+        token=retry_token,
+    )
+    print_report("Attempt 2 (with token if available)", report2)
+
+if __name__ == "__main__":
+    asyncio.run(main())