Regenerate clients with compression test support

features/generated-test/test-server

@@ -10,12 +10,15 @@ from __future__ import annotations
 
 import argparse
 import base64
+import ctypes
 import json
 import os
 import re
 import tempfile
 import threading
 import uuid
+import zlib
+from ctypes.util import find_library
 from datetime import datetime, timezone
 from http import HTTPStatus
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
@@ -44,6 +47,9 @@ SAFE_BROWSER_HEADERS = {
     "content-security-policy": "default-src 'none'; sandbox",
     "x-content-type-options": "nosniff",
 }
+ZSTD_CONTENTSIZE_UNKNOWN = (1 << 64) - 1
+ZSTD_CONTENTSIZE_ERROR = (1 << 64) - 2
+MAX_DECOMPRESSED_BODY_SIZE = 256 * 1024 * 1024
 
 
 class RecordingDatabase:
@@ -52,9 +58,11 @@ class RecordingDatabase:
         self.lock = threading.RLock()
         self.shards: dict[tuple[str, str], dict[str, Any]] = {}
         self.shard_paths: dict[tuple[str, str], Path] = {}
+        self.request_plans: dict[tuple[str, str, str], dict[str, Any]] = {}
         self.sessions: dict[str, dict[str, Any]] = {}
         self.fallback_consumed: set[tuple[str, str, str, int]] = set()
         self._load()
+        self._load_request_plans()
 
     def _load(self) -> None:
         manifest_path = self.root / "manifest.json"
@@ -70,6 +78,17 @@ class RecordingDatabase:
             self.shards[key] = shard
             self.shard_paths[key] = path
 
+    def _load_request_plans(self) -> None:
+        root = self.root.parent / "test-runner-data"
+        manifest_path = root / "manifest.json"
+        if not manifest_path.exists():
+            return
+        manifest = _read_json(manifest_path)
+        for item in manifest.get("scenarios", []):
+            plan = _read_json(root / item["file"])
+            key = (item["version"], item["feature"], item["scenario"])
+            self.request_plans[key] = plan.get("request", {})
+
     def start(self, version: str, feature: str, scenario: str, mode: str) -> dict[str, Any]:
         with self.lock:
             key = (version, feature)
@@ -91,6 +110,7 @@ class RecordingDatabase:
                 "key": key,
                 "scenario": scenario,
                 "recording": recording,
+                "request_plan": self.request_plans.get((version, feature, scenario)),
                 "cursor": 0,
                 "captures": [],
                 "frozen_at": frozen_at,
@@ -109,7 +129,7 @@ class RecordingDatabase:
                 if cursor >= len(interactions):
                     raise LookupError(f"Recording has no interaction #{cursor + 1}")
                 expected = interactions[cursor]
-                if not _requests_match(expected["request"], actual):
+                if not _requests_match(expected["request"], actual, session["request_plan"]):
                     raise RequestMismatchError(expected["request"], actual, cursor)
                 session["cursor"] += 1
                 return expected["response"]
@@ -118,8 +138,9 @@ class RecordingDatabase:
                 for recording in shard["recordings"]:
                     for index, interaction in enumerate(recording["interactions"]):
                         consumed_key = (version, feature, recording["scenario"], index)
+                        request_plan = self.request_plans.get((version, feature, recording["scenario"]))
                         if consumed_key not in self.fallback_consumed and _requests_match(
-                            interaction["request"], actual
+                            interaction["request"], actual, request_plan
                         ):
                             self.fallback_consumed.add(consumed_key)
                             return interaction["response"]
@@ -298,7 +319,13 @@ class TestRequestHandler(BaseHTTPRequestHandler):
 
     def _handle_api_request(self) -> None:
         body = self._read_body()
-        actual = _normalise_request(self.command, self.path, self.headers.get("content-type", ""), body)
+        actual = _normalise_request(
+            self.command,
+            self.path,
+            self.headers.get("content-type", ""),
+            self.headers.get("content-encoding", ""),
+            body,
+        )
         session_id = self.headers.get(SESSION_HEADER)
         if self.server.mode == "replay":
             response = self.server.database.replay(session_id, actual)
@@ -372,21 +399,50 @@ class TestRequestHandler(BaseHTTPRequestHandler):
             self.wfile.write(body)
 
 
-def _normalise_request(method: str, raw_path: str, content_type: str, body: bytes) -> dict[str, Any]:
+def _normalise_request(
+    method: str,
+    raw_path: str,
+    content_type: str,
+    content_encoding: str,
+    body: bytes,
+) -> dict[str, Any]:
     parsed = urlsplit(raw_path)
-    normalised_body = _normalise_body(body, content_type)
-    return {
+    compression = content_encoding.strip().casefold()
+    normalised_body = _normalise_body(body, content_type, compression)
+    request = {
         "method": method.upper(),
         "path": _normalise_path(parsed.path),
         "query": sorted([list(pair) for pair in parse_qsl(parsed.query, keep_blank_values=True)]),
         "content_type": _normalise_content_type(content_type, normalised_body),
         "body": normalised_body,
     }
+    if compression:
+        request["compression"] = compression
+    return request
 
 
-def _normalise_body(body: bytes, content_type: str) -> dict[str, Any]:
+def _normalise_body(body: bytes, content_type: str, compression: str = "") -> dict[str, Any]:
     if not body:
+        if compression in {"gzip", "deflate", "zstd1"}:
+            return {"type": "invalid-compression", "value": ""}
         return {"type": "empty", "value": None}
+    if compression == "zstd1":
+        decompressed = _decompress_zstd(body)
+        if decompressed is None:
+            return {
+                "type": "invalid-compression",
+                "value": base64.b64encode(body).decode("ascii"),
+            }
+        body = decompressed
+    elif compression in {"gzip", "deflate"}:
+        wbits = zlib.MAX_WBITS | 16 if compression == "gzip" else zlib.MAX_WBITS
+        decompressed = _decompress_zlib(body, wbits=wbits)
+        if decompressed is None:
+            return {
+                "type": "invalid-compression",
+                "value": base64.b64encode(body).decode("ascii"),
+            }
+        body = decompressed
     media_type = _media_type(content_type)
     text = body.decode("utf-8", errors="surrogateescape")
     if media_type.endswith("json"):
@@ -401,6 +457,92 @@ def _normalise_body(body: bytes, content_type: str) -> dict[str, Any]:
     return {"type": "text", "value": text}
 
 
+def _decompress_zlib(
+    body: bytes,
+    *,
+    wbits: int,
+    max_size: int = MAX_DECOMPRESSED_BODY_SIZE,
+) -> bytes | None:
+    """Decompress exactly one complete zlib stream with bounded output."""
+    decompressor = zlib.decompressobj(wbits)
+    output = bytearray()
+    pending = body
+    try:
+        while pending:
+            remaining = max_size - len(output) + 1
+            chunk = decompressor.decompress(pending, remaining)
+            output.extend(chunk)
+            if len(output) > max_size:
+                return None
+            unconsumed = decompressor.unconsumed_tail
+            if not unconsumed:
+                break
+            if len(unconsumed) == len(pending) and not chunk:
+                return None
+            pending = unconsumed
+    except zlib.error:
+        return None
+    if not decompressor.eof or decompressor.unused_data or decompressor.unconsumed_tail:
+        return None
+    return bytes(output)
+
+
+def _decompress_zstd(body: bytes) -> bytes | None:
+    """Decompress one complete Zstandard frame with the platform libzstd."""
+    candidates = (
+        find_library("zstd"),
+        "libzstd.so.1",
+        "libzstd.dylib",
+        "/opt/homebrew/lib/libzstd.dylib",
+        "/usr/local/lib/libzstd.dylib",
+        "libzstd.dll",
+        "zstd.dll",
+    )
+    library = None
+    for candidate in dict.fromkeys(item for item in candidates if item):
+        try:
+            library = ctypes.CDLL(candidate)
+            break
+        except OSError:
+            continue
+    if library is None:
+        return None
+
+    try:
+        library.ZSTD_isError.argtypes = [ctypes.c_size_t]
+        library.ZSTD_isError.restype = ctypes.c_uint
+        library.ZSTD_findFrameCompressedSize.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+        library.ZSTD_findFrameCompressedSize.restype = ctypes.c_size_t
+        library.ZSTD_getFrameContentSize.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+        library.ZSTD_getFrameContentSize.restype = ctypes.c_ulonglong
+        library.ZSTD_decompressBound.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+        library.ZSTD_decompressBound.restype = ctypes.c_ulonglong
+        library.ZSTD_decompress.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_void_p, ctypes.c_size_t]
+        library.ZSTD_decompress.restype = ctypes.c_size_t
+
+        source = ctypes.create_string_buffer(body)
+        source_pointer = ctypes.cast(source, ctypes.c_void_p)
+        compressed_size = library.ZSTD_findFrameCompressedSize(source_pointer, len(body))
+        if library.ZSTD_isError(compressed_size) or compressed_size != len(body):
+            return None
+
+        capacity = library.ZSTD_getFrameContentSize(source_pointer, len(body))
+        if capacity == ZSTD_CONTENTSIZE_ERROR:
+            return None
+        if capacity == ZSTD_CONTENTSIZE_UNKNOWN:
+            capacity = library.ZSTD_decompressBound(source_pointer, len(body))
+        if capacity > MAX_DECOMPRESSED_BODY_SIZE:
+            return None
+
+        destination = ctypes.create_string_buffer(max(1, capacity))
+        decompressed_size = library.ZSTD_decompress(destination, capacity, source_pointer, len(body))
+        if library.ZSTD_isError(decompressed_size) or decompressed_size > capacity:
+            return None
+        return destination.raw[:decompressed_size]
+    except (AttributeError, OSError, OverflowError, TypeError):
+        return None
+
+
 def _normalise_json(value: Any) -> Any:
     if isinstance(value, dict):
         return {key: _normalise_json(item) for key, item in value.items()}
@@ -415,13 +557,33 @@ def _normalise_content_type(content_type: str, body: dict[str, Any]) -> str:
     return "" if body["type"] == "empty" else _media_type(content_type)
 
 
-def _requests_match(expected: dict[str, Any], actual: dict[str, Any]) -> bool:
+def _requests_match(
+    expected: dict[str, Any],
+    actual: dict[str, Any],
+    request_plan: dict[str, Any] | None = None,
+) -> bool:
     comparable_fields = ("method", "path", "query", "content_type")
     if any(expected[field] != actual[field] for field in comparable_fields):
         return False
+    expected_compression = expected.get("compression")
+    if request_plan and _request_matches_plan(expected, request_plan):
+        expected_compression = request_plan.get("compression", expected_compression)
+    if expected_compression is not None and actual.get("compression", "") != expected_compression:
+        return False
     return _bodies_match(expected["body"], actual["body"])
 
 
+def _request_matches_plan(request: dict[str, Any], plan: dict[str, Any]) -> bool:
+    if request["method"] != plan.get("method"):
+        return False
+    path = plan.get("path")
+    if not path:
+        return False
+    parts = re.split(r"(\{[^/{}]+\})", path)
+    pattern = "".join(r"[^/]+" if part.startswith("{") else re.escape(part) for part in parts)
+    return re.fullmatch(pattern, request["path"]) is not None
+
+
 def _bodies_match(expected: dict[str, Any], actual: dict[str, Any]) -> bool:
     if expected == actual:
         return True