diff --git a/proton/vpn/core/cache_handler.py b/proton/vpn/core/cache_handler.py index 0b16997..0b2b8f5 100644 --- a/proton/vpn/core/cache_handler.py +++ b/proton/vpn/core/cache_handler.py @@ -23,6 +23,7 @@ import json import os from pathlib import Path +from typing import Callable, Optional from proton.vpn import logging @@ -31,8 +32,15 @@ class CacheHandler: """Used to save, load, and remove cache files.""" - def __init__(self, filepath: str): + def __init__( + self, + filepath: str, + object_hook_factory: Optional[ + Callable[[], Callable[[dict], dict]] + ] = None + ): self._fp = Path(filepath) + self._object_hook_factory = object_hook_factory @property def exists(self): @@ -55,7 +63,11 @@ def load(self): try: with open(self._fp, "r", encoding="utf-8") as f: # pylint: disable=C0103 - return json.load(f) # pylint: disable=C0103 + object_hook = self._object_hook_factory() \ + if self._object_hook_factory else None + return json.load( # pylint: disable=C0103 + f, object_hook=object_hook + ) except (json.decoder.JSONDecodeError, UnicodeDecodeError): filename = os.path.basename(self._fp) logger.warning( diff --git a/proton/vpn/session/servers/logicals.py b/proton/vpn/session/servers/logicals.py index 98fd884..a69d617 100644 --- a/proton/vpn/session/servers/logicals.py +++ b/proton/vpn/session/servers/logicals.py @@ -35,6 +35,43 @@ UNIX_EPOCH = "Thu, 01 Jan 1970 00:00:00 GMT" +_SHARED_SERVER_STRING_FIELDS = frozenset({ + "EntryCountry", "ExitCountry", "HostCountry", "City", "State", "Domain", + "EntryIP", "ExitIP", "Generation", "Label", "ServicesDownReason", + "X25519PublicKey", +}) + + +def _server_string_object_hook() -> Callable[[dict], dict]: + """Creates a JSON object hook that shares immutable server strings.""" + shared_strings = {} + + def share_strings(item: dict) -> dict: + for field, value in item.items(): + if ( + field in _SHARED_SERVER_STRING_FIELDS + and isinstance(value, str) + ): + item[field] = shared_strings.setdefault(value, value) + return item + + return share_strings + + +def _deduplicate_server_strings(logicals: List[dict]) -> None: + """Shares common immutable strings within a decoded server list. + + This remains necessary for freshly downloaded responses, which have + already been decoded by the HTTP client. Cached lists use the same helper + as a JSON object hook so duplicate strings need not coexist in memory. + """ + share_strings = _server_string_object_hook() + + for logical in logicals: + share_strings(logical) + for physical in logical.get("Servers", ()): + share_strings(physical) + class PersistenceKeys(Enum): """JSON Keys used to persist the ServerList to disk.""" @@ -426,7 +463,9 @@ def from_dict( """ try: user_tier = data[PersistenceKeys.USER_TIER.value] - logicals = [LogicalServer(logical_dict) for logical_dict in data["LogicalServers"]] + logical_dicts = data[PersistenceKeys.LOGICALS.value] + _deduplicate_server_strings(logical_dicts) + logicals = [LogicalServer(logical_dict) for logical_dict in logical_dicts] except KeyError as error: raise ServerListDecodeError("Error building server list from dict") from error diff --git a/proton/vpn/session/servers/server_list_fetcher.py b/proton/vpn/session/servers/server_list_fetcher.py index f79490f..1337e7c 100644 --- a/proton/vpn/session/servers/server_list_fetcher.py +++ b/proton/vpn/session/servers/server_list_fetcher.py @@ -26,7 +26,9 @@ from proton.vpn.core.cache_handler import CacheHandler from proton.vpn.session.exceptions import ServerListDecodeError from proton.vpn.session.servers.types import ServerLoad -from proton.vpn.session.servers.logicals import ServerList, PersistenceKeys +from proton.vpn.session.servers.logicals import ( + ServerList, PersistenceKeys, _server_string_object_hook +) from proton.vpn.session.dataclasses import VPNLocation from proton.vpn.session.utils import rest_api_request from proton.vpn.platform.core import ServerStatus # pylint: disable=E0401, E0611 @@ -173,7 +175,10 @@ def __init__( ): self._session = session self._server_list = server_list - self._cache_file = cache_file or CacheHandler(self.CACHE_PATH) + self._cache_file = cache_file or CacheHandler( + self.CACHE_PATH, + object_hook_factory=_server_string_object_hook, + ) def clear_cache(self): """Discards the cache, if existing.""" diff --git a/tests/python/core/test_cachehandler.py b/tests/python/core/test_cachehandler.py index 7f97044..fb95d35 100644 --- a/tests/python/core/test_cachehandler.py +++ b/tests/python/core/test_cachehandler.py @@ -52,6 +52,49 @@ def test_load_stored_cache(self, cache_filepath): assert "load_cache" in data assert "dummy-data" == data["load_cache"] + def test_load_stored_cache_with_object_hook(self, cache_filepath): + def object_hook_factory(): + def object_hook(item): + item["decoded"] = True + return item + + return object_hook + + cache_handler = CacheHandler( + cache_filepath, + object_hook_factory=object_hook_factory, + ) + with open(cache_filepath, "w") as f: + json.dump({"load_cache": "dummy-data"}, f) + + assert cache_handler.load() == { + "load_cache": "dummy-data", + "decoded": True, + } + + def test_object_hook_factory_is_scoped_to_each_load(self, cache_filepath): + load_count = 0 + + def object_hook_factory(): + nonlocal load_count + load_count += 1 + + def object_hook(item): + item["load"] = load_count + return item + + return object_hook + + cache_handler = CacheHandler( + cache_filepath, + object_hook_factory=object_hook_factory, + ) + with open(cache_filepath, "w") as f: + json.dump({"load_cache": "dummy-data"}, f) + + assert cache_handler.load()["load"] == 1 + assert cache_handler.load()["load"] == 2 + def test_load_cache_with_missing_file(self, cache_filepath): cache_handler = CacheHandler(cache_filepath) assert not cache_handler.load() diff --git a/tests/python/session/servers/test_logicals.py b/tests/python/session/servers/test_logicals.py index 4a830f3..e05196a 100644 --- a/tests/python/session/servers/test_logicals.py +++ b/tests/python/session/servers/test_logicals.py @@ -17,18 +17,21 @@ along with ProtonVPN. If not, see . """ import functools +import json +from copy import deepcopy from typing import List import pytest -from unittest.mock import Mock from proton.vpn.session.servers import LogicalServer, ServerFeatureEnum from proton.vpn.session.servers.logicals import ( + _server_string_object_hook, sort_servers_alphabetically_by_country_and_server_name, sort_servers_by_country_and_location_and_enabled_and_load, ServerList ) from proton.vpn.session.location_names_fetcher import LocationTranslations +from proton.vpn.core.cache_handler import CacheHandler def _compact_features(features: List[ServerFeatureEnum]) -> ServerFeatureEnum: @@ -130,6 +133,125 @@ def fixture_api_response() -> str: } +def test_from_dict_deduplicates_common_server_strings(): + country_1 = bytes((67, 72)).decode() + country_2 = bytes((67, 72)).decode() + city_1 = bytes((90, 117, 114, 105, 99, 104)).decode() + city_2 = bytes((90, 117, 114, 105, 99, 104)).decode() + label_1 = bytes((110, 111, 110, 101)).decode() + label_2 = bytes((110, 111, 110, 101)).decode() + domain_bytes = (110, 111, 100, 101, 46, 101, 120, 97, 109, 112, 108, 101) + domain_1 = bytes(domain_bytes).decode() + domain_2 = bytes(domain_bytes).decode() + entry_ip_1 = bytes((49, 57, 50, 46, 48, 46, 50, 46, 49)).decode() + entry_ip_2 = bytes((49, 57, 50, 46, 48, 46, 50, 46, 49)).decode() + assert country_1 is not country_2 + assert city_1 is not city_2 + assert label_1 is not label_2 + assert domain_1 is not domain_2 + assert entry_ip_1 is not entry_ip_2 + + payload = { + "MaxTier": 2, + "LogicalServers": [ + { + "ID": "1", "Name": "CH#1", "ExitCountry": country_1, + "City": city_1, "Domain": domain_1, + "Servers": [{ + "Label": label_1, "Domain": domain_1, + "EntryIP": entry_ip_1, + }], + }, + { + "ID": "2", "Name": "CH#2", "ExitCountry": country_2, + "City": city_2, "Domain": domain_2, + "Servers": [{ + "Label": label_2, "Domain": domain_2, + "EntryIP": entry_ip_2, + }], + }, + ], + } + expected_values = deepcopy(payload) + + server_list = ServerList.from_dict(payload) + + assert server_list[0].exit_country is server_list[1].exit_country + assert server_list[0].city is server_list[1].city + assert ( + server_list[0].physical_servers[0].label + is server_list[1].physical_servers[0].label + ) + assert ( + server_list[0].physical_servers[0].domain + is server_list[1].physical_servers[0].domain + ) + assert ( + server_list[0].physical_servers[0].entry_ip + is server_list[1].physical_servers[0].entry_ip + ) + assert server_list[0].id is not server_list[1].id + assert payload == expected_values + + +def test_json_object_hook_deduplicates_strings_during_decode(): + data = json.loads( + '{"LogicalServers": [' + '{"ExitCountry": "CH", "Servers": [{"Domain": "node.example"}]},' + '{"ExitCountry": "CH", "Servers": [{"Domain": "node.example"}]}' + ']}', + object_hook=_server_string_object_hook(), + ) + + first, second = data["LogicalServers"] + assert first["ExitCountry"] is second["ExitCountry"] + assert first["Servers"][0]["Domain"] is second["Servers"][0]["Domain"] + + +def test_string_sharing_preserves_missing_and_non_string_fields(): + payload = { + "MaxTier": 2, + "LogicalServers": [{ + "ID": "1", + "Name": "CH#1", + "ExitCountry": "CH", + "City": None, + "HostCountry": 42, + "Servers": [{"Label": False}], + }], + } + + server_list = ServerList.from_dict(payload) + + assert server_list[0].city is None + assert server_list[0].host_country == 42 + assert server_list[0].physical_servers[0].label is False + assert "State" not in payload["LogicalServers"][0] + + +def test_cached_string_pool_is_not_retained_between_loads(tmp_path): + shared_value = "shared-value-too-long-for-python-interning" + payload = { + "MaxTier": 2, + "LogicalServers": [ + {"ID": "1", "Name": "CH#1", "ExitCountry": shared_value}, + {"ID": "2", "Name": "CH#2", "ExitCountry": shared_value}, + ], + } + cache = CacheHandler( + tmp_path / "server-list.json", + object_hook_factory=_server_string_object_hook, + ) + cache.save(payload) + + first = cache.load()["LogicalServers"] + second = cache.load()["LogicalServers"] + + assert first[0]["ExitCountry"] is first[1]["ExitCountry"] + assert second[0]["ExitCountry"] is second[1]["ExitCountry"] + assert first[0]["ExitCountry"] is not second[0]["ExitCountry"] + + def test_set_location_translations_applies_to_every_logical(api_response: str): server_list = ServerList( user_tier=2,