From e41d6099216cd5ba33766fbfef04d809e6091e4e Mon Sep 17 00:00:00 2001
From: Sauyon Lee <sauyon@github.com>
Date: Wed, 22 Sep 2021 11:16:19 -0700
Subject: [PATCH] Use newtype for SourceOrSinkElement

---
 ql/lib/semmle/go/dataflow/ExternalFlow.qll    | 10 ++---
 .../internal/FlowSummaryImplSpecific.qll      | 45 +++++++++++++------
 2 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/ql/lib/semmle/go/dataflow/ExternalFlow.qll b/ql/lib/semmle/go/dataflow/ExternalFlow.qll
index 889600de9c2..68b88118355 100644
--- a/ql/lib/semmle/go/dataflow/ExternalFlow.qll
+++ b/ql/lib/semmle/go/dataflow/ExternalFlow.qll
@@ -356,23 +356,23 @@ private string paramsStringPart(Function f, int i) {
 string paramsString(Function f) { result = concat(int i | | paramsStringPart(f, i) order by i) }
 
 /** Gets the source/sink/summary element corresponding to the supplied parameters. */
-AstNode interpretElement(
+SourceOrSinkElement interpretElement(
   string pkg, string type, boolean subtypes, string name, string signature, string ext
 ) {
   elementSpec(pkg, type, subtypes, name, signature, ext) and
   // Go does not need to distinguish functions with signature
   signature = "" and
   (
-    exists(Field f | f.hasQualifiedName(pkg, type, name) | result = f.getDeclaration())
+    exists(Field f | f.hasQualifiedName(pkg, type, name) | result.asEntity() = f)
     or
     exists(Method m | m.hasQualifiedName(pkg, type, name) |
-      result = m.getFuncDecl()
+      result.asEntity() = m
       or
-      subtypes = true and result.(FuncDecl).getFunction().(Method).implements(m)
+      subtypes = true and result.asEntity().(Method).implements(m)
     )
     or
     type = "" and
-    exists(Entity e | e.hasQualifiedName(pkg, name) | result = e.getDeclaration())
+    exists(Entity e | e.hasQualifiedName(pkg, name) | result.asEntity() = e)
   )
 }
 
diff --git a/ql/lib/semmle/go/dataflow/internal/FlowSummaryImplSpecific.qll b/ql/lib/semmle/go/dataflow/internal/FlowSummaryImplSpecific.qll
index 5acdf5cd482..5c1427edf36 100644
--- a/ql/lib/semmle/go/dataflow/internal/FlowSummaryImplSpecific.qll
+++ b/ql/lib/semmle/go/dataflow/internal/FlowSummaryImplSpecific.qll
@@ -52,7 +52,11 @@ predicate summaryElement(DataFlowCallable c, string input, string output, string
     string namespace, string type, boolean subtypes, string name, string signature, string ext
   |
     summaryModel(namespace, type, subtypes, name, signature, ext, input, output, kind) and
-    c = interpretElement(namespace, type, subtypes, name, signature, ext)
+    c =
+      interpretElement(namespace, type, subtypes, name, signature, ext)
+          .asEntity()
+          .(Function)
+          .getFuncDecl()
   )
 }
 
@@ -62,7 +66,24 @@ SummaryComponent interpretComponentSpecific(string c) {
   exists(Content content | parseContent(c, content) and result = SummaryComponent::content(content))
 }
 
-class SourceOrSinkElement = AstNode;
+private newtype TSourceOrSinkElement =
+  TEntityElement(Entity e) or
+  TAstElement(AstNode n)
+
+class SourceOrSinkElement extends TSourceOrSinkElement {
+  Entity asEntity() { this = TEntityElement(result) }
+
+  AstNode asAstNode() { this = TAstElement(result) }
+
+  string toString() {
+    result = "element representing " + [this.asEntity().toString(), this.asAstNode().toString()]
+  }
+
+  predicate hasLocationInfo(string fp, int sl, int sc, int el, int ec) {
+    this.asEntity().hasLocationInfo(fp, sl, sc, el, ec) or
+    this.asAstNode().hasLocationInfo(fp, sl, sc, el, ec)
+  }
+}
 
 /**
  * Holds if an external source specification exists for `e` with output specification
@@ -106,16 +127,14 @@ class InterpretNode extends TInterpretNode {
   Node asNode() { this = TNode(result) }
 
   /** Gets the call that this node corresponds to, if any. */
-  DataFlowCall asCall() { result = this.asElement() }
+  DataFlowCall asCall() { result = this.asElement().asAstNode() }
 
   /** Gets the callable that this node corresponds to, if any. */
-  DataFlowCallable asCallable() { result = this.asElement() }
+  DataFlowCallable asCallable() { result = this.asElement().asEntity().(Function).getFuncDecl() }
 
   /** Gets the target of this call, if any. */
-  Ident getCallTarget() {
-    exists(Function f | f = this.asCall().getNode().(DataFlow::CallNode).getTarget() |
-      result = f.getDeclaration()
-    )
+  SourceOrSinkElement getCallTarget() {
+    result.asEntity() = this.asCall().getNode().(DataFlow::CallNode).getTarget()
   }
 
   /** Gets a textual representation of this node. */
@@ -136,15 +155,15 @@ class InterpretNode extends TInterpretNode {
 /** Provides additional sink specification logic required for annotations. */
 pragma[inline]
 predicate interpretOutputSpecific(string c, InterpretNode mid, InterpretNode node) {
-  exists(Node n, AstNode ast |
+  exists(Node n, SourceOrSinkElement e |
     n = node.asNode() and
-    ast = mid.asElement()
+    e = mid.asElement()
   |
     (c = "Parameter" or c = "") and
-    node.asNode().asParameter().getDeclaration() = ast
+    node.asNode().asParameter() = e.asEntity()
     or
     c = "" and
-    n.(DataFlow::FieldReadNode).getField().getDeclaration() = ast
+    n.(DataFlow::FieldReadNode).getField() = e.asEntity()
   )
 }
 
@@ -153,7 +172,7 @@ pragma[inline]
 predicate interpretInputSpecific(string c, InterpretNode mid, InterpretNode n) {
   exists(DataFlow::Write fw, Field f |
     c = "" and
-    f.getDeclaration() = mid.asElement() and
+    f = mid.asElement().asEntity() and
     fw.writesField(_, f, n.asNode())
   )
 }