diff --git a/src/ipinfo_mcp/cache.py b/src/ipinfo_mcp/cache.py index 7663f4c..03d38ed 100644 --- a/src/ipinfo_mcp/cache.py +++ b/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 diff --git a/src/ipinfo_mcp/tools/asn.py b/src/ipinfo_mcp/tools/asn.py index 1d3ae80..224cb06 100644 --- a/src/ipinfo_mcp/tools/asn.py +++ b/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) diff --git a/src/ipinfo_mcp/tools/geolocate.py b/src/ipinfo_mcp/tools/geolocate.py index 582e294..5242eff 100644 --- a/src/ipinfo_mcp/tools/geolocate.py +++ b/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) diff --git a/src/ipinfo_mcp/tools/lookup.py b/src/ipinfo_mcp/tools/lookup.py index 48999be..a8d7025 100644 --- a/src/ipinfo_mcp/tools/lookup.py +++ b/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) diff --git a/src/ipinfo_mcp/tools/privacy.py b/src/ipinfo_mcp/tools/privacy.py index c60b487..9dfed9f 100644 --- a/src/ipinfo_mcp/tools/privacy.py +++ b/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) diff --git a/src/ipinfo_mcp/tools/resproxy.py b/src/ipinfo_mcp/tools/resproxy.py index e3b1af9..5bf4327 100644 --- a/src/ipinfo_mcp/tools/resproxy.py +++ b/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) diff --git a/tests/test_cache.py b/tests/test_cache.py index 5ba7106..b5501c0 100644 --- a/tests/test_cache.py +++ b/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 diff --git a/tests/tools/test_asn.py b/tests/tools/test_asn.py index 4f18e85..fe5c530 100644 --- a/tests/tools/test_asn.py +++ b/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 ) diff --git a/tests/tools/test_geolocate.py b/tests/tools/test_geolocate.py index 2d5ed04..bfbb4b9 100644 --- a/tests/tools/test_geolocate.py +++ b/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 ) diff --git a/tests/tools/test_lookup.py b/tests/tools/test_lookup.py index 6e6123f..00217e3 100644 --- a/tests/tools/test_lookup.py +++ b/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 ) diff --git a/tests/tools/test_privacy.py b/tests/tools/test_privacy.py index 26bf365..2a6b909 100644 --- a/tests/tools/test_privacy.py +++ b/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 ) diff --git a/tests/tools/test_resproxy.py b/tests/tools/test_resproxy.py index 8e70efb..7d579a2 100644 --- a/tests/tools/test_resproxy.py +++ b/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 )