make seperate steps for YAML.parse* and use getAsuccessor*() to reach final to_ruby method call, All parts have Rewritten with API graphs exclusively

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

@@ -10,25 +10,48 @@ private import codeql.ruby.ApiGraphs
  * A taint step related to the result of `YAML.parse` calls, or similar.
  * In the following example, this step will propagate taint from
  * `source` to `sink`:
- *
+ * this contains two seperate steps:
  * ```rb
  * x = source
- * result = YAML.parse(x)
- * sink result.to_ruby # Unsafe call
+ * sink = YAML.parse(x)
+ * ```
+ * By second step
+ * source is a Successor of `YAML.parse(x)`
+ * which ends with `to_ruby` or an Element of `to_ruby`
+ * ```ruby
+ * sink source.to_ruby # Unsafe call
  * ```
  */
 private class YamlParseStep extends AdditionalTaintStep {
   override predicate step(DataFlow::Node pred, DataFlow::Node succ) {
-    exists(DataFlow::CallNode yamlParserMethod |
-      succ = yamlParserMethod.getAMethodCall("to_ruby") and
+    exists(API::Node yamlParserMethod |
+      succ = yamlParserMethod.getReturn().asSource() and
       (
-        yamlParserMethod = yamlNode().getAMethodCall(["parse", "parse_stream"]) and
-        pred = [yamlParserMethod.getArgument(0), yamlParserMethod.getKeywordArgument("yaml")]
+        yamlParserMethod = yamlNode().getMethod(["parse", "parse_stream"]) and
+        pred =
+          [yamlParserMethod.getParameter(0), yamlParserMethod.getKeywordParameter("yaml")].asSink()
         or
-        yamlParserMethod = yamlNode().getAMethodCall("parse_file") and
-        pred = [yamlParserMethod.getArgument(0), yamlParserMethod.getKeywordArgument("filename")]
+        yamlParserMethod = yamlNode().getMethod("parse_file") and
+        pred =
+          [yamlParserMethod.getParameter(0), yamlParserMethod.getKeywordParameter("filename")]
+              .asSink()
       )
     )
+    or
+    exists(API::Node yamlParserMethod |
+      succ =
+        [
+          yamlParserMethod.getASuccessor*().getMethod("to_ruby").getReturn().asSource(),
+          yamlParserMethod
+              .getASuccessor*()
+              .getMethod("to_ruby")
+              .getReturn()
+              .getAnElement()
+              .asSource()
+        ] and
+      yamlParserMethod = yamlNode().getMethod(["parse", "parse_stream", "parse_file"]) and
+      pred = yamlParserMethod.getReturn().asSource()
+    )
   }
 }
 

ruby/ql/lib/codeql/ruby/security/UnsafeDeserializationCustomizations.qll

@@ -97,12 +97,26 @@ module UnsafeDeserialization {
 
   /**
    * An argument in a call to `YAML.parse*`, considered a sink for unsafe deserialization
-   * if there is a call to `to_ruby` on the returned value.
+   * if there is a call to `to_ruby` on the returned value of any Successor.
    */
   class YamlParseArgument extends Sink {
     YamlParseArgument() {
       this =
-        yamlNode().getAMethodCall(["parse", "parse_stream", "parse_file"]).getAMethodCall("to_ruby")
+        [
+          yamlNode()
+              .getMethod(["parse", "parse_stream", "parse_file"])
+              .getASuccessor*()
+              .getMethod("to_ruby")
+              .getReturn()
+              .asSource(),
+          yamlNode()
+              .getMethod(["parse", "parse_stream", "parse_file"])
+              .getASuccessor*()
+              .getMethod("to_ruby")
+              .getReturn()
+              .getAnElement()
+              .asSource()
+        ]
     }
   }
 

ruby/ql/test/query-tests/security/cwe-502/unsafe-deserialization/UnsafeDeserialization.qlref

@@ -1 +1 @@
-queries/security/cwe-502/UnsafeDeserialization.ql
+queries/security/cwe-502/UnsafeDeserialization.ql

ruby/ql/test/query-tests/security/cwe-502/unsafe-deserialization/YAMLUnsafeDeserialization.rb

@@ -15,7 +15,23 @@ class UsersController < ActionController::Base
     parse_output.to_ruby
     Psych.parse(params[:yaml_string]).to_ruby
     Psych.parse_file(params[:yaml_file]).to_ruby
-
+    parsed_yaml.children.each do |child|
+      puts child.to_ruby
+    end
+    Psych.parse_stream(params[:yaml_string])  do |document|
+      puts document.to_ruby
+    end
+    parsed_yaml.children.first.to_ruby
+    parsed_yaml = Psych.parse_stream(params[:yaml_string])
+    content = parsed_yaml.children[0].children[0].children
+    parsed = parsed_yaml.to_ruby[0]
+    parsed = content.to_ruby[0]
+    Psych.parse(params[:yaml_string]).children[0].to_ruby
+    # FP
+    parsed_yaml = Psych2.parse_stream(params[:yaml_string])
+    content = parsed_yaml.children[0].children[0].children
+    parsed = parsed_yaml.to_ruby
+    parsed = parsed_yaml.to_ruby[0]
   end
 end