Ruby: Derive edge labels from {Argument,Parameter}Position

asgerf · asgerf · commit 5e350a0270c0 · 2022-02-04T11:20:42.000+01:00
diff --git a/ruby/ql/lib/codeql/ruby/ApiGraphs.qll b/ruby/ql/lib/codeql/ruby/ApiGraphs.qll
@@ -11,6 +11,8 @@ import codeql.ruby.DataFlow
 import codeql.ruby.typetracking.TypeTracker
 import codeql.ruby.ast.internal.Module
 private import codeql.ruby.controlflow.CfgNodes
+private import codeql.ruby.dataflow.internal.DataFlowPrivate as DataFlowPrivate
+private import codeql.ruby.dataflow.internal.DataFlowDispatch as DataFlowDispatch
 
 /**
  * Provides classes and predicates for working with APIs used in a database.
@@ -423,49 +425,56 @@ module API {
     /** Gets a data flow node that flows to the RHS of a def-node. */
     private DataFlow::LocalSourceNode defCand() { result = defCand(TypeBackTracker::end()) }
 
-    /**
-     * Holds if there should be a `lbl`-edge from the given call to an argument.
-     */
-    pragma[nomagic]
-    private predicate argumentStep(string lbl, DataFlow::CallNode call, DataFlow::Node argument) {
+    private string getLabelFromArgumentPosition(DataFlowDispatch::ArgumentPosition pos) {
+      exists(int n |
+        pos.isPositional(n) and
+        result = Label::parameter(n)
+      )
+      or
+      exists(string name |
+        pos.isKeyword(name) and
+        result = Label::keywordParameter(name)
+      )
+      or
+      pos.isBlock() and
+      result = Label::blockParameter()
+    }
+
+    private string getLabelFromParameterPosition(DataFlowDispatch::ParameterPosition pos) {
       exists(int n |
-        argument = call.getArgument(n) and
-        lbl = Label::parameter(n)
+        pos.isPositional(n) and
+        result = Label::parameter(n)
       )
       or
       exists(string name |
-        argument = call.getKeywordArgument(name) and
-        lbl = Label::keywordParameter(name)
+        pos.isKeyword(name) and
+        result = Label::keywordParameter(name)
       )
       or
-      argument = call.getBlock() and
-      lbl = Label::blockParameter()
+      pos.isBlock() and
+      result = Label::blockParameter()
     }
 
     /**
-     * Holds if there should be a `lbl`-edge from the given callable to a parameter.
+     * Holds if there should be a `lbl`-edge from the given call to an argument.
      */
     pragma[nomagic]
-    private predicate parameterStep(string lbl, DataFlow::Node callable, DataFlow::Node paramNode) {
-      exists(Parameter param |
-        paramNode.asParameter() = param and
-        callable.asExpr().getExpr().(Callable).getAParameter() = param and
-        lbl = getLabelFromParameter(param)
+    private predicate argumentStep(string lbl, DataFlow::CallNode call, DataFlowPrivate::ArgumentNode argument) {
+      exists(DataFlowDispatch::ArgumentPosition argPos |
+        argument.sourceArgumentOf(call.asExpr(), argPos) and
+        lbl = getLabelFromArgumentPosition(argPos)
       )
     }
 
-    private string getLabelFromParameter(Parameter param) {
-      result = Label::keywordParameter(param.(KeywordParameter).getName())
-      or
-      param instanceof BlockParameter and
-      result = Label::blockParameter()
-      or
-      (
-        param instanceof SimpleParameter
-        or
-        param instanceof OptionalParameter
-      ) and
-      result = Label::parameter(param.getPosition())
+    /**
+     * Holds if there should be a `lbl`-edge from the given callable to a parameter.
+     */
+    pragma[nomagic]
+    private predicate parameterStep(string lbl, DataFlow::Node callable, DataFlowPrivate::ParameterNodeImpl paramNode) {
+      exists(DataFlowDispatch::ParameterPosition paramPos |
+        paramNode.isSourceParameterOf(callable.asExpr().getExpr(), paramPos) and
+        lbl = getLabelFromParameterPosition(paramPos)
+      )
     }
 
     /**
diff --git a/ruby/ql/test/library-tests/dataflow/api-graphs/callbacks.rb b/ruby/ql/test/library-tests/dataflow/api-graphs/callbacks.rb
@@ -32,3 +32,8 @@ def getCallback()
     }
 end
 Something.indirectCallback(getCallback()) #$ use=getMember("Something").getMethod("indirectCallback").getReturn()
+
+Something.withMixed do |a, *args, b| #$ use=getMember("Something").getMethod("withMixed").getReturn()
+    a.something #$ use=getMember("Something").getMethod("withMixed").getBlock().getParameter(0).getMethod("something").getReturn()
+    b.something #$ use=getMember("Something").getMethod("withMixed").getBlock().getParameter(1).getMethod("something").getReturn()
+end

Original file line number	Diff line number	Diff line change
`@@ -32,3 +32,8 @@ def getCallback()`
`32`	`32`	`}`
`33`	`33`	`end`
`34`	`34`	`Something.indirectCallback(getCallback()) #$ use=getMember("Something").getMethod("indirectCallback").getReturn()`
	`35`	`+`
	`36`	`+Something.withMixed do \|a, *args, b\| #$ use=getMember("Something").getMethod("withMixed").getReturn()`
	`37`	`+ a.something #$ use=getMember("Something").getMethod("withMixed").getBlock().getParameter(0).getMethod("something").getReturn()`
	`38`	`+ b.something #$ use=getMember("Something").getMethod("withMixed").getBlock().getParameter(1).getMethod("something").getReturn()`
	`39`	`+end`