Add data-flow edge from -> to in the context to, ok := from.(*Type)

ql/src/semmle/go/controlflow/ControlFlowGraphImpl.qll

@@ -90,7 +90,7 @@ newtype TControlFlowNode =
    */
   MkCompoundAssignRhsNode(CompoundAssignStmt assgn) or
   /**
-   * A control-flow node that represents the `i`th component of a tuple expression `base`.
+   * A control-flow node that represents the `i`th component of a tuple expression `s`.
    */
   MkExtractNode(AstNode s, int i) {
     // in an assignment `x, y, z = tuple`

ql/src/semmle/go/controlflow/IR.qll

@@ -628,7 +628,8 @@ module IR {
 
   /**
    * An instruction selecting one of multiple values returned by a function, or either the key
-   * or the value of the iterator in a range loop.
+   * or the value of the iterator in a range loop, or the result or success value from a type
+   * assertion.
    */
   class ExtractTupleElementInstruction extends Instruction, MkExtractNode {
     AstNode s;
@@ -660,6 +661,10 @@ module IR {
         result = c.getTarget().getResultType(i)
       )
       or
+      exists(TypeAssertExpr tae | getBase() = evalExprInstruction(tae) |
+        result = tae.getType().(TupleType).getComponentType(i)
+      )
+      or
       exists(Type rangeType | rangeType = s.(RangeStmt).getDomain().getType().getUnderlyingType() |
         exists(Type baseType |
           baseType = rangeType.(ArrayType).getElementType() or

ql/src/semmle/go/dataflow/internal/DataFlowUtil.qll

@@ -945,7 +945,7 @@ SsaNode ssaNode(SsaVariable v) { result.getDefinition() = v.getDefinition() }
 
 /**
  * Gets the data-flow node corresponding to the `i`th element of tuple `t` (which is either a call
- * with multiple results or an iterator in a range loop).
+ * with multiple results, an iterator in a range loop, or the result of a type assertion).
  */
 Node extractTupleElement(Node t, int i) {
   exists(IR::Instruction insn | t = instructionNode(insn) |
@@ -981,13 +981,23 @@ private predicate basicLocalFlowStep(Node nodeFrom, Node nodeTo) {
   // Instruction -> Instruction
   exists(Expr pred, Expr succ |
     succ.(LogicalBinaryExpr).getAnOperand() = pred or
-    succ.(ConversionExpr).getOperand() = pred or
-    succ.(TypeAssertExpr).getExpr() = pred
+    succ.(ConversionExpr).getOperand() = pred
   |
     nodeFrom = exprNode(pred) and
     nodeTo = exprNode(succ)
   )
   or
+  // Type assertion: if in the context `checked, ok := e.(*Type)` (in which
+  // case tuple-extraction instructions exist), flow from `e` to `e.(*Type)[0]`;
+  // otherwise flow from `e` to `e.(*Type)`.
+  exists(IR::Instruction evalAssert, TypeAssertExpr assert |
+    nodeFrom.asExpr() = assert.getExpr() and
+    evalAssert = IR::evalExprInstruction(assert) and
+    if exists(IR::extractTupleElement(evalAssert, _))
+    then nodeTo.asInstruction() = IR::extractTupleElement(evalAssert, 0)
+    else nodeTo.asInstruction() = evalAssert
+  )
+  or
   // Instruction -> SSA
   exists(IR::Instruction pred, SsaExplicitDefinition succ |
     succ.getRhs() = pred and

ql/test/library-tests/semmle/go/dataflow/FlowSteps/LocalFlowStep.expected

@@ -76,7 +76,7 @@
 | main.go:26:2:26:17 | ... := ...[0] | main.go:26:2:26:2 | definition of n |
 | main.go:26:2:26:17 | ... := ...[1] | main.go:26:5:26:6 | definition of ok |
 | main.go:26:5:26:6 | definition of ok | main.go:27:5:27:6 | ok |
-| main.go:26:11:26:11 | x | main.go:26:11:26:17 | type assertion |
+| main.go:26:11:26:11 | x | main.go:26:2:26:17 | ... := ...[0] |
 | main.go:38:2:38:2 | definition of s | main.go:39:15:39:15 | s |
 | main.go:38:2:38:2 | definition of s | main.go:40:15:40:15 | s |
 | main.go:38:2:38:2 | definition of s | main.go:42:7:42:7 | s |

ql/test/library-tests/semmle/go/dataflow/TypeAssertions/DataFlow.ql

@@ -0,0 +1,29 @@
+import go
+import TestUtilities.InlineExpectationsTest
+
+class Configuration extends DataFlow::Configuration {
+  Configuration() { this = "test-configuration" }
+
+  override predicate isSource(DataFlow::Node source) {
+    source = any(DataFlow::CallNode c | c.getCalleeName() = "src").getResult(0)
+  }
+
+  override predicate isSink(DataFlow::Node sink) {
+    sink = any(DataFlow::CallNode c | c.getCalleeName() = "sink").getArgument(0)
+  }
+}
+
+class DataFlowTest extends InlineExpectationsTest {
+  DataFlowTest() { this = "DataFlowTest" }
+
+  override string getARelevantTag() { result = "dataflow" }
+
+  override predicate hasActualResult(string file, int line, string element, string tag, string value) {
+    tag = "dataflow" and
+    exists(DataFlow::Node sink | any(Configuration c).hasFlow(_, sink) |
+      element = sink.toString() and
+      value = sink.toString() and
+      sink.hasLocationInfo(file, line, _, _, _)
+    )
+  }
+}

ql/test/library-tests/semmle/go/dataflow/TypeAssertions/test.go

@@ -0,0 +1,30 @@
+package main
+
+func src() interface{} {
+	return "hi"
+}
+
+func sink(p interface{}) {}
+
+func test() (bool, *string) {
+	ptr := src()
+	sink(ptr) // $dataflow=ptr
+	cast := ptr.(*string)
+	sink(cast) // $dataflow=cast
+	cast2, ok := ptr.(*string)
+	if !ok {
+		return true, nil
+	}
+	sink(cast2) // $dataflow=cast2
+	var cast3, ok2 = ptr.(*string)
+	if !ok2 {
+		return true, nil
+	}
+	sink(cast3) // $dataflow=cast3
+	cast2, ok = ptr.(*string)
+	if !ok {
+		return true, nil
+	}
+	sink(cast2) // $dataflow=cast2
+	return true, nil
+}