pickles-in-pytorch.yaml

rules:
  - id: pickles-in-pytorch
    message: >-
      Functions reliant on pickle can result in arbitrary code execution. 
      Consider loading from `state_dict`, using fickling, or switching to a safer serialization method like ONNX
    languages: [python]
    severity: ERROR
    metadata:
      category: security
      cwe: "CWE-502: Deserialization of Untrusted Data"
      subcategory: [vuln]
      confidence: MEDIUM
      likelihood: MEDIUM
      impact: HIGH
      technology: [pytorch]
      description: "Potential arbitrary code execution from `PyTorch` functions reliant on pickling"
      references:
        - https://blog.trailofbits.com/2021/03/15/never-a-dill-moment-exploiting-machine-learning-pickle-files/

    patterns:
      - pattern-either:
          - pattern: torch.save(...)
          - pattern: torch.load(...)
      - pattern-not: torch.load("...")
      - pattern-not: torch.save(..., "...")
      - pattern-not: torch.save($M.state_dict(), ...)
      - pattern-not-inside: $M.load_state_dict(...)

      - pattern-not:
          patterns:
            - pattern: torch.save($STATE_DICT, ...)
            - pattern-inside: |
                $STATE_DICT = $M.state_dict()
                ...