Merge pull request github#10783 from yoff/python/subscript-nodes

Python: API graph improvements for subscripts

Author: tausbn

python/ql/lib/change-notes/2022-10-04-api-subscript-nodes.md

@@ -0,0 +1,7 @@
+---
+category: minorAnalysis
+---
+* Fixed labels in the API graph pertaining to definitions of subscripts. Previously, these were found by `getMember` rather than `getASubscript`.
+* Added edges for indices of subscripts to the API graph. Now a subscripted API node will have an edge to the API node for the index expression. So if `foo` is matched by API node `A`, then `"key"` in `foo["key"]` will be matched by the API node `A.getIndex()`. This can be used to track the origin of the index.
+* Added member predicate `getSubscriptAt(API::Node index)` to `API::Node`. Like `getASubscript()`, this will return an API node that matches a subscript of the node, but here it will be restricted to subscripts where the index matches the `index` parameter.
+* Added convenience predicate `getSubscript("key")` to obtain a subscript at a specific index, when the index happens to be a statically known string.

python/ql/lib/semmle/python/ApiGraphs.qll

@@ -249,6 +249,60 @@ module API {
      */
     Node getASubscript() { result = this.getASuccessor(Label::subscript()) }
 
+    /**
+     * Gets a node representing an index of a subscript of this node.
+     * For example, in `obj[x]`, `x` is an index of `obj`.
+     */
+    Node getIndex() { result = this.getASuccessor(Label::index()) }
+
+    /**
+     * Gets a node representing a subscript of this node at (string) index `key`.
+     * This requires that the index can be statically determined.
+     *
+     * For example, the subscripts of `a` and `b` below would be found using
+     * the index `foo`:
+     * ```py
+     * a["foo"]
+     * x = "foo" if cond else "bar"
+     * b[x]
+     * ```
+     */
+    Node getSubscript(string key) {
+      exists(API::Node index | result = this.getSubscriptAt(index) |
+        key = index.getAValueReachingSink().asExpr().(PY::StrConst).getText()
+      )
+    }
+
+    /**
+     * Gets a node representing a subscript of this node at index `index`.
+     */
+    Node getSubscriptAt(API::Node index) {
+      result = this.getASubscript() and
+      index = this.getIndex() and
+      (
+        // subscripting
+        exists(PY::SubscriptNode subscript |
+          subscript.getObject() = this.getAValueReachableFromSource().asCfgNode() and
+          subscript.getIndex() = index.asSink().asCfgNode()
+        |
+          // reading
+          subscript = result.asSource().asCfgNode()
+          or
+          // writing
+          subscript.(PY::DefinitionNode).getValue() = result.asSink().asCfgNode()
+        )
+        or
+        // dictionary literals
+        exists(PY::Dict dict, PY::KeyValuePair item |
+          dict = this.getAValueReachingSink().asExpr() and
+          dict.getItem(_) = item and
+          item.getKey() = index.asSink().asExpr()
+        |
+          item.getValue() = result.asSink().asExpr()
+        )
+      )
+    }
+
     /**
      * Gets a string representation of the lexicographically least among all shortest access paths
      * from the root to this node.
@@ -405,7 +459,7 @@ module API {
   Node builtin(string n) { result = moduleImport("builtins").getMember(n) }
 
   /**
-   * An `CallCfgNode` that is connected to the API graph.
+   * A `CallCfgNode` that is connected to the API graph.
    *
    * Can be used to reason about calls to an external API in which the correlation between
    * parameters and/or return values must be retained.
@@ -694,12 +748,31 @@ module API {
           rhs = aw.getValue()
         )
         or
-        // TODO: I had expected `DataFlow::AttrWrite` to contain the attribute writes from a dict, that's how JS works.
+        // dictionary literals
         exists(PY::Dict dict, PY::KeyValuePair item |
           dict = pred.(DataFlow::ExprNode).getNode().getNode() and
-          dict.getItem(_) = item and
-          lbl = Label::member(item.getKey().(PY::StrConst).getS()) and
-          rhs.(DataFlow::ExprNode).getNode().getNode() = item.getValue()
+          dict.getItem(_) = item
+        |
+          // from `x` to `{ "key": x }`
+          // TODO: once convenient, this should be done at a higher level than the AST,
+          // at least at the CFG layer, to take splitting into account.
+          rhs.(DataFlow::ExprNode).getNode().getNode() = item.getValue() and
+          lbl = Label::subscript()
+          or
+          // from `"key"` to `{ "key": x }`
+          // TODO: once convenient, this should be done at a higher level than the AST,
+          // at least at the CFG layer, to take splitting into account.
+          rhs.(DataFlow::ExprNode).getNode().getNode() = item.getKey() and
+          lbl = Label::index()
+        )
+        or
+        // list literals, from `x` to `[x]`
+        // TODO: once convenient, this should be done at a higher level than the AST,
+        // at least at the CFG layer, to take splitting into account.
+        // Also consider `SequenceNode for generality.
+        exists(PY::List list | list = pred.(DataFlow::ExprNode).getNode().getNode() |
+          rhs.(DataFlow::ExprNode).getNode().getNode() = list.getAnElt() and
+          lbl = Label::subscript()
         )
         or
         exists(PY::CallableExpr fn | fn = pred.(DataFlow::ExprNode).getNode().getNode() |
@@ -720,6 +793,20 @@ module API {
         lbl = Label::memberFromRef(aw)
       )
       or
+      // subscripting
+      exists(DataFlow::LocalSourceNode src, DataFlow::Node subscript, DataFlow::Node index |
+        use(base, src) and
+        subscript = trackUseNode(src).getSubscript(index)
+      |
+        // from `x` to a definition of `x[...]`
+        rhs.asCfgNode() = subscript.asCfgNode().(PY::DefinitionNode).getValue() and
+        lbl = Label::subscript()
+        or
+        // from `x` to `"key"` in `x["key"]`
+        rhs = index and
+        lbl = Label::index()
+      )
+      or
       exists(EntryPoint entry |
         base = root() and
         lbl = Label::entryPoint(entry) and
@@ -757,7 +844,8 @@ module API {
         or
         // Subscripting a node that is a use of `base`
         lbl = Label::subscript() and
-        ref = pred.getASubscript()
+        ref = pred.getSubscript(_) and
+        ref.asCfgNode().isLoad()
         or
         // Subclassing a node
         lbl = Label::subclass() and
@@ -973,8 +1061,7 @@ module API {
           member = any(DataFlow::AttrRef pr).getAttributeName() or
           exists(Builtins::likelyBuiltin(member)) or
           ImportStar::namePossiblyDefinedInImportStar(_, member, _) or
-          Impl::prefix_member(_, member, _) or
-          member = any(PY::Dict d).getAnItem().(PY::KeyValuePair).getKey().(PY::StrConst).getS()
+          Impl::prefix_member(_, member, _)
         } or
         MkLabelUnknownMember() or
         MkLabelParameter(int i) {
@@ -992,6 +1079,7 @@ module API {
         MkLabelSubclass() or
         MkLabelAwait() or
         MkLabelSubscript() or
+        MkLabelIndex() or
         MkLabelEntryPoint(EntryPoint ep)
 
       /** A label for a module. */
@@ -1072,6 +1160,11 @@ module API {
         override string toString() { result = "getASubscript()" }
       }
 
+      /** A label that gets the index of a subscript. */
+      class LabelIndex extends ApiLabel, MkLabelIndex {
+        override string toString() { result = "getIndex()" }
+      }
+
       /** A label for entry points. */
       class LabelEntryPoint extends ApiLabel, MkLabelEntryPoint {
         private EntryPoint entry;
@@ -1120,6 +1213,9 @@ module API {
     /** Gets the `subscript` edge label. */
     LabelSubscript subscript() { any() }
 
+    /** Gets the `subscript` edge label. */
+    LabelIndex index() { any() }
+
     /** Gets the label going from the root node to the nodes associated with the given entry point. */
     LabelEntryPoint entryPoint(EntryPoint ep) { result = MkLabelEntryPoint(ep) }
   }

python/ql/lib/semmle/python/dataflow/new/internal/LocalSources.qll

@@ -104,7 +104,7 @@ class LocalSourceNode extends Node {
   /**
    * Gets a subscript of this node.
    */
-  Node getASubscript() { Cached::subscript(this, result) }
+  Node getSubscript(Node index) { Cached::subscript(this, result, index) }
 
   /**
    * Gets a call to the method `methodName` on this node.
@@ -249,13 +249,14 @@ private module Cached {
   }
 
   /**
-   * Holds if `node` flows to a sequence/mapping of which `subscript` is a subscript.
+   * Holds if `node` flows to a sequence/mapping of which `subscript` is a subscript with index/key `index`.
    */
   cached
-  predicate subscript(LocalSourceNode node, CfgNode subscript) {
+  predicate subscript(LocalSourceNode node, CfgNode subscript, CfgNode index) {
     exists(CfgNode seq, SubscriptNode subscriptNode | subscriptNode = subscript.getNode() |
       node.flowsTo(seq) and
-      seq.getNode() = subscriptNode.getObject()
+      seq.getNode() = subscriptNode.getObject() and
+      index.getNode() = subscriptNode.getIndex()
     )
   }
 }

python/ql/lib/semmle/python/frameworks/Aiohttp.qll

@@ -621,15 +621,12 @@ module AiohttpWebModel {
     DataFlow::Node value;
 
     AiohttpResponseCookieSubscriptWrite() {
-      exists(SubscriptNode subscript |
+      exists(API::Node i |
+        value = aiohttpResponseInstance().getMember("cookies").getSubscriptAt(i).asSink() and
+        index = i.asSink() and
         // To give `this` a value, we need to choose between either LHS or RHS,
-        // and just go with the LHS
-        this.asCfgNode() = subscript
-      |
-        subscript.getObject() =
-          aiohttpResponseInstance().getMember("cookies").getAValueReachableFromSource().asCfgNode() and
-        value.asCfgNode() = subscript.(DefinitionNode).getValue() and
-        index.asCfgNode() = subscript.getIndex()
+        // and just go with the RHS as it is readily available
+        this = value
       )
     }
 

python/ql/src/experimental/semmle/python/frameworks/Django.qll

@@ -91,14 +91,10 @@ private module ExperimentalPrivateDjango {
             result = baseClassRef().getReturn().getAMember()
           }
 
-          /** Gets a reference to a header instance call with `__setitem__`. */
-          API::Node headerSetItem() {
-            result = headerInstance() and
-            result.asSource().(DataFlow::AttrRead).getAttributeName() = "__setitem__"
-          }
-
           class DjangoResponseSetItemCall extends DataFlow::CallCfgNode, HeaderDeclaration::Range {
-            DjangoResponseSetItemCall() { this = headerSetItem().getACall() }
+            DjangoResponseSetItemCall() {
+              this = baseClassRef().getReturn().getMember("__setitem__").getACall()
+            }
 
             override DataFlow::Node getNameArg() { result = this.getArg(0) }
 
@@ -109,8 +105,7 @@ private module ExperimentalPrivateDjango {
             DataFlow::Node headerInput;
 
             DjangoResponseDefinition() {
-              this.asCfgNode().(DefinitionNode) =
-                headerInstance().getAValueReachableFromSource().asCfgNode() and
+              headerInput = headerInstance().asSink() and
               headerInput.asCfgNode() = this.asCfgNode().(DefinitionNode).getValue()
             }
 

python/ql/src/experimental/semmle/python/frameworks/Sendgrid.qll

@@ -26,7 +26,7 @@ private module Sendgrid {
   }
 
   /** Gets a reference to a `SendGridAPIClient` instance call with `send` or `post`. */
-  private DataFlow::CallCfgNode sendgridApiSendCall() {
+  private API::CallNode sendgridApiSendCall() {
     result = sendgridApiClient().getMember("send").getACall()
     or
     result =
@@ -62,7 +62,7 @@ private module Sendgrid {
    * * `getFrom()`'s result would be `"[email protected]"`.
    * * `getSubject()`'s result would be `"Sending with SendGrid is Fun"`.
    */
-  private class SendGridMail extends DataFlow::CallCfgNode, EmailSender::Range {
+  private class SendGridMail extends API::CallNode, EmailSender::Range {
     SendGridMail() { this = sendgridApiSendCall() }
 
     private DataFlow::CallCfgNode getMailCall() {
@@ -118,40 +118,28 @@ private module Sendgrid {
       or
       result = this.sendgridWrite("html_content")
       or
-      exists(KeyValuePair content, Dict generalDict, KeyValuePair typePair, KeyValuePair valuePair |
-        content.getKey().(StrConst).getText() = "content" and
-        content.getValue().(List).getAnElt() = generalDict and
-        // declare KeyValuePairs keys and values
-        typePair.getKey().(StrConst).getText() = "type" and
-        typePair.getValue().(StrConst).getText() = ["text/html", "text/x-amp-html"] and
-        valuePair.getKey().(StrConst).getText() = "value" and
-        result.asExpr() = valuePair.getValue() and
-        // correlate generalDict with previously set KeyValuePairs
-        generalDict.getAnItem() in [typePair, valuePair] and
-        [this.getArg(0), this.getArgByName("request_body")].getALocalSource().asExpr() =
-          any(Dict d | d.getAnItem() = content)
+      exists(API::Node contentElement |
+        contentElement =
+          this.getKeywordParameter("request_body").getSubscript("content").getASubscript()
+      |
+        contentElement.getSubscript("type").getAValueReachingSink().asExpr().(StrConst).getText() =
+          ["text/html", "text/x-amp-html"] and
+        result = contentElement.getSubscript("value").getAValueReachingSink()
       )
       or
-      exists(KeyValuePair footer, Dict generalDict, KeyValuePair enablePair, KeyValuePair htmlPair |
-        footer.getKey().(StrConst).getText() = ["footer", "subscription_tracking"] and
-        footer.getValue() = generalDict and
-        // check footer is enabled
-        enablePair.getKey().(StrConst).getText() = "enable" and
-        exists(enablePair.getValue().(True)) and
-        // get html content
-        htmlPair.getKey().(StrConst).getText() = "html" and
-        result.asExpr() = htmlPair.getValue() and
-        // correlate generalDict with previously set KeyValuePairs
-        generalDict.getAnItem() in [enablePair, htmlPair] and
-        exists(KeyValuePair k |
-          k.getKey() =
-            [this.getArg(0), this.getArgByName("request_body")]
-                .getALocalSource()
-                .asExpr()
-                .(Dict)
-                .getAKey() and
-          k.getValue() = any(Dict d | d.getAKey() = footer.getKey())
-        )
+      exists(API::Node html |
+        html =
+          this.getKeywordParameter("request_body")
+              .getSubscript("tracking_settings")
+              .getSubscript("subscription_tracking")
+        or
+        html =
+          this.getKeywordParameter("request_body")
+              .getSubscript("mail_settings")
+              .getSubscript("footer")
+      |
+        html.getSubscript("enable").getAValueReachingSink().asExpr() instanceof True and
+        result = html.getSubscript("html").getAValueReachingSink()
       )
     }