Enhance Grape framework with additional data flow modeling and helper method support

ruby/ql/lib/codeql/ruby/frameworks/Grape.qll

@@ -12,6 +12,7 @@ private import codeql.ruby.typetracking.TypeTracking
 private import codeql.ruby.frameworks.Rails
 private import codeql.ruby.frameworks.internal.Rails
 private import codeql.ruby.dataflow.internal.DataFlowDispatch
+private import codeql.ruby.dataflow.FlowSteps
 
 /**
  * Provides modeling for Grape, a REST-like API framework for Ruby.
@@ -125,21 +126,17 @@ class GrapeParamsSource extends Http::Server::RequestInputAccess::Range {
 }
 
 /**
- * A call to `params` from within a Grape API endpoint.
+ * A call to `params` from within a Grape API endpoint or helper method.
  */
 private class GrapeParamsCall extends ParamsCallImpl {
   GrapeParamsCall() {
-    exists(GrapeEndpoint endpoint |
-      this.getParent+() = endpoint.getBody().asCallableAstNode() and
-      this.getMethodName() = "params"
+    // Simplified approach: find params calls that are descendants of Grape API class methods
+    exists(GrapeAPIClass api |
+      this.getMethodName() = "params" and
+      this.getParent+() = api.getADeclaration()
     )
-    or
-    // Also handle cases where params is called on an instance of a Grape API class
-    this = grapeAPIInstance().getAMethodCall("params").asExpr().getExpr()
   }
-}
-
-/**
+}/**
  * A call to `headers` from within a Grape API endpoint.
  * Headers can also be a source of user input.
  */
@@ -195,4 +192,44 @@ private class GrapeRequestCall extends MethodCall {
     // Also handle cases where request is called on an instance of a Grape API class
     this = grapeAPIInstance().getAMethodCall("request").asExpr().getExpr()
   }
+}
+
+/**
+ * A method defined within a `helpers` block in a Grape API class.
+ * These methods become available in endpoint contexts through Grape's DSL.
+ */
+private class GrapeHelperMethod extends Method {
+  private GrapeAPIClass apiClass;
+
+  GrapeHelperMethod() {
+    exists(DataFlow::CallNode helpersCall |
+      helpersCall = apiClass.getAModuleLevelCall("helpers") and
+      this.getParent+() = helpersCall.getBlock().asExpr().getExpr()
+    )
+  }
+
+  /**
+   * Gets the API class that contains this helper method.
+   */
+  GrapeAPIClass getAPIClass() { result = apiClass }
+}
+
+/**
+ * Additional taint step to model dataflow from method arguments to parameters
+ * for Grape helper methods defined in `helpers` blocks.
+ * This bridges the gap where standard dataflow doesn't recognize the Grape DSL semantics.
+ */
+private class GrapeHelperMethodTaintStep extends AdditionalTaintStep {
+  override predicate step(DataFlow::Node nodeFrom, DataFlow::Node nodeTo) {
+    exists(GrapeHelperMethod helperMethod, MethodCall call, int i |
+      // Find calls to helper methods from within Grape endpoints
+      call.getMethodName() = helperMethod.getName() and
+      exists(GrapeEndpoint endpoint |
+        call.getParent+() = endpoint.getBody().asExpr().getExpr()
+      ) and
+      // Map argument to parameter
+      nodeFrom.asExpr().getExpr() = call.getArgument(i) and
+      nodeTo.asParameter() = helperMethod.getParameter(i)
+    )
+  }
 }