trailofbits.python.pytorch-classes-load-library.pytorch-classes-load-library

profile photo of trailofbitstrailofbits
Author
unknown
Download Count*

Loading custom operator libraries can result in arbitrary code execution

Run Locally

Run in CI

Defintion

rules:
  - id: pytorch-classes-load-library
    message: Loading custom operator libraries can result in arbitrary code execution
    languages:
      - python
    severity: ERROR
    metadata:
      category: security
      cwe: "CWE-676: Use of Potentially Dangerous Function"
      subcategory:
        - audit
      confidence: MEDIUM
      likelihood: MEDIUM
      impact: HIGH
      technology:
        - pytorch
      description: Potential arbitrary code execution from `PyTorch` library loading
      references:
        - https://pytorch.org/tutorials/advanced/torch_script_custom_classes.html
      license: AGPL-3.0 license
      vulnerability_class:
        - Dangerous Method or Function
    patterns:
      - pattern: torch.classes.load_library(...)
      - pattern-not: torch.classes.load_library("...")

Examples

pytorch-classes-load-library.py

import torch

path = "lib.so"

# ok: pytorch-classes-load-library
torch.classes.load_library(path)

# ruleid: pytorch-classes-load-library
torch.classes.load_library(input())

def test1(p):
    # ruleid: pytorch-classes-load-library
    torch.classes.load_library(input() + p)

def test2(p):
    # ok: pytorch-classes-load-library
    load_library(p)

def test3(p):
    from torch.classes import load_library
    # ruleid: pytorch-classes-load-library
    load_library(p)

def test4(p):
    from torch.classes import load_library
    # ok: pytorch-classes-load-library
    load_library(path)