diff --git a/platformio/registry/mirror.py b/platformio/registry/mirror.py index 240e9a814a..4d48791471 100644 --- a/platformio/registry/mirror.py +++ b/platformio/registry/mirror.py @@ -43,10 +43,17 @@ def __next__(self): if result is not None: try: headers = json.loads(result) - return ( - headers["Location"], - headers["X-PIO-Content-SHA256"], - ) + mirror = headers.get("X-PIO-Mirror") + # Only use the cached entry if it identifies a mirror we + # have not already tried. Without this, a cached redirect + # to a failing mirror would be returned forever because + # _visited_mirrors never advances on the cache-hit path. + if mirror and mirror not in self._visited_mirrors: + self._visited_mirrors.append(mirror) + return ( + headers["Location"], + headers["X-PIO-Content-SHA256"], + ) except (ValueError, KeyError): pass @@ -84,6 +91,7 @@ def __next__(self): "X-PIO-Content-SHA256": response.headers.get( "X-PIO-Content-SHA256" ), + "X-PIO-Mirror": response.headers.get("X-PIO-Mirror"), } ), "1h", diff --git a/tests/registry/__init__.py b/tests/registry/__init__.py new file mode 100644 index 0000000000..b051490361 --- /dev/null +++ b/tests/registry/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2014-present PlatformIO +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/registry/test_mirror.py b/tests/registry/test_mirror.py new file mode 100644 index 0000000000..2720add7d3 --- /dev/null +++ b/tests/registry/test_mirror.py @@ -0,0 +1,180 @@ +# Copyright (c) 2014-present PlatformIO +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# pylint: disable=protected-access,redefined-outer-name + +import json +from typing import Any, Dict, List, Optional, Tuple, Type + +import pytest + +from platformio.cache import ContentCache +from platformio.registry import mirror +from platformio.registry.mirror import RegistryFileMirrorIterator + + +class _FakeContentCache: + """In memory stand in for ContentCache used by the iterator tests.""" + + _store: Dict[str, str] = {} + + def __init__(self, _namespace: Optional[str] = None) -> None: + pass + + def __enter__(self) -> "_FakeContentCache": + return self + + def __exit__(self, *_exc: Any) -> bool: + return False + + def get(self, key: str) -> Optional[str]: + return type(self)._store.get(key) + + def set(self, key: str, data: str, _valid: str) -> None: + type(self)._store[key] = data + + @staticmethod + def key_from_args(*args: Any) -> str: + return ContentCache.key_from_args(*args) + + +class _FakeResponse: + def __init__(self, status_code: int, headers: Dict[str, str]) -> None: + self.status_code = status_code + self.headers = headers + + +class _FakeHTTPClient: + def __init__(self, responses: List[_FakeResponse]) -> None: + self._responses = list(responses) + self.requests: List[Tuple[str, str, Dict[str, Any]]] = [] + + def send_request(self, method: str, path: str, **kwargs: Any) -> _FakeResponse: + self.requests.append((method, path, kwargs)) + return self._responses.pop(0) + + +@pytest.fixture +def fake_cache(monkeypatch: pytest.MonkeyPatch) -> Type[_FakeContentCache]: + _FakeContentCache._store = {} + monkeypatch.setattr(mirror, "ContentCache", _FakeContentCache) + return _FakeContentCache + + +def _install_fake_http( + monkeypatch: pytest.MonkeyPatch, responses: List[_FakeResponse] +) -> _FakeHTTPClient: + http_client = _FakeHTTPClient(responses) + monkeypatch.setattr( + RegistryFileMirrorIterator, "get_http_client", lambda _self: http_client + ) + return http_client + + +def test_cache_hit_advances_visited_mirrors( + fake_cache: Type[_FakeContentCache], monkeypatch: pytest.MonkeyPatch +) -> None: + # A cached mirror must be recorded as visited on the cache hit path, + # otherwise the next iteration would recompute the same cache key and + # replay the same redirect forever when the mirror keeps failing. + url = "https://dl.example.com/packages/tool.tar.gz" + fake_cache._store[ContentCache.key_from_args("head", url, [])] = json.dumps( + { + "Location": "https://mirror-a.example.com/tool.tar.gz", + "X-PIO-Content-SHA256": "aaa", + "X-PIO-Mirror": "mirror-a", + } + ) + http_client = _install_fake_http( + monkeypatch, + [ + _FakeResponse( + 302, + { + "Location": "https://mirror-b.example.com/tool.tar.gz", + "X-PIO-Mirror": "mirror-b", + "X-PIO-Content-SHA256": "bbb", + }, + ) + ], + ) + + iterator = RegistryFileMirrorIterator(url) + first = next(iterator) + second = next(iterator) + + assert first == ("https://mirror-a.example.com/tool.tar.gz", "aaa") + assert second == ("https://mirror-b.example.com/tool.tar.gz", "bbb") + assert http_client.requests, "second iteration should issue a HEAD request" + assert http_client.requests[-1][2]["params"] == {"bypass": "mirror-a"} + + +def test_cache_entry_without_mirror_header_falls_through( + fake_cache: Type[_FakeContentCache], monkeypatch: pytest.MonkeyPatch +) -> None: + # Entries written before X-PIO-Mirror was cached lack the field; the + # iterator must treat them as unusable and request a fresh redirect so + # the fix self heals existing caches. + url = "https://dl.example.com/packages/tool.tar.gz" + fake_cache._store[ContentCache.key_from_args("head", url, [])] = json.dumps( + { + "Location": "https://mirror-a.example.com/tool.tar.gz", + "X-PIO-Content-SHA256": "aaa", + } + ) + _install_fake_http( + monkeypatch, + [ + _FakeResponse( + 302, + { + "Location": "https://mirror-b.example.com/tool.tar.gz", + "X-PIO-Mirror": "mirror-b", + "X-PIO-Content-SHA256": "bbb", + }, + ) + ], + ) + + iterator = RegistryFileMirrorIterator(url) + + assert next(iterator) == ("https://mirror-b.example.com/tool.tar.gz", "bbb") + + +def test_head_response_caches_mirror_header( + fake_cache: Type[_FakeContentCache], monkeypatch: pytest.MonkeyPatch +) -> None: + # The iterator must persist X-PIO-Mirror so later runs can tell whether + # the cached redirect still names a mirror that has not been tried. + url = "https://dl.example.com/packages/tool.tar.gz" + _install_fake_http( + monkeypatch, + [ + _FakeResponse( + 302, + { + "Location": "https://mirror-a.example.com/tool.tar.gz", + "X-PIO-Mirror": "mirror-a", + "X-PIO-Content-SHA256": "aaa", + }, + ) + ], + ) + + iterator = RegistryFileMirrorIterator(url) + next(iterator) + + cached = json.loads(fake_cache._store[ContentCache.key_from_args("head", url, [])]) + assert cached["X-PIO-Mirror"] == "mirror-a" + assert cached["Location"] == "https://mirror-a.example.com/tool.tar.gz"