Validate Welib fallback URLs (#977)

shelfmark/release_sources/direct_download.py

@@ -222,6 +222,35 @@ def _is_configured_zlib_link(url: str) -> bool:
     return False
 
 
+def _get_direct_download_allowed_hosts() -> set[str]:
+    """Return configured hosts allowed for direct-download fallback fetches."""
+    from shelfmark.core import mirrors
+
+    candidate_urls = [
+        network.get_aa_base_url(),
+        *mirrors.get_aa_mirrors(),
+        *mirrors.get_welib_mirrors(),
+    ]
+    hosts: set[str] = set()
+    for candidate_url in candidate_urls:
+        parsed = urlparse(candidate_url)
+        if parsed.scheme not in {"http", "https"}:
+            continue
+        if parsed.hostname:
+            hosts.add(parsed.hostname.lower().rstrip("."))
+    return hosts
+
+
+def _is_allowed_direct_download_url(url: str) -> bool:
+    """Return True when URL scheme and host match configured direct-download mirrors."""
+    parsed = urlparse(url)
+    if parsed.scheme not in {"http", "https"} or not parsed.hostname:
+        return False
+
+    hostname = parsed.hostname.lower().rstrip(".")
+    return hostname in _get_direct_download_allowed_hosts()
+
+
 def _get_md5_url_template(source_id: str) -> str | None:
     """Get URL template for MD5-based sources from centralized config."""
     from shelfmark.core import mirrors
@@ -897,6 +926,10 @@ def _try_download_url(
     try:
         logger.info("Trying download source [%s]: %s", source_id, url)
 
+        if source_id == "welib" and not _is_allowed_direct_download_url(url):
+            logger.warning("Blocked Welib fallback URL outside configured mirrors: %s", url)
+            return None
+
         if status_callback:
             status_callback("resolving", f"Trying {source_context}")
 
@@ -906,6 +939,12 @@ def _try_download_url(
         if not download_url:
             _raise_runtime_error("No download URL resolved")
 
+        if source_id == "welib" and not _is_allowed_direct_download_url(download_url):
+            logger.warning(
+                "Blocked Welib download URL outside configured mirrors: %s", download_url
+            )
+            return None
+
         logger.info("Resolved download URL [%s]: %s", source_id, download_url)
 
         data = downloader.download_url(
@@ -986,9 +1025,12 @@ def _get_download_urls_from_welib(
 
     soup = BeautifulSoup(_html_response_text(html), "html.parser")
     links = [
-        downloader.get_absolute_url(url, href)
+        absolute_url
         for a in soup.find_all("a", href=True)
-        if (href := _get_attr(a, "href")) and "/slow_download/" in href
+        if (href := _get_attr(a, "href"))
+        and "/slow_download/" in href
+        and (absolute_url := downloader.get_absolute_url(url, href))
+        and _is_allowed_direct_download_url(absolute_url)
     ]
     return list(dict.fromkeys(links))  # Dedupe while preserving order
 

tests/direct_download/test_welib_url_allowlist.py

@@ -0,0 +1,86 @@
+from io import BytesIO
+
+from shelfmark.release_sources import BrowseRecord
+
+
+def _book() -> BrowseRecord:
+    return BrowseRecord(
+        id="abc123",
+        title="Test Book",
+        source="direct_download",
+        size="20 KB",
+    )
+
+
+def _enable_welib_only(monkeypatch, dd):
+    monkeypatch.setattr(dd, "_get_source_priority", lambda: [{"id": "welib", "enabled": True}])
+    monkeypatch.setattr(dd, "_is_source_enabled", lambda source_id: source_id == "welib")
+    monkeypatch.setattr(dd.config, "USE_CF_BYPASS", True)
+    monkeypatch.setattr(
+        "shelfmark.core.mirrors.get_welib_url_template",
+        lambda: "https://welib.example/md5/{md5}",
+    )
+    monkeypatch.setattr(
+        "shelfmark.core.mirrors.get_welib_mirrors",
+        lambda: ["https://welib.example"],
+    )
+    monkeypatch.setattr("shelfmark.core.mirrors.get_aa_mirrors", lambda: [])
+    monkeypatch.setattr(dd.network, "get_aa_base_url", lambda: "https://annas.example")
+
+
+def test_welib_rejects_hostile_returned_url_before_fetch(monkeypatch, tmp_path):
+    import shelfmark.release_sources.direct_download as dd
+
+    _enable_welib_only(monkeypatch, dd)
+    fetched_urls: list[str] = []
+
+    def fake_html_get_page(url: str, **_kwargs):
+        fetched_urls.append(url)
+        if url == "https://welib.example/md5/abc123":
+            return '<a href="http://169.254.169.254/slow_download/abc123">Download</a>'
+        raise AssertionError(f"unexpected fetch: {url}")
+
+    def unexpected_download(*_args, **_kwargs):
+        raise AssertionError("hostile URL must not reach file download")
+
+    monkeypatch.setattr(dd.downloader, "html_get_page", fake_html_get_page)
+    monkeypatch.setattr(dd.downloader, "download_url", unexpected_download)
+
+    result = dd._download_book(_book(), tmp_path / "book.epub")
+
+    assert result is None
+    assert fetched_urls == ["https://welib.example/md5/abc123"]
+
+
+def test_welib_allows_configured_host_returned_url(monkeypatch, tmp_path):
+    import shelfmark.release_sources.direct_download as dd
+
+    _enable_welib_only(monkeypatch, dd)
+    fetched_pages: list[str] = []
+    downloaded: list[tuple[str, str | None]] = []
+
+    def fake_html_get_page(url: str, **_kwargs):
+        fetched_pages.append(url)
+        if url == "https://welib.example/md5/abc123":
+            return '<a href="/files/book.epub">Download</a>'
+        raise AssertionError(f"unexpected fetch: {url}")
+
+    def fake_download_url(url: str, *_args, referer: str | None = None, **_kwargs):
+        downloaded.append((url, referer))
+        payload = BytesIO(b"x" * (11 * 1024))
+        payload.seek(0, 2)
+        return payload
+
+    book_path = tmp_path / "book.epub"
+
+    monkeypatch.setattr(dd.downloader, "html_get_page", fake_html_get_page)
+    monkeypatch.setattr(dd.downloader, "download_url", fake_download_url)
+
+    result = dd._download_book(_book(), book_path)
+
+    assert result == "https://welib.example/files/book.epub"
+    assert fetched_pages == ["https://welib.example/md5/abc123"]
+    assert downloaded == [
+        ("https://welib.example/files/book.epub", "https://welib.example/md5/abc123")
+    ]
+    assert book_path.read_bytes() == b"x" * (11 * 1024)