From 70f87b59d2fed87ee2ccfff4347a46069f4b839f Mon Sep 17 00:00:00 2001
From: Max Schaefer <max@semmle.com>
Date: Wed, 6 May 2020 09:34:23 +0100
Subject: [PATCH] Data flow: Support stores into nodes that are not
 `PostUpdateNode`s.

cf https://github.com/github/codeql/pull/3312
---
 .../go/dataflow/internal/DataFlowImpl.qll     |  4 +--
 .../dataflow/internal/DataFlowImplCommon.qll  | 25 ++++++++++++++-----
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/ql/src/semmle/go/dataflow/internal/DataFlowImpl.qll b/ql/src/semmle/go/dataflow/internal/DataFlowImpl.qll
index db0fbcf7130..9587ea5f274 100644
--- a/ql/src/semmle/go/dataflow/internal/DataFlowImpl.qll
+++ b/ql/src/semmle/go/dataflow/internal/DataFlowImpl.qll
@@ -1060,8 +1060,8 @@ private module LocalFlowBigStep {
       jumpStep(_, node, config) or
       additionalJumpStep(_, node, config) or
       node instanceof ParameterNode or
-      node instanceof OutNode or
-      node instanceof PostUpdateNode or
+      node instanceof OutNodeExt or
+      store(_, _, node) or
       read(_, _, node) or
       node instanceof CastNode
     )
diff --git a/ql/src/semmle/go/dataflow/internal/DataFlowImplCommon.qll b/ql/src/semmle/go/dataflow/internal/DataFlowImplCommon.qll
index b241a574c97..852f54974e2 100644
--- a/ql/src/semmle/go/dataflow/internal/DataFlowImplCommon.qll
+++ b/ql/src/semmle/go/dataflow/internal/DataFlowImplCommon.qll
@@ -415,8 +415,7 @@ class CastingNode extends Node {
   CastingNode() {
     this instanceof ParameterNode or
     this instanceof CastNode or
-    this instanceof OutNode or
-    this.(PostUpdateNode).getPreUpdateNode() instanceof ArgumentNode
+    this instanceof OutNodeExt
   }
 }
 
@@ -564,6 +563,18 @@ class ReturnNodeExt extends Node {
   }
 }
 
+/**
+ * A node to which data can flow from a call. Either an ordinary out node
+ * or a post-update node associated with a call argument.
+ */
+class OutNodeExt extends Node {
+  OutNodeExt() {
+    this instanceof OutNode
+    or
+    this.(PostUpdateNode).getPreUpdateNode() instanceof ArgumentNode
+  }
+}
+
 /**
  * An extended return kind. A return kind describes how data can be returned
  * from a callable. This can either be through a returned value or an updated
@@ -574,7 +585,7 @@ abstract class ReturnKindExt extends TReturnKindExt {
   abstract string toString();
 
   /** Gets a node corresponding to data flow out of `call`. */
-  abstract Node getAnOutNode(DataFlowCall call);
+  abstract OutNodeExt getAnOutNode(DataFlowCall call);
 }
 
 class ValueReturnKind extends ReturnKindExt, TValueReturn {
@@ -586,7 +597,9 @@ class ValueReturnKind extends ReturnKindExt, TValueReturn {
 
   override string toString() { result = kind.toString() }
 
-  override Node getAnOutNode(DataFlowCall call) { result = getAnOutNode(call, this.getKind()) }
+  override OutNodeExt getAnOutNode(DataFlowCall call) {
+    result = getAnOutNode(call, this.getKind())
+  }
 }
 
 class ParamUpdateReturnKind extends ReturnKindExt, TParamUpdate {
@@ -598,9 +611,9 @@ class ParamUpdateReturnKind extends ReturnKindExt, TParamUpdate {
 
   override string toString() { result = "param update " + pos }
 
-  override PostUpdateNode getAnOutNode(DataFlowCall call) {
+  override OutNodeExt getAnOutNode(DataFlowCall call) {
     exists(ArgumentNode arg |
-      result.getPreUpdateNode() = arg and
+      result.(PostUpdateNode).getPreUpdateNode() = arg and
       arg.argumentOf(call, this.getPosition())
     )
   }