Taint underlying aggregates of protobuf messages when an element is written

For example, writing to a[b].c[d] taints 'a'.

Author: smowton

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

@@ -593,11 +593,21 @@ class ReadNode extends InstructionNode {
 /**
  * A data-flow node that reads the value of a field from a struct, or an element from an array, slice, map or string.
  */
-class ReadFromAggregateNode extends ReadNode {
+abstract class ReadFromAggregateNode extends ReadNode {
   /** Gets the data-flow node representing the base from which the field or element is read. */
   abstract Node getBase();
 }
 
+/**
+ * Gets the data-flow node representing the bottom of a stack of zero or more `ReadFromAggregateNode`s.
+ *
+ * For example, in the expression a.b[c].d[e], this would return the dataflow node for the read from `a`.
+ */
+Node getUnderlyingNode(ReadNode read) {
+  (result = read or result = read.(ReadFromAggregateNode).getBase+()) and
+  not result instanceof ReadFromAggregateNode
+}
+
 /**
  * A data-flow node that reads an element of an array, map, slice or string.
  */

ql/src/semmle/go/frameworks/Protobuf.qll

@@ -146,12 +146,13 @@ module Protobuf {
   }
 
   /**
-   * Additional taint step tainting a Message when taint is written to any of its fields.
+   * Additional taint step tainting a Message when taint is written to any of its fields and/or elements.
    */
   private class WriteMessageFieldStep extends TaintTracking::AdditionalTaintStep {
     override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
-      any(DataFlow::Write w)
-          .writesField(succ.(DataFlow::PostUpdateNode).getPreUpdateNode(), getAMessageField(), pred)
+      exists(DataFlow::ReadNode base | succ = DataFlow::getUnderlyingNode(base) |
+        any(DataFlow::Write w).writesField(base, getAMessageField(), pred)
+      )
     }
   }
 }