#!/usr/bin/env python3 """A protocol peer, not a RustProxy emulator; used only as the test upstream. Performance comparisons use nginx as the upstream, never this Python fixture. """ from __future__ import annotations import base64 import collections import hashlib import json import socket import struct import threading import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from urllib.parse import urlsplit, parse_qs class Peer: def __init__(self, port, site): self.site = site self.counts = collections.Counter() self.received = {} self.lock = threading.Lock() self.started_upload = threading.Event() self.release_stream = threading.Event() peer = self class Handler(BaseHTTPRequestHandler): protocol_version = "HTTP/1.1" def log_message(self, *_): pass def emit(self, status, body=b"", headers=None): self.send_response(status) self.send_header("Content-Length", str(len(body))) if "Content-Type" not in (headers or {}): self.send_header("Content-Type", "application/octet-stream") for key, value in (headers or {}).items(): self.send_header(key, value) self.end_headers() if self.command != "HEAD" and body: self.wfile.write(body) def upload(self): chunks = [] if self.headers.get("Transfer-Encoding", "").lower() == "chunked": while True: length = int(self.rfile.readline().split(b";", 1)[0], 16) if length == 0: while self.rfile.readline().strip(): pass break payload = self.rfile.read(length) self.rfile.read(2) chunks.append(payload) peer.started_upload.set() else: remaining = int(self.headers.get("Content-Length", 0)) while remaining: payload = self.rfile.read(min(remaining, 1024)) if not payload: break chunks.append(payload) remaining -= len(payload) peer.started_upload.set() return b"".join(chunks) def do_POST(self): self.handle_method() def do_PUT(self): self.handle_method() def do_HEAD(self): self.handle_method() def do_GET(self): self.handle_method() def handle_method(self): try: self.route() except (BrokenPipeError, ConnectionResetError, socket.timeout): self.close_connection = True def route(self): parsed = urlsplit(self.path) path, query = parsed.path, parse_qs(parsed.query) with peer.lock: peer.counts[(self.command, path)] += 1 count = peer.counts[(self.command, path)] body = self.upload() if self.command in ("POST", "PUT") or self.headers.get("Content-Length") or self.headers.get("Transfer-Encoding") else b"" if path == "/no-replay": self.close_connection = True self.connection.shutdown(socket.SHUT_RDWR) self.connection.close() return if path in ("/echo", "/stream-upload"): data = {"method": self.command, "path": self.path, "headers": {key.lower(): value for key, value in self.headers.items()}, "body_bytes": len(body), "body_sha256": hashlib.sha256(body).hexdigest()} with peer.lock: peer.received[path] = data self.emit(200, json.dumps(data).encode(), {"Content-Type": "application/json", "Cache-Control": "no-store", "X-Origin-Fixture": "yes"}) return if path == "/stream-download": self.send_response(200) self.send_header("Transfer-Encoding", "chunked") self.send_header("Content-Type", "application/octet-stream") self.send_header("Cache-Control", "no-store") self.end_headers() self.wfile.write(b"5\r\nfirst\r\n") self.wfile.flush() peer.release_stream.wait(5) self.wfile.write(b"4\r\nlast\r\n0\r\n\r\n") self.wfile.flush() return if path == "/ws": key = self.headers.get("Sec-WebSocket-Key", "") accept = base64.b64encode(hashlib.sha1( (key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11").encode()).digest()).decode() self.send_response(101) self.send_header("Upgrade", "websocket") self.send_header("Connection", "Upgrade") self.send_header("Sec-WebSocket-Accept", accept) self.end_headers() self.wfile.flush() frame = self.rfile.read(2) if len(frame) == 2: length = frame[1] & 127 if length == 126: length = struct.unpack("!H", self.rfile.read(2))[0] mask = self.rfile.read(4) if frame[1] & 128 else b"\0\0\0\0" data = self.rfile.read(length) data = bytes(byte ^ mask[i % 4] for i, byte in enumerate(data)) self.wfile.write(bytes([0x81, len(data)]) + data) self.wfile.flush() self.close_connection = True return if path.startswith("/status/"): self.emit(int(path.rsplit("/", 1)[1]), b"origin status payload\n", {"Cache-Control": "no-store", "X-Origin-Fixture": "yes"}) return if path == "/redirect": self.emit(302, b"redirect body\n", {"Location": "/small.html", "Cache-Control": "no-store"}) return if path in ("/private", "/cookie", "/no-store", "/vary", "/fresh", "/revalidate", "/concurrent"): cache_control = "public, max-age=60" headers = {"Content-Type": "text/plain"} if body: headers["X-Received-Body-Bytes"] = str(len(body)) if path == "/private": cache_control = "private, max-age=60" if path == "/no-store": cache_control = "no-store" if path == "/cookie": headers["Set-Cookie"] = "application-session=fixture; Path=/" if path == "/vary": headers["Vary"] = "Accept-Language" if path == "/concurrent": time.sleep(.08) if path == "/revalidate": cache_control = "public, max-age=1, stale-if-error=10" headers["ETag"] = '"fixture-v1"' if self.headers.get("If-None-Match") == headers["ETag"]: self.emit(304, b"", {**headers, "Cache-Control": cache_control}) return data = f"{path}:count={count}:lang={self.headers.get('Accept-Language', '')}".encode() if path == "/revalidate": data = b"stable revalidated object" self.emit(200, data, {**headers, "Cache-Control": cache_control}) return if path == "/no-transform": self.emit(200, b"unchanged compression fixture " * 1024, {"Content-Type": "text/plain", "Cache-Control": "no-transform, public, max-age=60"}) return file = peer.site / path.lstrip("/") if not file.is_file() or ".." in file.parts: self.emit(404, b"fixture not found") return payload = file.read_bytes() tag = '"' + hashlib.sha256(payload).hexdigest() + '"' headers = {"Cache-Control": "public, max-age=60", "ETag": tag, "Accept-Ranges": "bytes", "X-Origin-Fixture": "yes"} headers["Content-Type"] = {".html": "text/html", ".txt": "text/plain", ".css": "text/css", ".json": "application/json"}.get(file.suffix, "application/octet-stream") if self.headers.get("If-None-Match") == tag: self.emit(304, b"", headers) return range_value = self.headers.get("Range") if range_value and self.headers.get("If-Range", tag) == tag: try: first, last = range_value.removeprefix("bytes=").split("-", 1) start = int(first) if first else max(0, len(payload) - int(last)) end = min(int(last), len(payload) - 1) if first and last else len(payload) - 1 if start >= len(payload) or end < start: self.emit(416, b"", {**headers, "Content-Range": f"bytes */{len(payload)}"}) else: self.emit(206, payload[start:end + 1], {**headers, "Content-Range": f"bytes {start}-{end}/{len(payload)}"}) return except ValueError: pass self.emit(200, payload, headers) self.server = ThreadingHTTPServer(("127.0.0.1", port), Handler) self.server.daemon_threads = True self.thread = threading.Thread(target=self.server.serve_forever, daemon=True) self.thread.start() def count(self, method, path): with self.lock: return self.counts[(method, path)] def close(self): self.release_stream.set() self.server.shutdown() self.server.server_close() self.thread.join(timeout=2)