Harden model weight acquisition and deserialization

src/deeprhythm/batch_infer.py

@@ -14,8 +14,8 @@
     AudioTooShortError,
     class_to_bpm,
     get_device,
-    get_weights,
     load_and_split_audio,
+    load_weights,
 )
 
 NUM_WORKERS = 8
@@ -120,7 +120,7 @@ def consume_and_process(
     if not quiet:
         print('made kernels')
     model = DeepRhythmModel()
-    model.load_state_dict(torch.load(get_weights(quiet=quiet), map_location=torch.device(device), weights_only=False))
+    model.load_state_dict(load_weights(torch.device(device), quiet=quiet))
     model = model.to(device=device)
     model.eval()
     if not quiet:

src/deeprhythm/model/predictor.py

@@ -10,7 +10,7 @@
 from deeprhythm.batch_infer import get_audio_files
 from deeprhythm.batch_infer import main as batch_infer_main
 from deeprhythm.model.frame_cnn import DeepRhythmModel
-from deeprhythm.utils import class_to_bpm, get_device, get_weights, load_and_split_audio, split_audio
+from deeprhythm.utils import class_to_bpm, get_device, load_and_split_audio, load_weights, split_audio
 
 
 class DeepRhythmPredictor:
@@ -22,15 +22,15 @@ def __init__(self, device: Optional[str] = None, quiet: bool = False):
             device: Device to run inference on ('cuda' or 'cpu')
             quiet: Whether to suppress progress messages
         """
-        self.model_path = get_weights(quiet=quiet)
         self.device = torch.device(device if device else get_device())
+        self.quiet = quiet
         self.model = self._load_model()
         self.specs = self._make_kernels()
 
     def _load_model(self) -> DeepRhythmModel:
         """Load and initialize the model."""
         model = DeepRhythmModel()
-        model.load_state_dict(torch.load(self.model_path, map_location=self.device, weights_only=False))
+        model.load_state_dict(load_weights(self.device, quiet=self.quiet))
         model = model.to(device=self.device)
         model.eval()
         return model

src/deeprhythm/utils.py

@@ -1,17 +1,46 @@
+import hashlib
 import os
+import tempfile
 
 import librosa
 import requests
 import torch
 
-model_url = 'https://github.com/bleugreen/deeprhythm/raw/main/weights/'
+MODEL_WEIGHTS_FILENAME = "deeprhythm-0.7.pth"
+MODEL_WEIGHTS_URL = (
+    "https://raw.githubusercontent.com/bleugreen/deeprhythm/"
+    "c3548590d12f5d97ccd39dee27e9591e522c05e2/weights/deeprhythm-0.7.pth"
+)
+MODEL_WEIGHTS_SHA256 = "c7cc8cc0425929cd2bf695474d7ec1fd63ed0d0a4a68f361d4e4b57bd9b3d9c4"
+MODEL_WEIGHTS_SIZE = 5_443_228
+MODEL_DOWNLOAD_CONNECT_TIMEOUT = 10
+MODEL_DOWNLOAD_READ_TIMEOUT = 60
+MODEL_DOWNLOAD_TIMEOUT = (MODEL_DOWNLOAD_CONNECT_TIMEOUT, MODEL_DOWNLOAD_READ_TIMEOUT)
+MODEL_DOWNLOAD_CHUNK_SIZE = 64 * 1024
 
 
 class AudioTooShortError(ValueError):
     """Raised when audio file is shorter than minimum required length."""
     pass
 
 
+class ModelWeightsError(RuntimeError):
+    """Raised when verified model weights cannot be acquired."""
+
+
+def _weights_are_valid(path):
+    try:
+        if os.path.getsize(path) != MODEL_WEIGHTS_SIZE:
+            return False
+        digest = hashlib.sha256()
+        with open(path, "rb") as weights_file:
+            for chunk in iter(lambda: weights_file.read(MODEL_DOWNLOAD_CHUNK_SIZE), b""):
+                digest.update(chunk)
+        return digest.hexdigest() == MODEL_WEIGHTS_SHA256
+    except OSError:
+        return False
+
+
 class AudioLoadError(IOError):
     """Raised when audio file cannot be loaded."""
     pass
@@ -26,30 +55,76 @@ def get_device():
         return 'cpu'
 
 
-def get_weights(filename="deeprhythm-0.7.pth", quiet=False):
+def get_weights(quiet=False):
     home_dir = os.path.expanduser("~")
     model_dir = os.path.join(home_dir, ".local", "share", "deeprhythm")
-    if not os.path.exists(model_dir):
-        os.makedirs(model_dir, exist_ok=True)
-    model_path = os.path.join(model_dir, filename)
-
-    if not os.path.isfile(model_path):
-        print("Downloading model weights...")
-        try:
-            r = requests.get(model_url + filename, allow_redirects=True)
-            if r.status_code == 200:
-                with open(model_path, 'wb') as f:
-                    f.write(r.content)
-                print("Model weights downloaded successfully.")
-            else:
-                print(f"Failed to download model weights. HTTP Error: {r.status_code}")
-        except Exception as e:
-            print(f"An error occurred during the download: {e}")
-    else:
+    os.makedirs(model_dir, exist_ok=True)
+    model_path = os.path.join(model_dir, MODEL_WEIGHTS_FILENAME)
+
+    if _weights_are_valid(model_path):
         if not quiet:
-            print("Model weights already exist.")
+            print("Model weights already exist and are verified.")
+        return model_path
 
-    return model_path
+    temporary_path = None
+    try:
+        if not quiet:
+            print("Downloading model weights...")
+        with requests.get(
+            MODEL_WEIGHTS_URL,
+            stream=True,
+            timeout=MODEL_DOWNLOAD_TIMEOUT,
+        ) as response:
+            response.raise_for_status()
+            digest = hashlib.sha256()
+            downloaded_size = 0
+            with tempfile.NamedTemporaryFile(
+                mode="wb",
+                dir=model_dir,
+                prefix=f".{MODEL_WEIGHTS_FILENAME}.",
+                suffix=".tmp",
+                delete=False,
+            ) as temporary_file:
+                temporary_path = temporary_file.name
+                for chunk in response.iter_content(chunk_size=MODEL_DOWNLOAD_CHUNK_SIZE):
+                    if not chunk:
+                        continue
+                    downloaded_size += len(chunk)
+                    if downloaded_size > MODEL_WEIGHTS_SIZE:
+                        raise ModelWeightsError(
+                            f"Model weights exceed the expected size of {MODEL_WEIGHTS_SIZE} bytes"
+                        )
+                    digest.update(chunk)
+                    temporary_file.write(chunk)
+
+        if downloaded_size != MODEL_WEIGHTS_SIZE:
+            raise ModelWeightsError(
+                f"Model weights have size {downloaded_size} bytes; expected {MODEL_WEIGHTS_SIZE}"
+            )
+        if digest.hexdigest() != MODEL_WEIGHTS_SHA256:
+            raise ModelWeightsError("Model weights failed SHA-256 verification")
+
+        os.replace(temporary_path, model_path)
+        temporary_path = None
+        if not quiet:
+            print("Model weights downloaded and verified successfully.")
+        return model_path
+    except requests.RequestException as exc:
+        raise ModelWeightsError(f"Failed to download model weights: {exc}") from exc
+    except OSError as exc:
+        raise ModelWeightsError(f"Failed to install model weights: {exc}") from exc
+    finally:
+        if temporary_path is not None:
+            try:
+                os.unlink(temporary_path)
+            except FileNotFoundError:
+                pass
+
+
+def load_weights(device, quiet=False):
+    """Acquire verified weights and safely deserialize their state dictionary."""
+    path = get_weights(quiet=quiet)
+    return torch.load(path, map_location=device, weights_only=True)
 
 
 def split_audio(audio, sr, clip_length=8, share_mem=False):

tests/test_utils.py

@@ -1,14 +1,160 @@
+import hashlib
+from pathlib import Path
+from unittest.mock import Mock
+
 import numpy as np
+import pytest
+import requests
 
+import deeprhythm.utils as utils
 from deeprhythm.utils import (
     AudioLoadError,
     AudioTooShortError,
     bpm_to_class,
     class_to_bpm,
     get_device,
+    get_weights,
+    load_weights,
     split_audio,
 )
 
+
+class StreamedResponse:
+    def __init__(self, chunks=(), error=None):
+        self.chunks = chunks
+        self.error = error
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        return False
+
+    def raise_for_status(self):
+        if self.error:
+            raise self.error
+
+    def iter_content(self, chunk_size):
+        assert chunk_size == utils.MODEL_DOWNLOAD_CHUNK_SIZE
+        yield from self.chunks
+
+
+@pytest.fixture
+def model_artifact(monkeypatch, tmp_path):
+    payload = b"verified model bytes"
+    monkeypatch.setenv("HOME", str(tmp_path))
+    monkeypatch.setattr(utils, "MODEL_WEIGHTS_SIZE", len(payload))
+    monkeypatch.setattr(utils, "MODEL_WEIGHTS_SHA256", hashlib.sha256(payload).hexdigest())
+    cache_dir = tmp_path / ".local" / "share" / "deeprhythm"
+    destination = cache_dir / utils.MODEL_WEIGHTS_FILENAME
+    return payload, cache_dir, destination
+
+
+def test_get_weights_accepts_valid_cache_without_network(model_artifact, monkeypatch):
+    payload, cache_dir, destination = model_artifact
+    cache_dir.mkdir(parents=True)
+    destination.write_bytes(payload)
+    request = Mock(side_effect=AssertionError("network must not be used"))
+    monkeypatch.setattr(utils.requests, "get", request)
+
+    assert get_weights(quiet=True) == str(destination)
+    request.assert_not_called()
+
+
+def test_get_weights_downloads_from_pinned_url_and_atomically_replaces_corrupt_cache(
+    model_artifact, monkeypatch
+):
+    payload, cache_dir, destination = model_artifact
+    cache_dir.mkdir(parents=True)
+    destination.write_bytes(b"corrupt")
+    request = Mock(return_value=StreamedResponse([payload[:5], payload[5:]]))
+    monkeypatch.setattr(utils.requests, "get", request)
+    real_replace = utils.os.replace
+    replace = Mock(side_effect=real_replace)
+    monkeypatch.setattr(utils.os, "replace", replace)
+
+    assert get_weights(quiet=True) == str(destination)
+    request.assert_called_once_with(
+        utils.MODEL_WEIGHTS_URL,
+        stream=True,
+        timeout=(utils.MODEL_DOWNLOAD_CONNECT_TIMEOUT, utils.MODEL_DOWNLOAD_READ_TIMEOUT),
+    )
+    source, target = replace.call_args.args
+    assert target == str(destination)
+    assert destination.read_bytes() == payload
+    assert not Path(source).exists()
+    assert list(cache_dir.glob("*.tmp")) == []
+
+
+@pytest.mark.parametrize(
+    ("response", "expected_message"),
+    [
+        (StreamedResponse(error=requests.HTTPError("500 Server Error")), "Failed to download"),
+        (StreamedResponse([b"partial"], requests.Timeout("timed out")), "Failed to download"),
+    ],
+)
+def test_get_weights_network_failures_do_not_publish(
+    response, expected_message, model_artifact, monkeypatch
+):
+    _, cache_dir, destination = model_artifact
+    monkeypatch.setattr(utils.requests, "get", Mock(return_value=response))
+
+    with pytest.raises(utils.ModelWeightsError, match=expected_message):
+        get_weights(quiet=True)
+
+    assert not destination.exists()
+    assert list(cache_dir.glob("*.tmp")) == []
+
+
+def test_get_weights_interrupted_stream_preserves_existing_cache(model_artifact, monkeypatch):
+    payload, cache_dir, destination = model_artifact
+    cache_dir.mkdir(parents=True)
+    destination.write_bytes(b"existing corrupt bytes")
+
+    class InterruptedResponse(StreamedResponse):
+        def iter_content(self, chunk_size):
+            yield payload[:5]
+            raise requests.ConnectionError("stream interrupted")
+
+    monkeypatch.setattr(utils.requests, "get", Mock(return_value=InterruptedResponse()))
+
+    with pytest.raises(utils.ModelWeightsError, match="stream interrupted"):
+        get_weights(quiet=True)
+
+    assert destination.read_bytes() == b"existing corrupt bytes"
+    assert list(cache_dir.glob("*.tmp")) == []
+
+
+@pytest.mark.parametrize(
+    ("download", "expected_message"),
+    [
+        (b"verified model bytes!", "exceed the expected size"),
+        (b"short", "have size"),
+        (b"tampered model bytes", "SHA-256"),
+    ],
+)
+def test_get_weights_rejects_invalid_downloads(download, expected_message, model_artifact, monkeypatch):
+    _, cache_dir, destination = model_artifact
+    monkeypatch.setattr(utils.requests, "get", Mock(return_value=StreamedResponse([download])))
+
+    with pytest.raises(utils.ModelWeightsError, match=expected_message):
+        get_weights(quiet=True)
+
+    assert not destination.exists()
+    assert list(cache_dir.glob("*.tmp")) == []
+
+
+def test_load_weights_uses_safe_torch_deserialization(monkeypatch):
+    state_dict = {"layer": "weights"}
+    monkeypatch.setattr(utils, "get_weights", Mock(return_value="/verified/model.pth"))
+    torch_load = Mock(return_value=state_dict)
+    monkeypatch.setattr(utils.torch, "load", torch_load)
+
+    assert load_weights("cpu", quiet=True) == state_dict
+    torch_load.assert_called_once_with(
+        "/verified/model.pth", map_location="cpu", weights_only=True
+    )
+
 # ---------------------------------------------------------------------------
 # BPM <-> class conversion
 # ---------------------------------------------------------------------------