Improve http client request handling

codegen/http/generate_http.py

@@ -92,6 +92,49 @@ def serialize_value(value):
         serialized = serialize_value(body)
         return json.dumps(serialized)
 
+    def _canonical_url(self, url: str) -> str:
+        prepared = requests.models.PreparedRequest()
+        prepared.prepare_url(url, None)
+        return str(prepared.url)
+
+    def _url_origin(self, url: str) -> str:
+        parts = urlsplit(url)
+        scheme = parts.scheme.lower()
+        port = parts.port if parts.port is not None else {"http": 80, "https": 443}.get(scheme)
+        host = (parts.hostname or "").lower()
+        return f"{scheme}://{host}:{port}"
+
+    def _post(self, url: str, headers: Dict[str, str], data: str) -> requests.Response:
+        current_url = self._canonical_url(url)
+        origin = self._url_origin(current_url)
+        redirects = 0
+        while True:
+            response = requests.post(
+                current_url,
+                headers=headers,
+                data=data,
+                timeout=self.default_timeout,
+                allow_redirects=False
+            )
+            if not 300 <= response.status_code < 400:
+                return response
+
+            redirects += 1
+            if redirects > MAX_REDIRECTS:
+                raise requests.TooManyRedirects(f"Exceeded {MAX_REDIRECTS} redirects")
+
+            location = response.headers.get("Location")
+            refusal = f"Not following redirect ({response.status_code}) to {location!r}"
+            if response.status_code not in (307, 308) or not location:
+                raise requests.RequestException(refusal)
+            try:
+                next_url = self._canonical_url(urljoin(current_url, location))
+            except (requests.RequestException, ValueError) as exc:
+                raise requests.RequestException(refusal) from exc
+            if self._url_origin(next_url) != origin:
+                raise requests.RequestException(refusal)
+            current_url = next_url
+
     def _request(self, url: str, body: Dict[str, Any], response_type: type) -> Any:
         \"\"\"Make a request to the Turnkey API.
 
@@ -117,12 +160,7 @@ def _request(self, url: str, body: Dict[str, Any], response_type: type) -> Any:
         }
 
         try:
-            response = requests.post(
-                full_url,
-                headers=headers,
-                data=body_str,
-                timeout=self.default_timeout
-            )
+            response = self._post(full_url, headers, body_str)
         except requests.RequestException as exc:
             raise TurnkeyNetworkError(
                 "Request failed",
@@ -253,12 +291,7 @@ def send_signed_request(self, signed_request: SignedRequest, response_type: type
         }
 
         try:
-            response = requests.post(
-                signed_request.url,
-                headers=headers,
-                data=signed_request.body,
-                timeout=self.default_timeout
-            )
+            response = self._post(signed_request.url, headers, signed_request.body)
         except requests.RequestException as exc:
             raise TurnkeyNetworkError(
                 "Signed request failed",
@@ -543,12 +576,13 @@ def main():
 
     # Build full output
     output = f"{COMMENT_HEADER}\n\n"
-    output += "import json\nimport time\nfrom typing import Any, Callable, Dict, Optional, TypeVar, overload\nimport requests\n"
+    output += "import json\nimport time\nfrom typing import Any, Callable, Dict, Optional, TypeVar, overload\nfrom urllib.parse import urljoin, urlsplit\nimport requests\n"
     output += "from turnkey_api_key_stamper import ApiKeyStamper\n"
     output += "from turnkey_sdk_types import *\n"
     output += "from ..version import VERSION\n\n"
     output += "T = TypeVar('T')\n\n"
     output += f"TERMINAL_ACTIVITY_STATUSES = {TERMINAL_ACTIVITY_STATUSES}\n\n"
+    output += "MAX_REDIRECTS = 10\n\n"
     output += client_code
 
     # Ensure output directory exists

packages/http/src/turnkey_http/generated/client.py

@@ -3,6 +3,7 @@
 import json
 import time
 from typing import Any, Callable, Dict, Optional, TypeVar, overload
+from urllib.parse import urljoin, urlsplit
 import requests
 from turnkey_api_key_stamper import ApiKeyStamper
 from turnkey_sdk_types import *
@@ -17,6 +18,8 @@
     "ACTIVITY_STATUS_REJECTED",
 ]
 
+MAX_REDIRECTS = 10
+
 
 class TurnkeyClient:
     """Turnkey API HTTP client with auto-generated methods."""
@@ -71,6 +74,53 @@ def serialize_value(value):
         serialized = serialize_value(body)
         return json.dumps(serialized)
 
+    def _canonical_url(self, url: str) -> str:
+        prepared = requests.models.PreparedRequest()
+        prepared.prepare_url(url, None)
+        return str(prepared.url)
+
+    def _url_origin(self, url: str) -> str:
+        parts = urlsplit(url)
+        scheme = parts.scheme.lower()
+        port = (
+            parts.port
+            if parts.port is not None
+            else {"http": 80, "https": 443}.get(scheme)
+        )
+        host = (parts.hostname or "").lower()
+        return f"{scheme}://{host}:{port}"
+
+    def _post(self, url: str, headers: Dict[str, str], data: str) -> requests.Response:
+        current_url = self._canonical_url(url)
+        origin = self._url_origin(current_url)
+        redirects = 0
+        while True:
+            response = requests.post(
+                current_url,
+                headers=headers,
+                data=data,
+                timeout=self.default_timeout,
+                allow_redirects=False,
+            )
+            if not 300 <= response.status_code < 400:
+                return response
+
+            redirects += 1
+            if redirects > MAX_REDIRECTS:
+                raise requests.TooManyRedirects(f"Exceeded {MAX_REDIRECTS} redirects")
+
+            location = response.headers.get("Location")
+            refusal = f"Not following redirect ({response.status_code}) to {location!r}"
+            if response.status_code not in (307, 308) or not location:
+                raise requests.RequestException(refusal)
+            try:
+                next_url = self._canonical_url(urljoin(current_url, location))
+            except (requests.RequestException, ValueError) as exc:
+                raise requests.RequestException(refusal) from exc
+            if self._url_origin(next_url) != origin:
+                raise requests.RequestException(refusal)
+            current_url = next_url
+
     def _request(self, url: str, body: Dict[str, Any], response_type: type) -> Any:
         """Make a request to the Turnkey API.
 
@@ -96,9 +146,7 @@ def _request(self, url: str, body: Dict[str, Any], response_type: type) -> Any:
         }
 
         try:
-            response = requests.post(
-                full_url, headers=headers, data=body_str, timeout=self.default_timeout
-            )
+            response = self._post(full_url, headers, body_str)
         except requests.RequestException as exc:
             raise TurnkeyNetworkError(
                 "Request failed", None, TurnkeyErrorCodes.NETWORK_ERROR, str(exc)
@@ -239,12 +287,7 @@ def send_signed_request(
         }
 
         try:
-            response = requests.post(
-                signed_request.url,
-                headers=headers,
-                data=signed_request.body,
-                timeout=self.default_timeout,
-            )
+            response = self._post(signed_request.url, headers, signed_request.body)
         except requests.RequestException as exc:
             raise TurnkeyNetworkError(
                 "Signed request failed", None, TurnkeyErrorCodes.NETWORK_ERROR, str(exc)

packages/http/tests/test_redirects.py

@@ -0,0 +1,123 @@
+import json
+from unittest.mock import patch
+
+import pytest
+from turnkey_api_key_stamper import TStamp
+from turnkey_http.generated.client import TurnkeyClient
+from turnkey_sdk_types import RequestType, SignedRequest, TurnkeyNetworkError
+
+BASE_URL = "https://api.example.com"
+ENDPOINT = "/public/v1/query/test"
+STAMP = TStamp(stamp_header_name="X-Stamp", stamp_header_value="stamp-value")
+
+
+class StaticStamper:
+    def stamp(self, content):
+        return STAMP
+
+
+class FakeResponse:
+    def __init__(self, status_code, location=None, payload=None):
+        self.status_code = status_code
+        self.headers = {"Location": location} if location else {}
+        self.ok = status_code < 400
+        self._payload = payload or {}
+        self.text = json.dumps(self._payload)
+        self.reason = ""
+
+    def json(self):
+        return self._payload
+
+
+@pytest.fixture
+def client():
+    return TurnkeyClient(BASE_URL, StaticStamper(), "org-id")
+
+
+@pytest.fixture
+def calls():
+    return []
+
+
+def send(client, calls, responses, signed=False):
+    def post(url, headers, data, timeout, allow_redirects):
+        calls.append((url, headers, data, allow_redirects))
+        return responses.pop(0)
+
+    with patch("turnkey_http.generated.client.requests.post", new=post):
+        if signed:
+            request = SignedRequest(
+                BASE_URL + ENDPOINT,
+                '{"organizationId": "org-id"}',
+                STAMP,
+                RequestType.QUERY,
+            )
+            return client.send_signed_request(request)
+        return client._request(ENDPOINT, {"organizationId": "org-id"}, dict)
+
+
+@pytest.mark.parametrize(
+    "signed,status_code,location",
+    [
+        (False, 307, "https://other.example.com" + ENDPOINT),
+        (True, 308, "http://api.example.com" + ENDPOINT),
+        (False, 307, "https://evil.example\\@api.example.com/steal"),
+        (False, 307, "https://[invalid" + ENDPOINT),
+        (False, 302, BASE_URL + ENDPOINT),
+        (False, 307, None),
+    ],
+)
+def test_disallowed_redirect_is_refused(client, calls, signed, status_code, location):
+    responses = [FakeResponse(status_code, location)]
+
+    with pytest.raises(TurnkeyNetworkError) as excinfo:
+        send(client, calls, responses, signed)
+
+    assert f"({status_code}) to {location!r}" in str(excinfo.value.cause)
+    assert len(calls) == 1
+    assert calls[0][3] is False
+
+
+def test_same_origin_redirect_preserves_request(client, calls):
+    responses = [
+        FakeResponse(307, BASE_URL + "/public/v1/query/other"),
+        FakeResponse(200, payload={"result": "ok"}),
+    ]
+
+    assert send(client, calls, responses) == {"result": "ok"}
+    assert calls[1][0] == BASE_URL + "/public/v1/query/other"
+    assert calls[1][1:] == calls[0][1:]
+
+
+def test_redirect_chain_stays_on_original_origin(client, calls):
+    responses = [
+        FakeResponse(307, BASE_URL + "/public/v1/query/hop"),
+        FakeResponse(308, "https://other.example.com" + ENDPOINT),
+    ]
+
+    with pytest.raises(TurnkeyNetworkError):
+        send(client, calls, responses)
+
+    assert [call[0] for call in calls] == [
+        BASE_URL + ENDPOINT,
+        BASE_URL + "/public/v1/query/hop",
+    ]
+
+
+@pytest.mark.parametrize("redirects,ok", [(10, True), (11, False)])
+def test_redirects_are_bounded_separately_from_initial_request(
+    client, calls, redirects, ok
+):
+    responses = [
+        FakeResponse(307, BASE_URL + f"/public/v1/query/hop{i}")
+        for i in range(redirects)
+    ]
+    responses.append(FakeResponse(200, payload={"result": "ok"}))
+
+    if ok:
+        assert send(client, calls, responses) == {"result": "ok"}
+        assert len(calls) == redirects + 1
+    else:
+        with pytest.raises(TurnkeyNetworkError):
+            send(client, calls, responses)
+        assert len(calls) == redirects

pyproject.toml

@@ -16,6 +16,9 @@ dev = [
     "pytest>=7.0.0",
 ]
 
+[tool.ruff]
+extend-exclude = ["*.md"]
+
 [tool.mypy]
 python_version = "3.10"
 warn_return_any = false