[Rule] Unsafe ML model deserialization (torch.load/joblib/dill/numpy allow_pickle) — CWE-502

[DipesThapa]

Rule Description

What does this rule intend to find?

Proposing a new audit rule — ml-unsafe-model-deserialization — that flags loading ML artifacts through pickle-backed deserialisers on non-safe paths. This is the model-file attack surface (CWE-502) behind a large share of recent AI/ML CVEs.

Existing registry rules cover generic pickle.loads / yaml.load. This rule targets the ML-specific loaders those miss, and recognises the safe form so it does not fire on it:

• torch.load(...) without weights_only=True
• joblib.load(...), dill.load / dill.loads(...)
• pandas.read_pickle(...) / pd.read_pickle(...)
• numpy.load(..., allow_pickle=True) / np.load(..., allow_pickle=True)

Remediation: torch.load(..., weights_only=True), safetensors for weights, numpy.load(..., allow_pickle=False), and don't unpickle untrusted joblib/dill/pandas artifacts.

Examples or references

• CWE-502: https://cwe.mitre.org/data/definitions/502.html
• torch.load docs (weights_only): https://pytorch.org/docs/stable/generated/torch.load.html

Proposed rule (passes semgrep --test):

rules:
  - id: ml-unsafe-model-deserialization
    languages: [python]
    severity: WARNING
    message: >
      This loads a machine-learning artifact through a pickle-backed
      deserialiser. If the file/URL is attacker-controlled, loading it executes
      arbitrary code (CWE-502) — the model-file attack surface behind many AI/ML
      CVEs. Prefer safe formats/flags: torch.load(..., weights_only=True),
      safetensors for weights, numpy.load(..., allow_pickle=False), and never
      unpickle untrusted joblib/dill/pandas artifacts.
    metadata:
      category: security
      subcategory:
        - audit
      confidence: MEDIUM
      likelihood: MEDIUM
      impact: HIGH
      cwe:
        - "CWE-502: Deserialization of Untrusted Data"
      owasp:
        - A08:2021 - Software and Data Integrity Failures
      references:
        - https://cwe.mitre.org/data/definitions/502.html
        - https://pytorch.org/docs/stable/generated/torch.load.html
      technology:
        - python
        - pytorch
        - ml
    patterns:
      - pattern-either:
          - patterns:
              - pattern: torch.load(...)
              - pattern-not: torch.load(..., weights_only=True)
          - pattern: joblib.load(...)
          - pattern: pandas.read_pickle(...)
          - pattern: pd.read_pickle(...)
          - pattern: dill.load(...)
          - pattern: dill.loads(...)
          - pattern: numpy.load(..., allow_pickle=True, ...)
          - pattern: np.load(..., allow_pickle=True, ...)

Test file (matches the rule name; # ruleid: / # ok:):

import dill
import joblib
import numpy
import numpy as np
import pandas
import pandas as pd
import torch


def true_positives(path, url, data):
    # ruleid: ml-unsafe-model-deserialization
    torch.load(path)

    # ruleid: ml-unsafe-model-deserialization
    torch.load(path, map_location="cpu")

    # ruleid: ml-unsafe-model-deserialization
    joblib.load(path)

    # ruleid: ml-unsafe-model-deserialization
    numpy.load(path, allow_pickle=True)

    # ruleid: ml-unsafe-model-deserialization
    np.load(path, allow_pickle=True)

    # ruleid: ml-unsafe-model-deserialization
    pandas.read_pickle(url)

    # ruleid: ml-unsafe-model-deserialization
    pd.read_pickle(path)

    # ruleid: ml-unsafe-model-deserialization
    dill.loads(data)


def true_negatives(path, data):
    # ok: ml-unsafe-model-deserialization
    torch.load(path, weights_only=True)

    # ok: ml-unsafe-model-deserialization
    numpy.load(path)

    # ok: ml-unsafe-model-deserialization
    np.load(path, allow_pickle=False)

    # ok: ml-unsafe-model-deserialization
    import json
    json.loads(data)

Verified locally: semgrep --test → 1/1: ✓ All tests passed.

Additional information

Happy to open the PR at python/ml/security/audit/ml-unsafe-model-deserialization.{yaml,py} (or wherever you prefer) and set the id + metadata to match your conventions. Would you accept this rule? If it overlaps something already planned or you'd prefer a narrower/broader scope, glad to adjust.

PR Checklist

• This ticket has links, references, or examples.
• The rule has true positive and true negative test cases in a file that matches the rule name.

[DipesThapa]

Quick update after some due diligence:

• Confirmed via code search that there's currently no coverage for torch.load, joblib.load, pandas.read_pickle, or numpy.load(..., allow_pickle=True).
• dill is already handled by avoid-dill, so I've dropped it from this proposal.
• Refactored into per-library avoid-* rules to match the existing python/lang/security/deserialization/ style, using the same metadata shape (OWASP 2017/2021/2025, cwe2021/2022-top25) and a pattern-not on constant-string args to reduce false positives.
• semgrep --test passes 4/4.

I'll open a PR adding python/lang/security/deserialization/ml-deserialization.{yaml,py} (base: develop) referencing this issue. Happy to adjust naming or scope to your preference.