Relay the calling AAuth agent to the IPinfo API

.env.example

@@ -25,6 +25,6 @@ IPINFO_TOKEN=
 # Generate one with: openssl genpkey -algorithm ed25519
 # IPINFO_AAUTH_PROVIDER_KEY=
 
-# Public URL of this server for AAuth, required when IPINFO_AAUTH_PROVIDER_KEY is set. The IPinfo API fetches
-# this server's provider documents from it, and AAuth agents its resource metadata, so it must be reachable from the internet.
+# Public URL of this server for AAuth, required when IPINFO_AAUTH_PROVIDER_KEY is set.
+# Its AAuth documents are served there, so it must be reachable from the internet.
 # IPINFO_AAUTH_ISSUER=https://mcp.ipinfo.io

README.md

@@ -146,7 +146,13 @@ With the `http` transport, each request authenticates with its own `Authorizatio
 
 Setting `IPINFO_OAUTH_ISSUER` turns on OAuth for the `http` transport. The server then advertises the issuer at `/.well-known/oauth-protected-resource`, answers requests without a valid credential with `401 Unauthorized`, and verifies access tokens against the issuer's signing keys. A verified access token is sent to the IPinfo API in place of an API token. API tokens keep working as before.
 
-Setting `IPINFO_AAUTH_PROVIDER_KEY` turns on AAuth for the `http` transport. Generate a key with `openssl genpkey -algorithm ed25519`. The server then describes itself as an AAuth resource at `/.well-known/aauth-resource.json`, under the origin of `IPINFO_MCP_RESOURCE_URL` (or `IPINFO_AAUTH_ISSUER` when unset), pointing agents at `/aauth`. That endpoint serves a separate MCP server whose only tool is `ipinfo_lite_lookup`. Listing its tools needs no signature, since agents discover tools before signing a call, but calling one needs an AAuth signature; a refused one gets a `401 Unauthorized` naming the reason in its `Signature-Error` header. When `IPINFO_MCP_RESOURCE_URL` is set, only signatures made for that host are accepted. The server sends Lite lookups to the IPinfo API signed as its own agent, `aauth:ipinfo-mcp@<issuer host>`, and publishes the documents the API verifies them with at `/.well-known/aauth-agent.json` and `/.well-known/aauth-jwks.json`. `IPINFO_AAUTH_ISSUER` must be the server's public URL: `https`, no port, path, or trailing slash. The main endpoint is unaffected.
+Setting `IPINFO_AAUTH_PROVIDER_KEY` turns on AAuth for the `http` transport, so agents that sign their requests with AAuth can use `ipinfo_lite_lookup` without an IPinfo account. Generate the key with `openssl genpkey -algorithm ed25519`, and give every replica the same one. `IPINFO_AAUTH_ISSUER` must be the server's public URL: `https`, no port, path, or trailing slash. With AAuth on, the server also serves:
+
+- `/aauth`, the MCP endpoint for AAuth agents, whose only tool is `ipinfo_lite_lookup`. Listing its tools needs no signature; calling one needs an AAuth signature. When `IPINFO_MCP_RESOURCE_URL` is set, only signatures made for that host are accepted.
+- `/.well-known/aauth-resource.json`, which points agents at `/aauth`. It names the origin of `IPINFO_MCP_RESOURCE_URL`, or of `IPINFO_AAUTH_ISSUER` when that is unset.
+- `/.well-known/aauth-agent.json` and `/.well-known/aauth-jwks.json`, the documents of the server's own AAuth agent, `aauth:ipinfo-mcp@<issuer host>`.
+
+The main endpoint is unaffected.
 
 ## Feedback
 

src/ipinfo_mcp/aauth/agent.py

@@ -12,7 +12,6 @@
     AGENT_METADATA_DOCUMENT,
     AGENT_TOKEN_TYPE,
     HTTPS_PREFIX,
-    REQUIRED_COMPONENTS,
     is_server_identifier,
     sign,
     signature_base,
@@ -32,6 +31,11 @@
 
 SIGNATURE_LABEL = "sig"
 
+# Where requests to the IPinfo API relay the agent this server acts for: that agent's own
+# agent token. AAuth doesn't define relaying an agent's identity, so this is IPinfo's own
+# convention, honored by the API only from the intermediaries it trusts.
+AGENT_TOKEN_HEADER = "ipinfo-agent-token"
+
 
 class AgentProvider:
     """
@@ -121,18 +125,21 @@ def sign(self, request: httpx.Request) -> httpx.Request:
 
         Pass it as a request's httpx auth: httpx calls it once the request is built, so the
         authority and path signed are exactly the ones sent, IPv6 addresses included.
+
+        When the request relays the agent this server acts for, in AGENT_TOKEN_HEADER,
+        the signature covers that header too, so it can't be swapped on the way.
         """
         signature_key = ser({SIGNATURE_LABEL: (Token("jwt"), {"jwt": self.agent_token()})})
-        signature_params = ser_innerlist(([(name, {}) for name in REQUIRED_COMPONENTS], {"created": int(time.time())}))
-        base = signature_base(
-            [
-                ("@method", request.method),
-                ("@authority", request.headers["host"].lower()),
-                ("@path", request.url.raw_path.split(b"?", 1)[0].decode() or "/"),
-                ("signature-key", signature_key),
-            ],
-            signature_params,
-        )
+        components = [
+            ("@method", request.method),
+            ("@authority", request.headers["host"].lower()),
+            ("@path", request.url.raw_path.split(b"?", 1)[0].decode() or "/"),
+            ("signature-key", signature_key),
+        ]
+        if AGENT_TOKEN_HEADER in request.headers:
+            components.append((AGENT_TOKEN_HEADER, request.headers[AGENT_TOKEN_HEADER].strip()))
+        signature_params = ser_innerlist(([(name, {}) for name, _ in components], {"created": int(time.time())}))
+        base = signature_base(components, signature_params)
         signature = sign(base.encode(), self._agent_key, ALGORITHM)
 
         request.headers["Signature-Key"] = signature_key

src/ipinfo_mcp/aauth/lite.py

@@ -10,7 +10,7 @@
 from fastmcp.server.dependencies import get_http_request
 from mcp.types import ToolAnnotations
 
-from ipinfo_mcp.aauth.agent import AgentProvider
+from ipinfo_mcp.aauth.agent import AGENT_TOKEN_HEADER, AgentProvider
 from ipinfo_mcp.aauth.resource import Agent, scope_agent
 from ipinfo_mcp.cache import CachedResponse, IPCache
 from ipinfo_mcp.client import USER_AGENT
@@ -66,17 +66,21 @@ def http(self) -> httpx.AsyncClient:
             raise RuntimeError("Client not initialized. Use 'async with' context manager.")
         return self._http
 
-    async def lookup(self, ips: list[str]) -> dict[str, LiteResponse]:
+    async def lookup(self, ips: list[str], *, agent_token: str | None = None) -> dict[str, LiteResponse]:
         """
         Look up IPs on the Lite endpoint. Returns each IP's response, keyed by the IP.
 
+        agent_token is the agent token of the agent the lookups are for. It is relayed in
+        AGENT_TOKEN_HEADER, so the API can attribute them to that agent rather than to us.
+
         HTTP errors (401, 429, etc.) are propagated as httpx.HTTPStatusError, the first one that happens.
         """
         semaphore = asyncio.Semaphore(MAX_CONCURRENT_REQUESTS)
+        headers = {AGENT_TOKEN_HEADER: agent_token} if agent_token else None
 
         async def fetch(ip: str) -> LiteResponse:
             async with semaphore:
-                response = await self.http.get(f"/lite/{ip}")
+                response = await self.http.get(f"/lite/{ip}", headers=headers)
             if response.status_code == 401:
                 # The API names what it refused in our signature, that's a fault on our side
                 logger.warning("GET /lite refused signature_error=%s", response.headers.get("signature-error"))
@@ -189,7 +193,8 @@ async def ipinfo_lite_lookup(
     errors: dict[str, ErrorResponse] = {}
     if misses:
         try:
-            fetched = await client.lookup(misses)
+            # Relaying the calling agent lets the API tell agents apart, though they all reach it through us
+            fetched = await client.lookup(misses, agent_token=agent.token if agent else None)
         except httpx.HTTPStatusError as exc:
             logger.warning("ipinfo_lite_lookup api_error status=%d", exc.response.status_code)
             return handle_lite_error(exc)

tests/test_aauth_agent.py

@@ -144,6 +144,32 @@ def test_signs_the_host_header_sent(self) -> None:
         agent_jwk = decode_token(provider, provider.agent_token())[1]["cnf"]["jwk"]
         assert verify(base.encode(), signature, agent_jwk)
 
+    def test_signature_covers_the_relayed_agent(self) -> None:
+        provider = make_agent_provider(ISSUER)
+
+        request = provider.sign(
+            httpx.Request(
+                "GET", "https://api.ipinfo.io/lite/8.8.8.8", headers={"IPinfo-Agent-Token": "caller.agent.token"}
+            )
+        )
+
+        signature_key = request.headers["signature-key"]
+        params = request.headers["signature-input"].removeprefix("sig=")
+        assert params.startswith('("@method" "@authority" "@path" "signature-key" "ipinfo-agent-token")')
+        base = "\n".join(
+            [
+                '"@method": GET',
+                '"@authority": api.ipinfo.io',
+                '"@path": /lite/8.8.8.8',
+                f'"signature-key": {signature_key}',
+                '"ipinfo-agent-token": caller.agent.token',
+                f'"@signature-params": {params}',
+            ]
+        )
+        signature = parse(request.headers["signature"].encode(), tltype="dictionary")["sig"][0]
+        agent_jwk = decode_token(provider, provider.agent_token())[1]["cnf"]["jwk"]
+        assert verify(base.encode(), signature, agent_jwk)
+
     def test_sends_no_token(self) -> None:
         provider = make_agent_provider(ISSUER)
 

tests/test_aauth_app.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import time
 from collections.abc import AsyncIterator, Callable
@@ -235,6 +236,12 @@ async def test_connector_lists_unsigned_then_calls_signed(
 
         assert [tool["name"] for tool in listed["result"]["tools"]] == ["ipinfo_lite_lookup"]
         assert called["result"]["structuredContent"]["result"]["results"] == {"8.8.8.8": LITE_8888}
+        # The API is told which agent the lookup is for
+        lookup = httpx_mock.get_request()
+        assert lookup is not None
+        relayed = lookup.headers["IPinfo-Agent-Token"]
+        claims = json.loads(base64.urlsafe_b64decode(relayed.split(".")[1] + "=="))
+        assert claims["sub"] == AGENT
 
 
 class TestConfiguration:

tests/test_aauth_lite.py

@@ -65,6 +65,26 @@ async def test_looks_up_each_ip_signed_as_our_agent(self, client: LiteClient, ht
             assert "Authorization" not in request.headers
             assert request.headers["User-Agent"] == USER_AGENT
 
+    async def test_relays_the_agent_it_looks_up_for(self, client: LiteClient, httpx_mock: HTTPXMock) -> None:
+        httpx_mock.add_response(url=f"{BASE_URL}/lite/8.8.8.8", method="GET", json=LITE_8888)
+
+        await client.lookup(["8.8.8.8"], agent_token="caller.agent.token")
+
+        request = httpx_mock.get_request()
+        assert request is not None
+        assert request.headers["IPinfo-Agent-Token"] == "caller.agent.token"
+        assert '"ipinfo-agent-token"' in request.headers["Signature-Input"]
+
+    async def test_relays_nothing_without_an_agent(self, client: LiteClient, httpx_mock: HTTPXMock) -> None:
+        httpx_mock.add_response(url=f"{BASE_URL}/lite/8.8.8.8", method="GET", json=LITE_8888)
+
+        await client.lookup(["8.8.8.8"])
+
+        request = httpx_mock.get_request()
+        assert request is not None
+        assert "IPinfo-Agent-Token" not in request.headers
+        assert "ipinfo-agent-token" not in request.headers["Signature-Input"]
+
     @pytest.mark.parametrize("status_code", [401, 403, 429])
     async def test_propagates_errors(self, client: LiteClient, httpx_mock: HTTPXMock, status_code: int) -> None:
         httpx_mock.add_response(url=f"{BASE_URL}/lite/8.8.8.8", method="GET", status_code=status_code)