Apply suggestions from code review

yoff · RasmusWL · web-flow · commit 3b9ea3a95854 · 2020-10-14T08:24:26.000+02:00
Co-authored-by: Rasmus Wriedt Larsen &lt;rasmuswriedtlarsen@gmail.com&gt;
diff --git a/python/ql/src/experimental/semmle/python/frameworks/Dill.qll b/python/ql/src/experimental/semmle/python/frameworks/Dill.qll
@@ -43,8 +43,8 @@ private module Dill {
  * See https://pypi.org/project/dill/ (which currently refers you
  * to https://docs.python.org/3/library/pickle.html#pickle.loads)
  */
-private class DillDeserialization extends UnmarshalingFunction::Range {
-  DillDeserialization() {
+private class DillLoadsCall extends UnmarshalingFunction::Range {
+  DillLoadsCall() {
     this.asCfgNode().(CallNode).getFunction() = Dill::dill::loads().asCfgNode()
   }
 
@@ -56,10 +56,5 @@ private class DillDeserialization extends UnmarshalingFunction::Range {
 
   override DataFlow::Node getOutput() { result = this }
 
-  override string getFormat() {
-    result = this.asCfgNode().(CallNode).getArgByName("encoding").(NameNode).getId()
-    or
-    not exists(this.asCfgNode().(CallNode).getArgByName("encoding")) and
-    result = "ASCII"
-  }
+  override string getFormat() { result = "dill" }
 }
diff --git a/python/ql/src/experimental/semmle/python/frameworks/Stdlib.qll b/python/ql/src/experimental/semmle/python/frameworks/Stdlib.qll
@@ -363,8 +363,8 @@ private module Stdlib {
    * A call to `marshal.loads`
    * See https://docs.python.org/3/library/marshal.html#marshal.loads
    */
-  private class MarshalDeserialization extends UnmarshalingFunction::Range {
-    MarshalDeserialization() {
+  private class MarshalLoadsCall extends UnmarshalingFunction::Range {
+    MarshalLoadsCall() {
       this.asCfgNode().(CallNode).getFunction() = marshal::loads().asCfgNode()
     }
 
@@ -376,13 +376,13 @@ private module Stdlib {
 
     override DataFlow::Node getOutput() { result = this }
 
-    override string getFormat() { none() }
+    override string getFormat() { result = "marshal" }
   }
 
   // ---------------------------------------------------------------------------
   // pickle
   // ---------------------------------------------------------------------------
-  private string pickleModuleName() { result in ["pickle", "cPickle"] }
+  private string pickleModuleName() { result in ["pickle", "cPickle", "_pickle"] }
 
   /** Gets a reference to the `pickle` module. */
   private DataFlow::Node pickle(DataFlow::TypeTracker t) {
@@ -416,8 +416,8 @@ private module Stdlib {
    * A call to `pickle.loads`
    * See https://docs.python.org/3/library/pickle.html#pickle.loads
    */
-  private class PickleDeserialization extends UnmarshalingFunction::Range {
-    PickleDeserialization() {
+  private class PickleLoadsCall extends UnmarshalingFunction::Range {
+    PickleLoadsCall() {
       this.asCfgNode().(CallNode).getFunction() = pickle::loads().asCfgNode()
     }
 
@@ -429,11 +429,6 @@ private module Stdlib {
 
     override DataFlow::Node getOutput() { result = this }
 
-    override string getFormat() {
-      result = this.asCfgNode().(CallNode).getArgByName("encoding").(NameNode).getId()
-      or
-      not exists(this.asCfgNode().(CallNode).getArgByName("encoding")) and
-      result = "ASCII"
-    }
+    override string getFormat() { result = "pickle" }
   }
 }
diff --git a/python/ql/src/experimental/semmle/python/frameworks/Yaml.qll b/python/ql/src/experimental/semmle/python/frameworks/Yaml.qll
@@ -60,5 +60,5 @@ private class YamlDeserialization extends UnmarshalingFunction::Range {
 
   override DataFlow::Node getOutput() { result = this }
 
-  override string getFormat() { none() }
+  override string getFormat() { result = "YAML" }
 }
diff --git a/python/ql/test/experimental/library-tests/frameworks/dill/UnmarshalFunction.py b/python/ql/test/experimental/library-tests/frameworks/dill/UnmarshalFunction.py
@@ -1,13 +1,4 @@
-import flask
 import dill
 
-from flask import Flask, request
-
-app = Flask(__name__)
-
-
-@app.route("/")
-def hello():
-    payload = request.args.get("payload")
-    dill.loads(payload)  # $UNSAFE_getAnInput=payload $UNSAFE_getOutput=Attribute() $UNSAFE_getFormat=ASCII
-    dill.loads(payload, encoding='latin1')  # $UNSAFE_getAnInput=payload $UNSAFE_getOutput=Attribute() $UNSAFE_getFormat=latin1
+dill.loads(payload)  # $UNSAFE_getAnInput=payload $UNSAFE_getOutput=Attribute() $UNSAFE_getFormat=ASCII
+dill.loads(payload, encoding='latin1')  # $UNSAFE_getAnInput=payload $UNSAFE_getOutput=Attribute() $UNSAFE_getFormat=latin1
diff --git a/python/ql/test/experimental/query-tests/Security-new-dataflow/CWE-502/unsafe_deserialization.py b/python/ql/test/experimental/query-tests/Security-new-dataflow/CWE-502/unsafe_deserialization.py
@@ -12,10 +12,10 @@
 @app.route("/")
 def hello():
     payload = request.args.get("payload")
-    pickle.loads(payload)
-    yaml.load(payload)
-    yaml.load(payload, Loader=SafeLoader)
-    marshal.loads(payload)
+    pickle.loads(payload) # NOT OK
+    yaml.load(payload) # NOT OK
+    yaml.load(payload, Loader=SafeLoader) # OK
+    marshal.loads(payload) # NOT OK
 
     import dill
-    dill.loads(payload)
+    dill.loads(payload) # NOT OK

Original file line number	Diff line number	Diff line change
`@@ -60,5 +60,5 @@ private class YamlDeserialization extends UnmarshalingFunction::Range {`
`60`	`60`
`61`	`61`	`override DataFlow::Node getOutput() { result = this }`
`62`	`62`
`63`		`- override string getFormat() { none() }`
	`63`	`+ override string getFormat() { result = "YAML" }`
`64`	`64`	`}`