diff --git a/python/ql/lib/semmle/python/Frameworks.qll b/python/ql/lib/semmle/python/Frameworks.qll index a6288dadc11..85f2ef45060 100644 --- a/python/ql/lib/semmle/python/Frameworks.qll +++ b/python/ql/lib/semmle/python/Frameworks.qll @@ -68,6 +68,7 @@ private import semmle.python.frameworks.SqlAlchemy private import semmle.python.frameworks.Starlette private import semmle.python.frameworks.Stdlib private import semmle.python.frameworks.Toml +private import semmle.python.frameworks.Torch private import semmle.python.frameworks.Tornado private import semmle.python.frameworks.Twisted private import semmle.python.frameworks.Ujson diff --git a/python/ql/lib/semmle/python/frameworks/Torch.qll b/python/ql/lib/semmle/python/frameworks/Torch.qll new file mode 100644 index 00000000000..70094b9b3ec --- /dev/null +++ b/python/ql/lib/semmle/python/frameworks/Torch.qll @@ -0,0 +1,76 @@ +/** + * Provides classes modeling security-relevant aspects of the `torch` PyPI package. + * See https://pypi.org/project/torch/. + */ + +private import python +private import semmle.python.Concepts +private import semmle.python.ApiGraphs + +/** + * Provides models for the `torch` PyPI package. + * See https://pypi.org/project/torch/. + */ +private module Torch { + /** + * A call to `torch.load` + * See https://pytorch.org/docs/stable/generated/torch.load.html#torch.load + */ + private class TorchLoadCall extends Decoding::Range, API::CallNode { + TorchLoadCall() { this = API::moduleImport("torch").getMember("load").getACall() } + + override predicate mayExecuteInput() { + not exists(this.getParameter(2, "pickle_module").asSink()) or + exists(this.getParameter(2, "pickle_module").asSink().asExpr().(None)) + } + + override DataFlow::Node getAnInput() { result = this.getParameter(0, "f").asSink() } + + override DataFlow::Node getOutput() { result = this } + + override string getFormat() { result = "pickle" } + } + + API::Node test() { + result = API::moduleImport("torch").getMember("package").getMember("PackageImporter") + } + + /** + * A call to `torch.package.PackageImporter` + * See https://pytorch.org/docs/stable/package.html#torch.package.PackageImporter + */ + private class TorchPackageImporter extends Decoding::Range, API::CallNode { + TorchPackageImporter() { + this = API::moduleImport("torch").getMember("package").getMember("PackageImporter").getACall() and + exists(this.getAMethodCall("load_pickle")) + } + + override predicate mayExecuteInput() { any() } + + override DataFlow::Node getAnInput() { + result = this.getParameter(0, "file_or_buffer").asSink() + } + + override DataFlow::Node getOutput() { result = this.getAMethodCall("load_pickle") } + + override string getFormat() { result = "pickle" } + } + + /** + * A call to `torch.jit.load` + * See https://pytorch.org/docs/stable/generated/torch.jit.load.html#torch.jit.load + */ + private class TorchJitLoad extends Decoding::Range, API::CallNode { + TorchJitLoad() { + this = API::moduleImport("torch").getMember("jit").getMember("load").getACall() + } + + override predicate mayExecuteInput() { any() } + + override DataFlow::Node getAnInput() { result = this.getParameter(0, "f").asSink() } + + override DataFlow::Node getOutput() { result = this } + + override string getFormat() { result = "pickle" } + } +} diff --git a/python/ql/test/library-tests/frameworks/torch/ConceptsTest.expected b/python/ql/test/library-tests/frameworks/torch/ConceptsTest.expected new file mode 100644 index 00000000000..8ec8033d086 --- /dev/null +++ b/python/ql/test/library-tests/frameworks/torch/ConceptsTest.expected @@ -0,0 +1,2 @@ +testFailures +failures diff --git a/python/ql/test/library-tests/frameworks/torch/ConceptsTest.ql b/python/ql/test/library-tests/frameworks/torch/ConceptsTest.ql new file mode 100644 index 00000000000..b557a0bccb6 --- /dev/null +++ b/python/ql/test/library-tests/frameworks/torch/ConceptsTest.ql @@ -0,0 +1,2 @@ +import python +import experimental.meta.ConceptsTest diff --git a/python/ql/test/library-tests/frameworks/torch/Decoding.py b/python/ql/test/library-tests/frameworks/torch/Decoding.py new file mode 100644 index 00000000000..ae3ef5eb766 --- /dev/null +++ b/python/ql/test/library-tests/frameworks/torch/Decoding.py @@ -0,0 +1,25 @@ +from io import BytesIO + +import torch + + +def someSafeMethod(): + pass + + +PicklePayload = BytesIO(b"payload") +torch.load(PicklePayload) # $ decodeInput=PicklePayload decodeOutput=torch.load(..) decodeFormat=pickle decodeMayExecuteInput +torch.load(PicklePayload, pickle_module=None) # $ decodeInput=PicklePayload decodeOutput=torch.load(..) decodeFormat=pickle decodeMayExecuteInput +torch.load(PicklePayload, pickle_module=someSafeMethod()) # $ decodeInput=PicklePayload decodeOutput=torch.load(..) decodeFormat=pickle + +from torch.package import PackageImporter + +importer = PackageImporter(PicklePayload) # $ decodeInput=PicklePayload PackageImporter(..) decodeFormat=pickle decodeMayExecuteInput +my_tensor = importer.load_pickle("my_resources", "tensor.pkl") # $ decodeOutput=importer.load_pickle(..) + +importer = PackageImporter(PicklePayload) + + +from torch import jit + +jit.load(PicklePayload) # $ decodeInput=PicklePayload decodeOutput=jit.load(..) decodeFormat=pickle decodeMayExecuteInput \ No newline at end of file