fix(solana): async re-sign payment guard + poll_url host pinning

blockrun_llm/solana_client.py

@@ -1391,9 +1391,22 @@ def _absolute_url(self, url: str) -> str:
         configured ``api_url`` already includes the trailing ``/api`` so
         we strip it once to avoid ``/api/api/...``.
         """
+        base = self._api_url[: -len("/api")] if self._api_url.endswith("/api") else self._api_url
         if url.startswith("http://") or url.startswith("https://"):
+            # The poll loop sends (and re-signs) the wallet's PAYMENT-SIGNATURE
+            # against this URL, so an absolute poll_url is pinned to the API
+            # host+scheme — a gateway response pointing it elsewhere would leak
+            # the signed payment off-host.
+            poll, api = httpx.URL(url), httpx.URL(base)
+            if (poll.scheme, poll.host) != (api.scheme, api.host):
+                raise APIError(
+                    "Refusing an absolute poll_url on a different host/scheme than "
+                    f"the API ({poll.scheme}://{poll.host} != {api.scheme}://{api.host}); "
+                    "the signed payment header must not be sent off-host.",
+                    502,
+                    {"poll_url": url},
+                )
             return url
-        base = self._api_url[: -len("/api")] if self._api_url.endswith("/api") else self._api_url
         return f"{base}{url}"
 
     def _request_image_with_payment(
@@ -3504,9 +3517,22 @@ async def image_edit(
     def _absolute_url(self, url: str) -> str:
         """Resolve a server-supplied relative ``poll_url`` against the API host
         (``api_url`` already includes the trailing ``/api`` — strip it once)."""
+        base = self._api_url[: -len("/api")] if self._api_url.endswith("/api") else self._api_url
         if url.startswith("http://") or url.startswith("https://"):
+            # The poll loop sends (and re-signs) the wallet's PAYMENT-SIGNATURE
+            # against this URL, so an absolute poll_url is pinned to the API
+            # host+scheme — a gateway response pointing it elsewhere would leak
+            # the signed payment off-host.
+            poll, api = httpx.URL(url), httpx.URL(base)
+            if (poll.scheme, poll.host) != (api.scheme, api.host):
+                raise APIError(
+                    "Refusing an absolute poll_url on a different host/scheme than "
+                    f"the API ({poll.scheme}://{poll.host} != {api.scheme}://{api.host}); "
+                    "the signed payment header must not be sent off-host.",
+                    502,
+                    {"poll_url": url},
+                )
             return url
-        base = self._api_url[: -len("/api")] if self._api_url.endswith("/api") else self._api_url
         return f"{base}{url}"
 
     # ── Video / music / speech / enrollment / market data (async) ──────────
@@ -3986,8 +4012,23 @@ async def _request_image_with_payment(
             return probe.json()
 
         # Step 2: sign x402 SVM payload (reuse the encoded signature on polls).
-        payment_headers, cost_usd = await self._sign_payment_from_response(probe)
-        encoded_payment = payment_headers["PAYMENT-SIGNATURE"]
+        # Inlined rather than _sign_payment_from_response so the original payment
+        # terms are captured for the mid-poll re-sign guard below.
+        probe_payment_header = SolanaLLMClient._extract_payment_header(probe)
+        if not probe_payment_header:
+            raise PaymentError("402 response but no payment requirements found")
+        payment_required = decode_payment_required_header(probe_payment_header)
+        payment_payload_obj = await self._sign_payment(payment_required)
+        encoded_payment = encode_payment_signature_header(payment_payload_obj)
+        cost_usd = float(payment_payload_obj.accepted.amount) / 1e6
+        # Terms this job is authorized to pay — any mid-poll re-sign must match.
+        orig_amount = payment_payload_obj.accepted.amount
+        orig_pay_to = payment_payload_obj.accepted.pay_to
+        payment_headers = {
+            "Content-Type": "application/json",
+            "User-Agent": _get_user_agent(),
+            "PAYMENT-SIGNATURE": encoded_payment,
+        }
 
         # Step 3: submit with signature.
         submit_resp = await self._client.post(
@@ -4075,20 +4116,32 @@ async def _request_image_with_payment(
                 # Base VideoClient.
                 if resigns_left > 0:
                     resigns_left -= 1
+                    resign_payload = None
                     try:
                         challenge = await self._client.get(
                             poll_url,
                             headers={"User-Agent": _get_user_agent()},
                             timeout=eff_timeout,
                         )
-                        if challenge.status_code == 402:
-                            resign_headers, _ = await self._sign_payment_from_response(challenge)
-                            poll_headers["PAYMENT-SIGNATURE"] = resign_headers["PAYMENT-SIGNATURE"]
-                            continue
+                        resign_header = SolanaLLMClient._extract_payment_header(challenge)
+                        if challenge.status_code == 402 and resign_header:
+                            resign_required = decode_payment_required_header(resign_header)
+                            resign_payload = await self._sign_payment(resign_required)
                     except (PaymentError, httpx.HTTPError):
                         # Challenge GET or re-sign failed — surface the gateway's
                         # real 402 reason, not a network/signing error.
-                        pass
+                        resign_payload = None
+                    if resign_payload is not None:
+                        # Refuse a re-challenge that reprices or redirects the
+                        # payment vs. what this job originally authorized. This
+                        # PaymentError must propagate (NOT fall through to the
+                        # generic 402); the guard also pins the amount, so the
+                        # submit-time cost_usd stays correct for the ledger.
+                        _assert_same_payment_terms(resign_payload, orig_amount, orig_pay_to)
+                        poll_headers["PAYMENT-SIGNATURE"] = encode_payment_signature_header(
+                            resign_payload
+                        )
+                        continue
                 raise build_payment_rejected_error(poll_resp)
 
             if last_status == "failed":

tests/unit/test_solana_media.py

@@ -16,8 +16,19 @@
 import httpx
 import pytest
 
-from blockrun_llm.solana_client import SolanaLLMClient, _assert_same_payment_terms
-from blockrun_llm.types import (
+# Solana x402 extras (x402[svm]) require Python >= 3.10; skip the whole module
+# on 3.9, where they aren't installed and the codec stubs below have nothing to
+# patch. Mirrors test_solana_timeout_routing.py.
+pytest.importorskip("x402")
+pytest.importorskip("solders")
+
+from blockrun_llm.solana_client import (  # noqa: E402
+    AsyncSolanaLLMClient,
+    SolanaLLMClient,
+    _assert_same_payment_terms,
+)
+from blockrun_llm.types import (  # noqa: E402
+    APIError,
     MusicResponse,
     PaymentError,
     SpeechResponse,
@@ -234,3 +245,161 @@ def test_network_with_traversal_rejected(self) -> None:
         client = _make_client(lambda r: httpx.Response(500))
         with pytest.raises(ValueError, match="network"):
             client.rpc("../evil", "eth_blockNumber")
+
+
+# ---------------------------------------------------------------------------
+# poll_url host pinning — the signed PAYMENT-SIGNATURE must not go off-host
+# ---------------------------------------------------------------------------
+
+
+class TestPollUrlHostPin:
+    def test_relative_poll_url_resolved_to_api_host(self) -> None:
+        client = _make_client(lambda r: httpx.Response(500))
+        assert (
+            client._absolute_url("/api/v1/videos/generations/JOB")
+            == "https://sol.blockrun.ai/api/v1/videos/generations/JOB"
+        )
+
+    def test_absolute_same_host_passes(self) -> None:
+        client = _make_client(lambda r: httpx.Response(500))
+        url = "https://sol.blockrun.ai/api/v1/videos/generations/JOB"
+        assert client._absolute_url(url) == url
+
+    def test_absolute_cross_host_rejected(self) -> None:
+        client = _make_client(lambda r: httpx.Response(500))
+        with pytest.raises(APIError, match="off-host"):
+            client._absolute_url("https://evil.example.com/api/v1/videos/generations/JOB")
+
+    def test_absolute_http_downgrade_rejected(self) -> None:
+        client = _make_client(lambda r: httpx.Response(500))
+        with pytest.raises(APIError, match="off-host"):
+            client._absolute_url("http://sol.blockrun.ai/api/v1/videos/generations/JOB")
+
+
+# ---------------------------------------------------------------------------
+# Mid-poll re-sign — end-to-end through the poll loop (sync + async parity)
+# ---------------------------------------------------------------------------
+
+
+def _make_async_client(handler: Any) -> AsyncSolanaLLMClient:
+    with (
+        mock.patch("blockrun_llm.solana_client.register_exact_svm_client"),
+        mock.patch("blockrun_llm.solana_client._create_signer"),
+    ):
+        client = AsyncSolanaLLMClient(
+            private_key="bogus_signer_is_patched",
+            api_url="https://sol.blockrun.ai/api",
+            rpc_url="http://test",
+        )
+
+    class _FakePayload:
+        class accepted:
+            amount = "1000000"
+            pay_to = "GsbwXfJraMomNxBcpR3DBNxnKwZbyq7YCoDdSLDwzxdV"
+
+    client._x402_client = mock.MagicMock()
+    # Async _sign_payment awaits create_payment_payload — must return a coroutine.
+    client._x402_client.create_payment_payload = mock.AsyncMock(return_value=_FakePayload())
+    client._client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
+    client._address = "11111111111111111111111111111111"
+    return client
+
+
+def _resign_handler(signed_poll_codes: List[int]):
+    """Drive a video job through the mid-poll re-sign path.
+
+    probe → 402; signed POST → 202 + poll_url; each *signed* GET poll returns
+    the next code from ``signed_poll_codes`` (402 = settlement failed, 200 =
+    completed); an *unsigned* GET is the re-challenge and always hands back a
+    fresh 402 payment-required so the client re-signs.
+    """
+    pr = {"content-type": "application/json", "payment-required": "stub"}
+    completed = {
+        "status": "completed",
+        "created": 1,
+        "model": "xai/grok-imagine-video",
+        "data": [{"url": "https://cdn/v.mp4"}],
+    }
+    state = {"i": 0}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        has_sig = "PAYMENT-SIGNATURE" in request.headers
+        if request.method == "POST":
+            if not has_sig:  # unsigned probe
+                return httpx.Response(402, headers=pr, json={"error": "Payment Required"})
+            return httpx.Response(  # signed submit
+                202,
+                json={
+                    "id": "JOB",
+                    "poll_url": "/api/v1/videos/generations/JOB",
+                    "status": "queued",
+                },
+            )
+        if not has_sig:  # unsigned re-challenge → trigger a re-sign
+            return httpx.Response(402, headers=pr, json={"error": "Payment Required"})
+        code = signed_poll_codes[min(state["i"], len(signed_poll_codes) - 1)]
+        state["i"] += 1
+        if code == 200:
+            return httpx.Response(200, json=completed, headers={"content-type": "application/json"})
+        return httpx.Response(402, headers=pr, json={"error": "settlement failed"})
+
+    return handler
+
+
+_HELPER_KW: Dict[str, Any] = {
+    "poll_budget_seconds": 5.0,
+    "poll_interval_seconds": 0.001,
+    "max_resigns": 2,
+    "label": "Video generation",
+}
+_VIDEO_BODY = {"model": "xai/grok-imagine-video", "prompt": "a cat"}
+
+
+class TestResignEndToEnd:
+    def test_sync_resign_same_terms_then_completes(self) -> None:
+        # poll 402 (stale blockhash) → re-challenge → re-sign (same terms, guard
+        # passes) → next poll 200 completed.
+        client = _make_client(_resign_handler([402, 200]))
+        data = client._request_image_with_payment(
+            "/v1/videos/generations", dict(_VIDEO_BODY), **_HELPER_KW
+        )
+        assert data["data"][0]["url"] == "https://cdn/v.mp4"
+
+    def test_sync_resign_reprice_propagates(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        # A guard rejection on the re-signed challenge must propagate, NOT be
+        # swallowed by the re-sign try/except and masked as a generic 402.
+        monkeypatch.setattr(
+            "blockrun_llm.solana_client._assert_same_payment_terms",
+            mock.Mock(side_effect=PaymentError("repriced")),
+        )
+        client = _make_client(_resign_handler([402, 200]))
+        with pytest.raises(PaymentError, match="repriced"):
+            client._request_image_with_payment(
+                "/v1/videos/generations", dict(_VIDEO_BODY), **_HELPER_KW
+            )
+
+    async def test_async_resign_same_terms_then_completes(self) -> None:
+        client = _make_async_client(_resign_handler([402, 200]))
+        try:
+            data = await client._request_image_with_payment(
+                "/v1/videos/generations", dict(_VIDEO_BODY), **_HELPER_KW
+            )
+            assert data["data"][0]["url"] == "https://cdn/v.mp4"
+        finally:
+            await client._client.aclose()
+
+    async def test_async_resign_reprice_propagates(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        # Async parity with the sync guard: a re-price is rejected and the
+        # PaymentError propagates out of the poll loop.
+        monkeypatch.setattr(
+            "blockrun_llm.solana_client._assert_same_payment_terms",
+            mock.Mock(side_effect=PaymentError("repriced")),
+        )
+        client = _make_async_client(_resign_handler([402, 200]))
+        try:
+            with pytest.raises(PaymentError, match="repriced"):
+                await client._request_image_with_payment(
+                    "/v1/videos/generations", dict(_VIDEO_BODY), **_HELPER_KW
+                )
+        finally:
+            await client._client.aclose()