diff --git a/blockrun_llm/solana_client.py b/blockrun_llm/solana_client.py index c6543c3..3717048 100644 --- a/blockrun_llm/solana_client.py +++ b/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": diff --git a/tests/unit/test_solana_media.py b/tests/unit/test_solana_media.py index f7a70b5..48b97e6 100644 --- a/tests/unit/test_solana_media.py +++ b/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()