Files
RustProxy/scripts/performance/checks_proxy.py
T
2026-10-17 11:57:16 +02:00

136 lines
7.3 KiB
Python
Executable File

"""Proxy streaming, status, header, request-ID and WebSocket tests."""
import copy
import hashlib
from http.client import HTTPConnection, HTTPResponse
import json
import re
import socket
import time
from common import http
def run(check, rust, peer):
port = rust.reverse_port
for name in ["small.html", "large.bin", "huge.bin"]:
expected = (rust.fixture.site / name).read_bytes()
s, h, b = http(port, "/" + name)
check("proxy.get." + name, s == 200 and b == expected, status=s, bytes=len(b))
s, h, b = http(port, "/" + name, "HEAD")
check("proxy.head." + name, s == 200 and b == b"" and h.get("content-length") == str(len(expected)), status=s)
s, h, b = http(port, "/large.bin", headers={"Range": "bytes=0-127"})
check("proxy.range", s == 206 and b == bytes(range(128)), status=s, bytes=len(b))
s, h, b = http(port)
s, h, b = http(port, headers={"If-None-Match": h["etag"]})
check("proxy.conditional304", s == 304 and b == b"", status=s)
for status in [201, 400, 404, 500, 503]:
s, h, b = http(port, f"/status/{status}")
check(f"proxy.origin_status.{status}", s == status and b == b"origin status payload\n"
and h.get("x-origin-fixture") == "yes", status=s)
s, h, b = http(port, "/redirect")
check("proxy.redirect", s == 302 and h.get("location") == "/small.html", status=s)
payload = bytes(range(256)) * 4096
ids = []
for method in ["POST", "PUT"]:
s, h, b = http(port, "/echo?query=preserved", method,
{"Content-Type": "application/octet-stream", "X-Request-ID": "untrusted-client-id"}, payload)
data = json.loads(b)
request_id = data["headers"].get("x-request-id", "")
ids.append(request_id)
check("proxy.upload." + method, s == 200 and data["method"] == method
and data["path"] == "/echo?query=preserved" and data["body_bytes"] == len(payload)
and data["body_sha256"] == hashlib.sha256(payload).hexdigest(), status=s, bytes=data["body_bytes"])
check("proxy.request_id." + method, bool(re.fullmatch("[0-9a-f]{64}", request_id)),
valid_64_hex=bool(re.fullmatch("[0-9a-f]{64}", request_id)))
check("proxy.request_id_unique", len(set(ids)) == 2)
s, h, b = http(port, "/echo", "GET", {"Content-Type": "application/octet-stream"}, b"nonempty-get-body")
data = json.loads(b)
check("proxy.nonempty_get_body_preserved", s == 200 and data["body_bytes"] == 17
and data["body_sha256"] == hashlib.sha256(b"nonempty-get-body").hexdigest(), status=s)
peer.started_upload.clear()
first, last = b"first streamed chunk" * 64, b"final streamed chunk" * 64
with socket.create_connection(("127.0.0.1", port), timeout=5) as sock:
sock.sendall((f"POST /stream-upload HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\n"
"Transfer-Encoding: chunked\r\nContent-Type: application/octet-stream\r\n\r\n").encode())
sock.sendall(f"{len(first):x}\r\n".encode() + first + b"\r\n")
forwarded_before_end = peer.started_upload.wait(2)
sock.sendall(f"{len(last):x}\r\n".encode() + last + b"\r\n0\r\n\r\n")
response = HTTPResponse(sock)
response.begin()
data = json.loads(response.read())
check("proxy.chunked_streaming_upload", forwarded_before_end and response.status == 200
and data["body_sha256"] == hashlib.sha256(first + last).hexdigest(),
first_chunk_arrived_before_end=forwarded_before_end, status=response.status)
peer.release_stream.clear()
connection = HTTPConnection("127.0.0.1", port, timeout=3)
connection.request("GET", "/stream-download")
started = time.monotonic()
response = connection.getresponse()
first = response.read(5)
first_seconds = time.monotonic() - started
peer.release_stream.set()
last = response.read()
connection.close()
check("proxy.streaming_response", response.status == 200 and first == b"first" and last == b"last"
and first_seconds < 2, first_chunk_seconds=first_seconds)
peer.release_stream.clear()
connection = HTTPConnection("127.0.0.1", port, timeout=3)
connection.request("GET", "/stream-download")
response = connection.getresponse()
response.read(5)
connection.close()
peer.release_stream.set()
s, h, b = http(port)
check("proxy.client_disconnect_survives", s == 200, status=s)
before = peer.count("POST", "/no-replay")
s, h, b = http(port, "/no-replay", "POST", {"Content-Type": "application/octet-stream"}, b"write-once")
after = peer.count("POST", "/no-replay")
check("proxy.failed_post_not_replayed", s in (502, 503, 504) and after - before == 1,
status=s, upstream_writes=after - before)
config = copy.deepcopy(rust.base)
config["headers"] = {"request": [{"op": "set", "name": "X-Fixture-Global", "value": "global"}],
"response": [{"op": "set", "name": "X-Fixture-Global", "value": "response"}]}
config["vhosts"][0]["headers"] = {"request": [{"op": "set", "name": "X-Fixture-Vhost", "value": "vhost"}]}
config["vhosts"][0]["routes"] = [{"id": "echo", "path_prefix": "/echo", "headers": {
"request": [{"op": "set", "name": "X-Fixture-Route", "value": "route"}],
"response": [{"op": "set", "name": "X-Fixture-Route", "value": "response"}]}}]
rust.apply(config, "proxy-headers")
s, h, b = http(port, "/echo", headers={"Connection": "keep-alive, x-hop-fixture", "X-Hop-Fixture": "must-be-stripped"})
observed = json.loads(b)["headers"]
check("proxy.header_scopes", s == 200 and observed.get("x-fixture-global") == "global"
and observed.get("x-fixture-vhost") == "vhost" and observed.get("x-fixture-route") == "route"
and h.get("x-fixture-global") == "response" and h.get("x-fixture-route") == "response", status=s)
check("proxy.hop_by_hop_removed", "x-hop-fixture" not in observed)
s, h, b = http(port, "/echo", headers={"X-Forwarded-For": "203.0.113.22", "X-Real-IP": "203.0.113.22"})
observed = json.loads(b)["headers"]
check("proxy.untrusted_forwarded_ip", observed.get("x-real-ip") == "127.0.0.1"
and "203.0.113.22" not in observed.get("x-forwarded-for", ""),
observed_x_real_ip=observed.get("x-real-ip"), observed_x_forwarded_for=observed.get("x-forwarded-for"))
rust.reset()
with socket.create_connection(("127.0.0.1", port), timeout=4) as sock:
sock.sendall((f"GET /ws HTTP/1.1\r\nHost: 127.0.0.1:{port}\r\nUpgrade: websocket\r\n"
"Connection: Upgrade\r\nSec-WebSocket-Key: dGhlIHNhbXBsZSBub25jZQ==\r\n"
"Sec-WebSocket-Version: 13\r\n\r\n").encode())
buffer = b""
while b"\r\n\r\n" not in buffer:
piece = sock.recv(4096)
if not piece:
break
buffer += piece
upgraded = buffer.startswith(b"HTTP/1.1 101")
echoed = b""
if upgraded:
message, mask = b"hello", b"\x11\x22\x33\x44"
sock.sendall(bytes([0x81, 0x80 | len(message)]) + mask +
bytes(c ^ mask[i % 4] for i, c in enumerate(message)))
frame = sock.recv(2)
if len(frame) == 2:
echoed = sock.recv(frame[1] & 127)
check("proxy.websocket_echo", upgraded and echoed == b"hello", upgraded=upgraded, bytes=len(echoed))