feat: Allow resolving IP to specific server on run-all-tests

Author: baalimago

README.md

@@ -310,6 +310,7 @@ This is needed when testing on edge, since `.../check` times out the workload.
 In conjunction with this, there is a convenience script which runs all tests each in a separate query.
 
 So you may also run `./run-all-tests-via-api.py --host <hostname> --port <port>`.
+If you need to hit a specific IP while preserving the original hostname (e.g., for edge testing or custom DNS), use `--resolve-ip <ip>` which is SNI-compatible for HTTPS and sets the HTTP `Host` header accordingly.
 This is intended to be run to validate package functionaltiy on edge, as each test becomes a separate workload.
 
 ### Notes

run-all-tests-via-api.py

@@ -6,41 +6,121 @@
 import sys
 import tempfile
 import time
+import socket
+import ssl
+import http.client
 from pathlib import Path
-from typing import List, Tuple
-from urllib.error import HTTPError, URLError
-from urllib.parse import quote
-from urllib.request import Request, urlopen
+from typing import List, Optional, Tuple
+from urllib.parse import quote, urlparse
 
 
-def fetch_tests(base_url: str) -> List[str]:
-    with urlopen(base_url + "/list", timeout=20) as resp:
-        if resp.status != 200:
-            raise RuntimeError(f"/list returned HTTP {resp.status}")
-        payload = json.loads(resp.read().decode("utf-8"))
-        tests = payload.get("tests", [])
-        if not isinstance(tests, list):
-            raise RuntimeError("Invalid /list payload: missing 'tests' list")
-        return [str(t) for t in tests]
+def _default_port_for_scheme(scheme: str) -> int:
+    return 443 if scheme == "https" else 80
 
 
-def run_single_test(base_url: str, test_name: str, timeout: float) -> Tuple[bool, str]:
+def _host_header(hostname: str, port: Optional[int], scheme: str) -> str:
+    default_port = _default_port_for_scheme(scheme)
+    if port and port != default_port:
+        return f"{hostname}:{port}"
+    return hostname
+
+
+class _ResolvedHTTPConnection(http.client.HTTPConnection):
+    def __init__(self, resolved_host: str, port: int, timeout: float) -> None:
+        super().__init__(host=resolved_host, port=port, timeout=timeout)
+        self._resolved_host = resolved_host
+
+    # For HTTP we just connect to the resolved host as usual.
+
+
+class _ResolvedHTTPSConnection(http.client.HTTPSConnection):
+    def __init__(self, resolved_host: str, port: int, timeout: float, *, server_hostname: str, context: Optional[ssl.SSLContext] = None) -> None:
+        # Pass the resolved host/IP to the base class so it doesn't try to resolve
+        super().__init__(host=resolved_host, port=port, timeout=timeout, context=context)
+        self._resolved_host = resolved_host
+        self._server_hostname = server_hostname
+
+    def connect(self) -> None:
+        # Largely mirrors the stdlib implementation but pins the TCP connect
+        # to the resolved host/IP and sets SNI to the original hostname.
+        self.sock = socket.create_connection((self._resolved_host, self.port), self.timeout, self.source_address)
+        if self._tunnel_host:
+            self._tunnel()
+        # Ensure we have a context
+        if self._context is None:
+            self._context = ssl.create_default_context()
+        # Enable hostname checking by default
+        self._context.check_hostname = True
+        self.sock = self._context.wrap_socket(self.sock, server_hostname=self._server_hostname)
+
+
+def http_get_text(url: str, *, resolve_ip: Optional[str], timeout: float) -> Tuple[int, str]:
+    """Perform a GET request with optional DNS override and SNI support.
+
+    If resolve_ip is given, connects to that IP, sets the Host header to the
+    original hostname, and (for HTTPS) uses SNI with the original hostname.
+    Returns (status_code, text_body).
+    """
+    parsed = urlparse(url)
+    scheme = (parsed.scheme or "http").lower()
+    hostname = parsed.hostname or ""
+    port = parsed.port or _default_port_for_scheme(scheme)
+    path = parsed.path or "/"
+    if parsed.query:
+        path += f"?{parsed.query}"
+
+    headers = {
+        "Accept": "*/*",
+    }
+
+    # If we override resolution, set the Host header explicitly
+    if resolve_ip:
+        headers["Host"] = _host_header(hostname, parsed.port, scheme)
+
+    try:
+        if resolve_ip:
+            if scheme == "https":
+                context = ssl.create_default_context()
+                conn = _ResolvedHTTPSConnection(resolve_ip, port, timeout, server_hostname=hostname, context=context)
+            else:
+                conn = _ResolvedHTTPConnection(resolve_ip, port, timeout)
+        else:
+            # No override — use stdlib conveniences
+            if scheme == "https":
+                conn = http.client.HTTPSConnection(hostname, port, timeout=timeout)
+            else:
+                conn = http.client.HTTPConnection(hostname, port, timeout=timeout)
+
+        conn.request("GET", path, headers=headers)
+        resp = conn.getresponse()
+        data = resp.read().decode("utf-8", errors="replace")
+        status = resp.status
+        conn.close()
+        return status, data
+    except Exception as e:
+        # Normalize into a network error string like urllib would give
+        return 0, f"Network error calling {url}: {e}"
+
+
+def fetch_tests(base_url: str, *, resolve_ip: Optional[str]) -> List[str]:
+    status, body = http_get_text(base_url + "/list", resolve_ip=resolve_ip, timeout=20)
+    if status != 200:
+        raise RuntimeError(f"/list returned HTTP {status}")
+    payload = json.loads(body)
+    tests = payload.get("tests", [])
+    if not isinstance(tests, list):
+        raise RuntimeError("Invalid /list payload: missing 'tests' list")
+    return [str(t) for t in tests]
+
+
+def run_single_test(base_url: str, test_name: str, timeout: float, *, resolve_ip: Optional[str]) -> Tuple[bool, str]:
     url = base_url + "/check/" + quote(test_name)
-    req = Request(url, method="GET")
     try:
         print(f"Checking: {url}")
-        with urlopen(req, timeout=timeout) as resp:
-            output = resp.read().decode("utf-8", errors="replace")
-            ok = resp.status == 200
-            return ok, output
-    except HTTPError as e:
-        try:
-            body = e.read().decode("utf-8", errors="replace")
-        except Exception:
-            body = str(e)
-        # 417 or 500 considered failure
-        return False, body
-    except URLError as e:
+        status, output = http_get_text(url, resolve_ip=resolve_ip, timeout=timeout)
+        ok = status == 200
+        return ok, output
+    except Exception as e:
         return False, f"Network error calling {url}: {e}"
 
 
@@ -72,13 +152,26 @@ def main() -> int:
         help="Timeout for each test in seconds (default: 30.0)",
     )
 
+    parser.add_argument(
+        "--resolve-ip",
+        default=os.environ.get("RESOLVE_IP"),
+        help=(
+            "Optional IP to resolve the server hostname to. "
+            "When set, connections go to this IP while preserving the original "
+            "hostname for HTTP Host and TLS SNI (SNI-compatible)."
+        ),
+    )
+
     args = parser.parse_args()
 
-    base_url = f"{args.host}:{args.port}"
+    host_value = args.host
+    if not (host_value.startswith("http://") or host_value.startswith("https://")):
+        host_value = "http://" + host_value
+    base_url = f"{host_value}:{args.port}"
     outdir = Path(args.outdir)
     outdir.mkdir(parents=True, exist_ok=True)
 
-    tests = fetch_tests(base_url)
+    tests = fetch_tests(base_url, resolve_ip=args.resolve_ip)
     if not tests:
         print("No tests returned by /list. Nothing to run.")
         return 0
@@ -92,7 +185,7 @@ def main() -> int:
 
     for idx, test in enumerate(tests, start=1):
         print(f"[{idx}/{len(tests)}] Running {test} ...", flush=True)
-        ok, output = run_single_test(base_url, test, timeout=args.test_timeout)
+        ok, output = run_single_test(base_url, test, timeout=args.test_timeout, resolve_ip=args.resolve_ip)
         # Write log
         safe_name = test.replace(os.sep, "_")
         log_path = outdir / f"{safe_name}.log"