Python: Prepare for general content in type-tracker

RasmusWL · RasmusWL · commit fc8caa66c8ac · 2024-03-14T10:42:37.000+01:00
Due to the char-pred of Content, this change should keep exactly the
same behavior as before.
diff --git a/python/ql/lib/semmle/python/dataflow/new/TypeTracking.qll b/python/ql/lib/semmle/python/dataflow/new/TypeTracking.qll
@@ -5,6 +5,7 @@
 
 private import internal.TypeTrackingImpl as Impl
 import Impl::Shared::TypeTracking<Impl::TypeTrackingInput>
+private import semmle.python.dataflow.new.internal.DataFlowPublic as DataFlowPublic
 
 /** A string that may appear as the name of an attribute or access path. */
 class AttributeName = Impl::TypeTrackingInput::Content;
@@ -40,17 +41,20 @@ class TypeTracker extends Impl::TypeTracker {
    * Holds if this is the starting point of type tracking, and the value starts in the attribute named `attrName`.
    * The type tracking only ends after the attribute has been loaded.
    */
-  predicate startInAttr(string attrName) { this.startInContent(attrName) }
+  predicate startInAttr(string attrName) {
+    exists(DataFlowPublic::AttributeContent content | content.getAttribute() = attrName |
+      this.startInContent(content)
+    )
+  }
 
   /**
    * INTERNAL. DO NOT USE.
    *
    * Gets the attribute associated with this type tracker.
    */
   string getAttr() {
-    result = this.getContent().asSome()
-    or
-    this.getContent().isNone() and
-    result = ""
+    if this.getContent().asSome() instanceof DataFlowPublic::AttributeContent
+    then result = this.getContent().asSome().(DataFlowPublic::AttributeContent).getAttribute()
+    else result = ""
   }
 }
diff --git a/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll b/python/ql/lib/semmle/python/dataflow/new/internal/TypeTrackingImpl.qll
@@ -97,24 +97,14 @@ private module SummaryTypeTrackerInput implements SummaryTypeTracker::Input {
 
 private module TypeTrackerSummaryFlow = SummaryTypeTracker::SummaryFlow<SummaryTypeTrackerInput>;
 
-/**
- * Gets the name of a possible piece of content. For Python, this is currently only attribute names,
- * using the name of the attribute for the corresponding content.
- */
-private string getPossibleContentName() {
-  Stages::TypeTracking::ref() and // the TypeTracking::append() etc. predicates that we want to cache depend on this predicate, so we can place the `ref()` call here to get around identical files.
-  result = any(DataFlowPublic::AttrRef a).getAttributeName()
-}
-
 module TypeTrackingInput implements Shared::TypeTrackingInput {
   class Node = DataFlowPublic::Node;
 
   class LocalSourceNode = DataFlowPublic::LocalSourceNode;
 
-  class Content instanceof string {
-    Content() { this = getPossibleContentName() }
-
-    string toString() { result = this }
+  class Content extends DataFlowPublic::Content {
+    // this char-pred is just a temporary restriction while transitioning to more general content
+    Content() { this instanceof DataFlowPublic::AttributeContent }
   }
 
   /**
@@ -181,46 +171,35 @@ module TypeTrackingInput implements Shared::TypeTrackingInput {
    * Holds if `nodeFrom` is being written to the `content` content of the object in `nodeTo`.
    */
   predicate storeStep(Node nodeFrom, Node nodeTo, Content content) {
-    exists(DataFlowPublic::AttrWrite a |
-      a.mayHaveAttributeName(content) and
+    exists(DataFlowPublic::AttrWrite a, string attrName |
+      content.(DataFlowPublic::AttributeContent).getAttribute() = attrName and
+      a.mayHaveAttributeName(attrName) and
       nodeFrom = a.getValue() and
       nodeTo = a.getObject()
     )
     or
-    exists(DataFlowPublic::ContentSet contents |
-      contents.(DataFlowPublic::AttributeContent).getAttribute() = content
-    |
-      TypeTrackerSummaryFlow::basicStoreStep(nodeFrom, nodeTo, contents)
-    )
+    TypeTrackerSummaryFlow::basicStoreStep(nodeFrom, nodeTo, content)
   }
 
   /**
    * Holds if `nodeTo` is the result of accessing the `content` content of `nodeFrom`.
    */
   predicate loadStep(Node nodeFrom, LocalSourceNode nodeTo, Content content) {
-    exists(DataFlowPublic::AttrRead a |
-      a.mayHaveAttributeName(content) and
+    exists(DataFlowPublic::AttrRead a, string attrName |
+      content.(DataFlowPublic::AttributeContent).getAttribute() = attrName and
+      a.mayHaveAttributeName(attrName) and
       nodeFrom = a.getObject() and
       nodeTo = a
     )
     or
-    exists(DataFlowPublic::ContentSet contents |
-      contents.(DataFlowPublic::AttributeContent).getAttribute() = content
-    |
-      TypeTrackerSummaryFlow::basicLoadStep(nodeFrom, nodeTo, contents)
-    )
+    TypeTrackerSummaryFlow::basicLoadStep(nodeFrom, nodeTo, content)
   }
 
   /**
    * Holds if the `loadContent` of `nodeFrom` is stored in the `storeContent` of `nodeTo`.
    */
   predicate loadStoreStep(Node nodeFrom, Node nodeTo, Content loadContent, Content storeContent) {
-    exists(DataFlowPublic::ContentSet loadContents, DataFlowPublic::ContentSet storeContents |
-      loadContents.(DataFlowPublic::AttributeContent).getAttribute() = loadContent and
-      storeContents.(DataFlowPublic::AttributeContent).getAttribute() = storeContent
-    |
-      TypeTrackerSummaryFlow::basicLoadStoreStep(nodeFrom, nodeTo, loadContents, storeContents)
-    )
+    TypeTrackerSummaryFlow::basicLoadStoreStep(nodeFrom, nodeTo, loadContent, storeContent)
   }
 
   /**
