224 lines
10 KiB
Python
Executable File
224 lines
10 KiB
Python
Executable File
#!/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)
|