Revamp the query to implement AdditionalTaintStep

java/ql/src/semmle/code/java/dataflow/FlowSources.qll

@@ -320,46 +320,31 @@ class AndroidIntentInput extends DataFlow::Node {
   }
 }
 
-/** Method access to external inputs of `android.content.Intent` object. */
+/** Method access to external inputs of `android.content.Intent` or `android.os.BaseBundle` object. */
 class IntentGetExtraMethodAccess extends MethodAccess {
   IntentGetExtraMethodAccess() {
     exists(AndroidComponent ac |
-      this.getEnclosingCallable().getDeclaringType() = ac and ac.isExported()
-    ) and
-    (
+      this.getEnclosingCallable().getDeclaringType() = ac and
+      ac.isExported() and
       this.getMethod().getName().regexpMatch("get\\w+Extra") and
       this.getMethod().getDeclaringType() instanceof TypeIntent
-      or
-      this.getMethod().getName().regexpMatch("get\\w+") and
-      this.getQualifier().(MethodAccess).getMethod().hasName("getExtras") and
-      this.getQualifier().(MethodAccess).getMethod().getDeclaringType() instanceof TypeIntent
     )
+    or
+    this.getMethod().getName().regexpMatch("get\\w+") and
+    this
+        .getMethod()
+        .getDeclaringType()
+        .getASupertype*()
+        .hasQualifiedName("android.os", "BaseBundle")
   }
 }
 
 /** Android intent extra source. */
 private class AndroidIntentExtraSource extends RemoteFlowSource {
   AndroidIntentExtraSource() {
-    exists(MethodAccess ma |
-      ma instanceof IntentGetExtraMethodAccess and
-      this.asExpr() = ma and
-      exists(AndroidIntentInput inode |
-        (
-          ma.getQualifier() = inode.asExpr() or // extra from intent
-          ma.getQualifier() = inode.asParameter().getAnAccess()
-        )
-        or
-        exists(
-          MethodAccess ema // extra from extras bundle of intent
-        |
-          ema.getMethod().hasName("getExtras") and
-          ma.getQualifier() = ema and
-          (
-            ema.getQualifier() = inode.asExpr() or
-            ema.getQualifier() = inode.asParameter().getAnAccess()
-          )
-        )
-      )
+    exists(AndroidIntentInput inode |
+      this.asExpr() = inode.asExpr() or
+      this.asExpr() = inode.asParameter().getAnAccess()
     )
   }
 

java/ql/src/semmle/code/java/dataflow/FlowSteps.qll

@@ -4,6 +4,7 @@
 
 private import java
 private import semmle.code.java.dataflow.DataFlow
+private import semmle.code.java.dataflow.FlowSources
 
 /**
  * A module importing the frameworks that implement additional flow steps,
@@ -139,3 +140,31 @@ private class StringBuilderTaintPreservingCallable extends TaintPreservingCallab
     sink = -1
   }
 }
+
+/**
+ * Holds if `n1` to `n2` is a dataflow step between the extra getter method and its caller Android `Intent` or `Bundle`.
+ */
+private predicate intentExtraStep(DataFlow::ExprNode n1, DataFlow::ExprNode n2) {
+  exists(IntentGetExtraMethodAccess ma |
+    n1.asExpr() = ma.getQualifier() and
+    n2.asExpr() = ma
+  )
+}
+
+/**
+ * Holds if `n1` to `n2` is a dataflow step from Android `Intent` to its `getExtras` method.
+ */
+private predicate bundleExtraStep(DataFlow::ExprNode n1, DataFlow::ExprNode n2) {
+  exists(MethodAccess ma | ma.getMethod().hasName("getExtras") |
+    n1.asExpr() = ma.getQualifier() and
+    n2.asExpr() = ma
+  )
+}
+
+/** A set of additional taint steps to consider when taint tracking Android intent extra related data flows. */
+class AndroidExtraSourceAdditionalTaintStep extends AdditionalTaintStep {
+  override predicate step(DataFlow::Node node1, DataFlow::Node node2) {
+    intentExtraStep(node1, node2) or
+    bundleExtraStep(node1, node2)
+  }
+}