Harden artifact loading for downloads and cache validation

src/audioclass/utils.py

@@ -7,8 +7,9 @@
 import hashlib
 import os
 import tempfile
+import time
 from pathlib import Path
-from typing import Generator, Optional, Union
+from typing import BinaryIO, Generator, Optional, Union
 from urllib.parse import urlparse
 
 import numpy as np
@@ -20,6 +21,10 @@
     "batched",
 ]
 
+DEFAULT_TIMEOUT: float = 30
+DEFAULT_RETRIES: int = 3
+DEFAULT_BACKOFF_FACTOR: float = 0.25
+
 
 def _hash_url(url: str) -> str:
     """Calculate the SHA256 hash of a URL.
@@ -70,13 +75,42 @@ def _is_url(path: str) -> bool:
     bool
         True if the path is a URL, False otherwise.
     """
-    return urlparse(path).scheme != ""
+    parsed = urlparse(path)
+    return parsed.scheme in {"http", "https"} and parsed.netloc != ""
+
+
+def _sha256_file(path: Path) -> str:
+    """Calculate the SHA256 hash of a file.
+
+    Parameters
+    ----------
+    path
+        The file to hash.
+
+    Returns
+    -------
+    str
+        The SHA256 digest in hexadecimal format.
+    """
+    digest = hashlib.sha256()
+
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(8192), b""):
+            digest.update(chunk)
+
+    return digest.hexdigest()
 
 
 def load_artifact(
     path: Union[Path, str],
     directory: Optional[Path] = None,
     download: bool = True,
+    *,
+    timeout: float = DEFAULT_TIMEOUT,
+    retries: int = DEFAULT_RETRIES,
+    backoff_factor: float = DEFAULT_BACKOFF_FACTOR,
+    expected_sha256: Optional[str] = None,
+    expected_size: Optional[int] = None,
 ) -> Path:
     """Load an artifact from a local path or a URL.
 
@@ -93,6 +127,19 @@ def load_artifact(
     download
         Whether to download the artifact if it is not found in the cache.
         Defaults to True.
+    timeout
+        Timeout in seconds for each HTTP request. Defaults to 30.
+    retries
+        Maximum number of HTTP attempts for transient failures. Defaults to 3.
+    backoff_factor
+        Multiplier used to compute exponential backoff between retries.
+        Defaults to 0.25.
+    expected_sha256
+        Expected SHA256 checksum of the downloaded artifact. If provided,
+        mismatch raises ValueError.
+    expected_size
+        Expected file size in bytes of the downloaded artifact. If provided,
+        mismatch raises ValueError.
 
     Returns
     -------
@@ -118,22 +165,112 @@ def load_artifact(
     new_path = directory / basename
 
     if new_path.exists():
-        return new_path
+        try:
+            _validate_artifact(
+                new_path,
+                expected_sha256=expected_sha256,
+                expected_size=expected_size,
+            )
+            return new_path
+        except ValueError:
+            if not download:
+                raise
+            new_path.unlink()
 
     if not download:
         raise FileNotFoundError(
             f"Could not find artifact at {new_path} corresponding to {path}"
             f" and download is disabled."
         )
 
-    with new_path.open("wb") as f:
-        response = requests.get(path)
-        response.raise_for_status()
-        f.write(response.content)
+    try:
+        with new_path.open("wb") as f:
+            _download_with_retry(
+                path,
+                f,
+                timeout=timeout,
+                retries=retries,
+                backoff_factor=backoff_factor,
+            )
+    except Exception:
+        if new_path.exists():
+            new_path.unlink()
+        raise
+
+    try:
+        _validate_artifact(
+            new_path,
+            expected_sha256=expected_sha256,
+            expected_size=expected_size,
+        )
+    except ValueError:
+        if new_path.exists():
+            new_path.unlink()
+        raise
 
     return new_path
 
 
+def _download_with_retry(
+    url: str,
+    fileobj: BinaryIO,
+    timeout: float,
+    retries: int,
+    backoff_factor: float,
+) -> None:
+    if retries < 1:
+        raise ValueError("retries must be at least one")
+
+    if backoff_factor < 0:
+        raise ValueError("backoff_factor must be non-negative")
+
+    last_error: Optional[Exception] = None
+
+    for attempt in range(1, retries + 1):
+        try:
+            with requests.get(url, timeout=timeout, stream=True) as response:
+                response.raise_for_status()
+
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        fileobj.write(chunk)
+
+            return
+        except requests.RequestException as error:
+            last_error = error
+
+            if attempt == retries:
+                raise
+
+            delay = backoff_factor * (2 ** (attempt - 1))
+            time.sleep(delay)
+
+    if last_error is not None:  # pragma: no cover
+        raise last_error
+
+
+def _validate_artifact(
+    path: Path,
+    expected_sha256: Optional[str],
+    expected_size: Optional[int],
+) -> None:
+    if expected_size is not None:
+        size = path.stat().st_size
+        if size != expected_size:
+            raise ValueError(
+                f"Invalid file size for artifact at {path}"
+                f" ({size} != {expected_size})"
+            )
+
+    if expected_sha256 is not None:
+        digest = _sha256_file(path)
+        if digest.lower() != expected_sha256.lower():
+            raise ValueError(
+                f"Invalid SHA256 checksum for artifact at {path}"
+                f" ({digest} != {expected_sha256})"
+            )
+
+
 def flat_sigmoid(
     x: np.ndarray,
     sensitivity: float = 1,
@@ -142,9 +279,9 @@ def flat_sigmoid(
 ) -> np.ndarray:
     """Apply a flattened sigmoid function to an array.
 
-    This function applies a sigmoid function to each element of the input array,
-    but with a flattened shape to prevent extreme values. The output values are
-    clipped between 0 and 1.
+    This function applies a sigmoid function to each element of the
+    input array, but with a flattened shape to prevent extreme values.
+    The output values are clipped between 0 and 1.
 
     Parameters
     ----------

tests/test_utils.py

@@ -12,6 +12,10 @@ def test_can_detect_birdnet_urls():
     assert _is_url(LABELS_PATH)
 
 
+def test_windows_like_paths_are_not_detected_as_urls():
+    assert not _is_url("C:\\Users\\name\\artifact.tflite")
+
+
 def test_load_artifact(tmp_path: Path):
     directory = tmp_path / "test"
     path = load_artifact(LABELS_PATH, directory=directory)
@@ -46,6 +50,128 @@ def test_load_artifact_does_not_download_if_file_exists(
     request_get.assert_not_called()
 
 
+def test_load_artifact_uses_timeout_and_streaming(mocker, tmp_path: Path):
+    directory = tmp_path / "test"
+
+    response = mocker.MagicMock()
+    response.iter_content.return_value = [b"abc"]
+    response.raise_for_status.return_value = None
+    response.__enter__.return_value = response
+    request_get = mocker.patch("requests.get", return_value=response)
+
+    load_artifact(
+        LABELS_PATH,
+        directory=directory,
+        timeout=12,
+    )
+
+    request_get.assert_called_once_with(
+        LABELS_PATH,
+        timeout=12,
+        stream=True,
+    )
+
+
+def test_load_artifact_retries_on_failure(mocker, tmp_path: Path):
+    directory = tmp_path / "test"
+
+    response = mocker.MagicMock()
+    response.iter_content.return_value = [b"abc"]
+    response.raise_for_status.return_value = None
+    response.__enter__.return_value = response
+
+    from requests import RequestException
+
+    request_get = mocker.patch("requests.get", side_effect=[
+        RequestException("temporary failure"),
+        response,
+    ])
+
+    sleep = mocker.patch("time.sleep")
+    load_artifact(
+        LABELS_PATH,
+        directory=directory,
+        retries=2,
+        backoff_factor=0.5,
+    )
+
+    assert request_get.call_count == 2
+    sleep.assert_called_once_with(0.5)
+
+
+def test_load_artifact_fails_for_invalid_size(mocker, tmp_path: Path):
+    directory = tmp_path / "test"
+
+    response = mocker.MagicMock()
+    response.iter_content.return_value = [b"abc"]
+    response.raise_for_status.return_value = None
+    response.__enter__.return_value = response
+    mocker.patch("requests.get", return_value=response)
+
+    with pytest.raises(ValueError, match="Invalid file size"):
+        load_artifact(
+            LABELS_PATH,
+            directory=directory,
+            expected_size=4,
+        )
+
+
+def test_load_artifact_fails_for_invalid_checksum(mocker, tmp_path: Path):
+    directory = tmp_path / "test"
+
+    response = mocker.MagicMock()
+    response.iter_content.return_value = [b"abc"]
+    response.raise_for_status.return_value = None
+    response.__enter__.return_value = response
+    mocker.patch("requests.get", return_value=response)
+
+    with pytest.raises(ValueError, match="Invalid SHA256 checksum"):
+        load_artifact(
+            LABELS_PATH,
+            directory=directory,
+            expected_sha256="deadbeef",
+        )
+
+
+def test_load_artifact_removes_partial_file_on_download_error(
+    mocker,
+    tmp_path: Path,
+):
+    directory = tmp_path / "test"
+
+    response = mocker.MagicMock()
+    response.raise_for_status.return_value = None
+    response.__enter__.return_value = response
+
+    def broken_chunks(*args, **kwargs):
+        yield b"partial"
+        raise RuntimeError("stream broke")
+
+    response.iter_content.side_effect = broken_chunks
+    mocker.patch("requests.get", return_value=response)
+
+    with pytest.raises(RuntimeError, match="stream broke"):
+        load_artifact(LABELS_PATH, directory=directory, retries=1)
+
+    path = directory / "BirdNET_GLOBAL_6K_V2.4_Labels_en_uk.txt"
+    assert not path.exists()
+
+
+def test_load_artifact_validates_existing_cached_file(tmp_path: Path):
+    directory = tmp_path / "test"
+    directory.mkdir()
+    path = directory / "BirdNET_GLOBAL_6K_V2.4_Labels_en_uk.txt"
+    path.write_text("abc")
+
+    with pytest.raises(ValueError, match="Invalid file size"):
+        load_artifact(
+            LABELS_PATH,
+            directory=directory,
+            expected_size=4,
+            download=False,
+        )
+
+
 def test_loaded_artifact_is_stored_in_tempdir(
     tmp_path: Path,
     monkeypatch,