Add token in cache key

src/ipinfo_mcp/cache.py

@@ -1,3 +1,5 @@
+import hashlib
+
 from ipinfo_mcp.types import LiteResponse, LookupResponse, ResproxyResponse
 
 # Any single-IP response that can be cached
@@ -6,35 +8,50 @@
 
 class IPCache:
     """
-    In-memory cache with namespace support.
+    In-memory cache of per-IP responses, scoped by token and namespace.
 
     Namespaces: "lite", "lookup", "resproxy".
+
+    Entries are scoped also using the token that fetched them. The cache is process
+    wide, so if we don't use the token as part of the key we would hand one token's
+    results to another.
     """
 
     def __init__(self) -> None:
-        self._store: dict[tuple[str, str], CachedResponse] = {}
-
-    def get(self, namespace: str, ip: str) -> CachedResponse | None:
-        """Get cached data for an IP in a namespace. Returns None on miss."""
-        return self._store.get((namespace, ip))
-
-    def put(self, namespace: str, ip: str, data: CachedResponse) -> None:
-        """Store data for an IP in a namespace."""
-        self._store[(namespace, ip)] = data
-
-    def get_many(self, namespace: str, ips: list[str]) -> tuple[dict[str, CachedResponse], list[str]]:
-        """Look up multiple IPs. Returns (cached_results, cache_misses)."""
+        self._store: dict[tuple[str, str, str], CachedResponse] = {}
+
+    @staticmethod
+    def _hash_token(token: str) -> str:
+        """
+        We don't want to store the token directly, though we still need
+        to scope by token so we derive an hash from it and use it as part
+        of the hash key.
+        """
+        return hashlib.sha256(token.encode()).hexdigest()[:32]
+
+    def get(self, token: str, namespace: str, ip: str) -> CachedResponse | None:
+        """Get this token's cached data for an IP in a namespace. Returns None on miss."""
+        return self._store.get((self._hash_token(token), namespace, ip))
+
+    def put(self, token: str, namespace: str, ip: str, data: CachedResponse) -> None:
+        """Store data for an IP in a namespace, scoped to this token."""
+        self._store[(self._hash_token(token), namespace, ip)] = data
+
+    def get_many(self, token: str, namespace: str, ips: list[str]) -> tuple[dict[str, CachedResponse], list[str]]:
+        """Look up multiple IPs for this token. Returns (cached_results, cache_misses)."""
+        token_hash = self._hash_token(token)
         cached: dict[str, CachedResponse] = {}
         misses: list[str] = []
         for ip in ips:
-            data = self._store.get((namespace, ip))
+            data = self._store.get((token_hash, namespace, ip))
             if data is not None:
                 cached[ip] = data
             else:
                 misses.append(ip)
         return cached, misses
 
-    def put_many(self, namespace: str, items: dict[str, CachedResponse]) -> None:
-        """Store multiple IP results in a namespace."""
+    def put_many(self, token: str, namespace: str, items: dict[str, CachedResponse]) -> None:
+        """Store multiple IP results in a namespace, scoped to this token."""
+        token_hash = self._hash_token(token)
         for ip, data in items.items():
-            self._store[(namespace, ip)] = data
+            self._store[(token_hash, namespace, ip)] = data

src/ipinfo_mcp/tools/asn.py

@@ -74,7 +74,7 @@ async def ipinfo_asn(
 
     page_ips, pagination = paginate_ips(valid_ips, page, page_size)
 
-    cached, misses = cache.get_many(namespace, page_ips)
+    cached, misses = cache.get_many(token, namespace, page_ips)
 
     api_calls = 0
     if misses:
@@ -84,7 +84,7 @@ async def ipinfo_asn(
             api_calls = 1
             for key, data in fetched.items():
                 ip = key.split("/", 1)[1]
-                cache.put(namespace, ip, data)
+                cache.put(token, namespace, ip, data)
                 cached[ip] = data
         except httpx.HTTPStatusError as exc:
             logger.warning("ipinfo_asn api_error status=%d", exc.response.status_code)

src/ipinfo_mcp/tools/geolocate.py

@@ -80,7 +80,7 @@ async def ipinfo_geolocate(
 
     page_ips, pagination = paginate_ips(valid_ips, page, page_size)
 
-    cached, misses = cache.get_many(namespace, page_ips)
+    cached, misses = cache.get_many(token, namespace, page_ips)
 
     api_calls = 0
     if misses:
@@ -90,7 +90,7 @@ async def ipinfo_geolocate(
             api_calls = 1
             for key, data in fetched.items():
                 ip = key.split("/", 1)[1]
-                cache.put(namespace, ip, data)
+                cache.put(token, namespace, ip, data)
                 cached[ip] = data
         except httpx.HTTPStatusError as exc:
             logger.warning("ipinfo_geolocate api_error status=%d", exc.response.status_code)

src/ipinfo_mcp/tools/lookup.py

@@ -65,7 +65,7 @@ async def ipinfo_lookup(
     page_ips, pagination = paginate_ips(valid_ips, page, page_size)
 
     # Check cache for this page's IPs only
-    cached, misses = cache.get_many(namespace, page_ips)
+    cached, misses = cache.get_many(token, namespace, page_ips)
 
     api_calls = 0
     if misses:
@@ -75,7 +75,7 @@ async def ipinfo_lookup(
             api_calls = 1
             for key, data in fetched.items():
                 ip = key.split("/", 1)[1]
-                cache.put(namespace, ip, data)
+                cache.put(token, namespace, ip, data)
                 cached[ip] = data
         except httpx.HTTPStatusError as exc:
             logger.warning("ipinfo_lookup api_error status=%d", exc.response.status_code)

src/ipinfo_mcp/tools/privacy.py

@@ -74,7 +74,7 @@ async def ipinfo_check_privacy(
     page_ips, pagination = paginate_ips(valid_ips, page, page_size)
 
     # Check cache for this page's IPs only
-    cached, misses = cache.get_many(namespace, page_ips)
+    cached, misses = cache.get_many(token, namespace, page_ips)
 
     api_calls = 0
     if misses:
@@ -84,7 +84,7 @@ async def ipinfo_check_privacy(
             api_calls = 1
             for key, data in fetched.items():
                 ip = key.split("/", 1)[1]
-                cache.put(namespace, ip, data)
+                cache.put(token, namespace, ip, data)
                 cached[ip] = data
         except httpx.HTTPStatusError as exc:
             logger.warning("ipinfo_check_privacy api_error status=%d", exc.response.status_code)

src/ipinfo_mcp/tools/resproxy.py

@@ -70,7 +70,7 @@ async def ipinfo_check_residential_proxy(
 
     page_ips, pagination = paginate_ips(valid_ips, page, page_size)
 
-    cached, misses = cache.get_many(namespace, page_ips)
+    cached, misses = cache.get_many(token, namespace, page_ips)
 
     api_calls = 0
     if misses:
@@ -80,7 +80,7 @@ async def ipinfo_check_residential_proxy(
             api_calls = 1
             for key, data in fetched.items():
                 ip = key.split("/", 1)[1]
-                cache.put(namespace, ip, data)
+                cache.put(token, namespace, ip, data)
                 cached[ip] = data
         except httpx.HTTPStatusError as exc:
             logger.warning("ipinfo_check_residential_proxy api_error status=%d", exc.response.status_code)

tests/test_cache.py

@@ -1,69 +1,103 @@
 from ipinfo_mcp.cache import IPCache
 
+TOKEN = "token_a"
+OTHER_TOKEN = "token_b"
+
 
 class TestIPCachePutGet:
     def test_get_returns_none_for_missing_key(self) -> None:
         cache = IPCache()
-        assert cache.get("lite", "8.8.8.8") is None
+        assert cache.get(TOKEN, "lite", "8.8.8.8") is None
 
     def test_put_and_get_single_entry(self) -> None:
         cache = IPCache()
         data = {"ip": "8.8.8.8", "country": "US"}
-        cache.put("lite", "8.8.8.8", data)
-        assert cache.get("lite", "8.8.8.8") == data
+        cache.put(TOKEN, "lite", "8.8.8.8", data)
+        assert cache.get(TOKEN, "lite", "8.8.8.8") == data
 
     def test_namespaces_are_isolated(self) -> None:
         cache = IPCache()
         lite_data = {"ip": "8.8.8.8", "country": "US"}
         lookup_data = {"ip": "8.8.8.8", "city": "Mountain View"}
-        cache.put("lite", "8.8.8.8", lite_data)
-        cache.put("lookup", "8.8.8.8", lookup_data)
+        cache.put(TOKEN, "lite", "8.8.8.8", lite_data)
+        cache.put(TOKEN, "lookup", "8.8.8.8", lookup_data)
 
-        assert cache.get("lite", "8.8.8.8") == lite_data
-        assert cache.get("lookup", "8.8.8.8") == lookup_data
+        assert cache.get(TOKEN, "lite", "8.8.8.8") == lite_data
+        assert cache.get(TOKEN, "lookup", "8.8.8.8") == lookup_data
 
     def test_put_overwrites_existing(self) -> None:
         cache = IPCache()
-        cache.put("lite", "8.8.8.8", {"old": True})
-        cache.put("lite", "8.8.8.8", {"new": True})
-        assert cache.get("lite", "8.8.8.8") == {"new": True}
+        cache.put(TOKEN, "lite", "8.8.8.8", {"old": True})
+        cache.put(TOKEN, "lite", "8.8.8.8", {"new": True})
+        assert cache.get(TOKEN, "lite", "8.8.8.8") == {"new": True}
+
+
+class TestIPCacheTokenIsolation:
+    def test_other_token_does_not_see_entry(self) -> None:
+        """One token's results must never be served to another token."""
+        cache = IPCache()
+        cache.put(TOKEN, "resproxy", "8.8.8.8", {"service": "NordVPN"})
+
+        assert cache.get(OTHER_TOKEN, "resproxy", "8.8.8.8") is None
+
+    def test_other_token_reports_a_miss(self) -> None:
+        cache = IPCache()
+        cache.put(TOKEN, "lite", "8.8.8.8", {"ip": "8.8.8.8"})
+
+        cached, misses = cache.get_many(OTHER_TOKEN, "lite", ["8.8.8.8"])
+        assert cached == {}
+        assert misses == ["8.8.8.8"]
+
+    def test_tokens_hold_independent_entries_for_same_ip(self) -> None:
+        cache = IPCache()
+        cache.put(TOKEN, "resproxy", "8.8.8.8", {})
+        cache.put(OTHER_TOKEN, "resproxy", "8.8.8.8", {"service": "NordVPN"})
+
+        assert cache.get(TOKEN, "resproxy", "8.8.8.8") == {}
+        assert cache.get(OTHER_TOKEN, "resproxy", "8.8.8.8") == {"service": "NordVPN"}
+
+    def test_store_does_not_hold_plaintext_tokens(self) -> None:
+        cache = IPCache()
+        cache.put(TOKEN, "lite", "8.8.8.8", {"ip": "8.8.8.8"})
+
+        assert all(TOKEN not in key for key in cache._store)
 
 
 class TestIPCacheGetMany:
     def test_all_cached(self) -> None:
         cache = IPCache()
-        cache.put("lite", "8.8.8.8", {"ip": "8.8.8.8"})
-        cache.put("lite", "1.1.1.1", {"ip": "1.1.1.1"})
+        cache.put(TOKEN, "lite", "8.8.8.8", {"ip": "8.8.8.8"})
+        cache.put(TOKEN, "lite", "1.1.1.1", {"ip": "1.1.1.1"})
 
-        cached, misses = cache.get_many("lite", ["8.8.8.8", "1.1.1.1"])
+        cached, misses = cache.get_many(TOKEN, "lite", ["8.8.8.8", "1.1.1.1"])
         assert cached == {"8.8.8.8": {"ip": "8.8.8.8"}, "1.1.1.1": {"ip": "1.1.1.1"}}
         assert misses == []
 
     def test_all_missing(self) -> None:
         cache = IPCache()
-        cached, misses = cache.get_many("lite", ["8.8.8.8", "1.1.1.1"])
+        cached, misses = cache.get_many(TOKEN, "lite", ["8.8.8.8", "1.1.1.1"])
         assert cached == {}
         assert misses == ["8.8.8.8", "1.1.1.1"]
 
     def test_partial_cache(self) -> None:
         cache = IPCache()
-        cache.put("lite", "8.8.8.8", {"ip": "8.8.8.8"})
+        cache.put(TOKEN, "lite", "8.8.8.8", {"ip": "8.8.8.8"})
 
-        cached, misses = cache.get_many("lite", ["8.8.8.8", "1.1.1.1"])
+        cached, misses = cache.get_many(TOKEN, "lite", ["8.8.8.8", "1.1.1.1"])
         assert cached == {"8.8.8.8": {"ip": "8.8.8.8"}}
         assert misses == ["1.1.1.1"]
 
     def test_empty_list(self) -> None:
         cache = IPCache()
-        cached, misses = cache.get_many("lite", [])
+        cached, misses = cache.get_many(TOKEN, "lite", [])
         assert cached == {}
         assert misses == []
 
     def test_uses_correct_namespace(self) -> None:
         cache = IPCache()
-        cache.put("lite", "8.8.8.8", {"ip": "8.8.8.8"})
+        cache.put(TOKEN, "lite", "8.8.8.8", {"ip": "8.8.8.8"})
 
-        cached, misses = cache.get_many("lookup", ["8.8.8.8"])
+        cached, misses = cache.get_many(TOKEN, "lookup", ["8.8.8.8"])
         assert cached == {}
         assert misses == ["8.8.8.8"]
 
@@ -75,12 +109,18 @@ def test_put_many_stores_all(self) -> None:
             "8.8.8.8": {"ip": "8.8.8.8"},
             "1.1.1.1": {"ip": "1.1.1.1"},
         }
-        cache.put_many("lite", items)
+        cache.put_many(TOKEN, "lite", items)
 
-        assert cache.get("lite", "8.8.8.8") == {"ip": "8.8.8.8"}
-        assert cache.get("lite", "1.1.1.1") == {"ip": "1.1.1.1"}
+        assert cache.get(TOKEN, "lite", "8.8.8.8") == {"ip": "8.8.8.8"}
+        assert cache.get(TOKEN, "lite", "1.1.1.1") == {"ip": "1.1.1.1"}
 
     def test_put_many_empty_dict(self) -> None:
         cache = IPCache()
-        cache.put_many("lite", {})
-        assert cache.get("lite", "8.8.8.8") is None
+        cache.put_many(TOKEN, "lite", {})
+        assert cache.get(TOKEN, "lite", "8.8.8.8") is None
+
+    def test_put_many_is_scoped_to_token(self) -> None:
+        cache = IPCache()
+        cache.put_many(TOKEN, "lite", {"8.8.8.8": {"ip": "8.8.8.8"}})
+
+        assert cache.get(OTHER_TOKEN, "lite", "8.8.8.8") is None

tests/tools/test_asn.py

@@ -106,9 +106,9 @@ class TestAsnCaching:
     async def test_uses_cache(
         self, client: IPinfoClient, cache: IPCache, httpx_mock: HTTPXMock
     ) -> None:
-        cache.put("lite", "8.8.8.8", LITE_8888)
+        cache.put("fake_token", "lite", "8.8.8.8", LITE_8888)
 
-        ctx = make_context(client, cache)
+        ctx = make_context(client, cache, api_token="fake_token")
         result = await ipinfo_asn(
             ips=["8.8.8.8"], detailed=False, page=1, page_size=5, ctx=ctx
         )

tests/tools/test_geolocate.py

@@ -106,9 +106,9 @@ class TestGeolocateCaching:
     async def test_uses_cache(
         self, client: IPinfoClient, cache: IPCache, httpx_mock: HTTPXMock
     ) -> None:
-        cache.put("lite", "8.8.8.8", LITE_8888)
+        cache.put("fake_token", "lite", "8.8.8.8", LITE_8888)
 
-        ctx = make_context(client, cache)
+        ctx = make_context(client, cache, api_token="fake_token")
         result = await ipinfo_geolocate(
             ips=["8.8.8.8"], detailed=False, page=1, page_size=5, ctx=ctx
         )

tests/tools/test_lookup.py

@@ -176,14 +176,14 @@ async def test_partial_cache_hit(
         self, client: IPinfoClient, cache: IPCache, httpx_mock: HTTPXMock
     ) -> None:
         # Pre-populate cache with one IP
-        cache.put("lite", "8.8.8.8", LITE_8888)
+        cache.put("fake_token", "lite", "8.8.8.8", LITE_8888)
 
         httpx_mock.add_response(
             url=f"{BASE_URL}/batch",
             method="POST",
             json={"lite/1.1.1.1": LITE_1111},
         )
-        ctx = make_context(client, cache)
+        ctx = make_context(client, cache, api_token="fake_token")
         result = await ipinfo_lookup(
             ips=["8.8.8.8", "1.1.1.1"], detailed=False, page=1, page_size=5, ctx=ctx
         )

tests/tools/test_privacy.py

@@ -134,9 +134,9 @@ async def test_uses_lookup_cache(
         self, client: IPinfoClient, cache: IPCache, httpx_mock: HTTPXMock
     ) -> None:
         """Privacy should reuse data already cached by a lookup call."""
-        cache.put("lookup", "8.8.8.8", LOOKUP_8888)
+        cache.put("fake_token", "lookup", "8.8.8.8", LOOKUP_8888)
 
-        ctx = make_context(client, cache)
+        ctx = make_context(client, cache, api_token="fake_token")
         result = await ipinfo_check_privacy(
             ips=["8.8.8.8"], page=1, page_size=5, ctx=ctx
         )

tests/tools/test_resproxy.py

@@ -71,9 +71,9 @@ class TestResproxyCaching:
     async def test_uses_resproxy_cache(
         self, client: IPinfoClient, cache: IPCache, httpx_mock: HTTPXMock
     ) -> None:
-        cache.put("resproxy", "1.2.3.4", RESPROXY_HIT)
+        cache.put("fake_token", "resproxy", "1.2.3.4", RESPROXY_HIT)
 
-        ctx = make_context(client, cache)
+        ctx = make_context(client, cache, api_token="fake_token")
         result = await ipinfo_check_residential_proxy(
             ips=["1.2.3.4"], page=1, page_size=5, ctx=ctx
         )