Unsafe torch.load deserialization in load_pretrained_params() allows remote code execution #2175

Affected versions: confirmed on commit 11941d2bf1f77a496fda44272ef3a7ce3a8a0f45 (tag v1.1.0); the vulnerable pattern (torch.load without weights_only=True) has been present at this call site across docTR's history and applies to any install satisfying docTR's own torch>=2.0.0,<3.0.0 dependency bound using torch versions older than 2.6 (torch's own default for weights_only only flipped to True in 2.6).

Summary

doctr.models.utils.pytorch.load_pretrained_params() loads a model checkpoint with torch.load(archive_path, map_location="cpu"), without weights_only=True. On any torch release docTR currently supports (torch>=2.0.0,<3.0.0, i.e. anything before 2.6) this uses Python's pickle module with no restriction on what object graph may be constructed, so a crafted checkpoint file executes arbitrary code the moment it is loaded, before any state_dict key is ever checked. The same function accepts a URL, downloads it, and loads it the same way, and the library's own documentation presents loading a checkpoint from an arbitrary URL or local path as a normal usage pattern.

Details

doctr/models/utils/pytorch.py:

def load_pretrained_params(
    model: nn.Module,
    path_or_url: str | None = None,
    hash_prefix: str | None = None,
    ignore_keys: list[str] | None = None,
    **kwargs: Any,
) -> None:
    """Load a set of parameters onto a model

    >>> from doctr.models import load_pretrained_params
    >>> load_pretrained_params(model, "https://yoursource.com/yourcheckpoint-yourhash.pt")
    ...
    """
    if path_or_url is None:
        logger.warning("No model URL or Path provided, using default initialization.")
        return

    archive_path = (
        download_from_url(path_or_url, hash_prefix=hash_prefix, cache_subdir="models", **kwargs)
        if validators.url(path_or_url)
        else path_or_url
    )

    # Read state_dict
    state_dict = torch.load(archive_path, map_location="cpu")

    # Remove weights from the state_dict
    if ignore_keys is not None and len(ignore_keys) > 0:
        ...
    else:
        # Load weights
        model.load_state_dict(state_dict)

path_or_url is either a URL (downloaded via doctr.utils.data.download_from_url, which performs no code-level validation of the fetched bytes) or a local path, taken as-is. The hash_prefix integrity check in download_from_url is optional (None by default) and only verifies that a downloaded file's SHA256 matches a value the caller must already know; it does nothing to restrict what torch.load is allowed to deserialize, and is never consulted at all when path_or_url is a local filesystem path rather than a URL.

torch.load() without weights_only=True unpickles the file with the standard pickle module. Any object whose class defines __reduce__ runs that method's returned callable during unpickling, so a checkpoint can execute an arbitrary function (e.g. os.system) as a side effect of being loaded, independent of whatever model.load_state_dict() does afterward. The docstring on this exact function shows loading a checkpoint from a caller-supplied URL as the intended usage, so any downstream application that exposes "load a custom/community checkpoint" (a common feature request for libraries like this one: letting a user point at their own fine-tuned weights) hands that same code-execution primitive to whoever controls the file at that URL/path.

This is the same code path reached by doctr.contrib.base.py's _BaseContrib (loads an ONNX model via the same download_from_url, not the torch.load sink) and by doctr/datasets/datasets/base.py's dataset downloader, but the PyTorch checkpoint path is the one with a direct pickle-based remote-code-execution sink.

Steps To Reproduce

(available upon request)

Impact

Any application built on docTR that loads a checkpoint from a path or URL not fully controlled by the application operator (a "load your own fine-tuned model" feature, a shared/community checkpoint link, a file a user uploads) gives whoever controls that file arbitrary code execution in the process loading it, with no dependency on what the checkpoint's tensor contents are or whether load_state_dict() ever succeeds. This reaches the full privileges of the process calling load_pretrained_params(), confirmed as root in an isolated test container. The same sink is reachable through the library's documented URL-based usage, so a victim does not need to be tricked into running an unusual code path, only into pointing the library's own advertised feature at an attacker-controlled file.

CVSS 3.1: 8.8 (High) - AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H
CWE-502: Deserialization of Untrusted Data

Suggested fix: pass weights_only=True to the torch.load() call at doctr/models/utils/pytorch.py:

-    state_dict = torch.load(archive_path, map_location="cpu")
+    state_dict = torch.load(archive_path, map_location="cpu", weights_only=True)