Ruby: Add Improper Memoization query

This query finds cases where a method memoizes its result but fails to
include one or more of its parameters in the memoization key (or doesn't
use memoization keys at all). This can lead to the method returning
incorrect results when subsequently called with different arguments.

Author: hmac

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

@@ -0,0 +1,86 @@
+/**
+ * Provides predicates for reasoning about improper memoization methods.
+ */
+
+import ruby
+import codeql.ruby.DataFlow
+import codeql.ruby.dataflow.internal.DataFlowDispatch
+
+/**
+ * A `||=` statement that memoizes a value by storing it in an instance variable.
+ */
+private class MemoStmt extends AssignLogicalOrExpr {
+  private InstanceVariableAccess instanceVariable;
+
+  MemoStmt() { instanceVariable.getParent*() = this.getLeftOperand() }
+
+  /**
+   * Gets the variable access on the LHS of this statement.
+   * This is the `@a` in `@a ||= e`.
+   */
+  VariableAccess getVariableAccess() { result = instanceVariable }
+}
+
+/**
+ * A `||=` statement that stores a value at a particular location in a Hash,
+ * which itself is stored in an instance variable.
+ * For example:
+ * ```rb
+ * @a[key] ||= e
+ * ```
+ */
+private class HashMemoStmt extends MemoStmt {
+  HashMemoStmt() {
+    exists(ElementReference e |
+      e = this.getLeftOperand() and this.getVariableAccess().getParent+() = e
+    )
+  }
+}
+
+/**
+ * A method that may be performing memoization.
+ */
+private class MemoCandidate extends Method {
+  MemoCandidate() { exists(MemoStmt m | m.getEnclosingMethod() = this) }
+}
+
+/**
+ * Holds if parameter `p` of `m` is read in the right hand side of `assign`.
+ */
+private predicate parameterUsedInMemoValue(Method m, Parameter p, MemoStmt a) {
+  p = m.getAParameter() and
+  a.getParent+() = m and
+  p.getAVariable().getAnAccess().getParent*() = a.getRightOperand()
+}
+
+/**
+ * Holds if parameter `p` of `m` is read in the left hand side of `assign`.
+ */
+private predicate parameterUsedInMemoKey(Method m, Parameter p, HashMemoStmt a) {
+  p = m.getAParameter() and
+  a.getParent+() = m and
+  p.getAVariable().getAnAccess().getParent+() = a.getLeftOperand()
+}
+
+/**
+ * Holds if the assignment `s` is returned from its parent method `m`.
+ */
+private predicate memoReturnedFromMethod(Method m, MemoStmt s) {
+  exists(DataFlow::Node n | n.asExpr().getExpr() = s and exprNodeReturnedFrom(n, m))
+  or
+  // If we don't have flow (e.g. due to the dataflow library not supporting instance variable flow yet),
+  // fall back to a syntactic heuristic: does the last statement in the method mention the memoization variable?
+  m.getLastStmt().getAChild*().(InstanceVariableReadAccess).getVariable() =
+    s.getVariableAccess().getVariable()
+}
+
+/**
+ * Holds if `m` is a memoization method with a parameter `p` which is not used in the memoization key.
+ * This can cause stale or incorrect values to be returned when the method is called with different arguments.
+ */
+predicate isImproperMemoizationMethod(Method m, Parameter p, AssignLogicalOrExpr s) {
+  m.getName() != "initialize" and
+  parameterUsedInMemoValue(m, p, s) and
+  not parameterUsedInMemoKey(m, p, s) and
+  memoReturnedFromMethod(m, s)
+}

ruby/ql/src/experimental/improper-memoization/ImproperMemoization.qhelp

@@ -0,0 +1,84 @@
+<!DOCTYPE qhelp PUBLIC
+"-//Semmle//qhelp//EN"
+"qhelp.dtd">
+<qhelp>
+  <overview>
+    <p>
+      A common pattern in Ruby is to memoize the result of a method by storing
+      it in an instance variable. If the method has no parameters, it can simply
+      be stored directly in a variable.
+    </p>
+
+    <sample language="ruby">
+      def answer
+        @answer ||= calculate_answer
+      end
+    </sample>
+
+    <p>
+    If the method takes parameters, then there are multiple results to store
+    (one for each combination of parameter value). In this case the values
+    should be stored in a hash, keyed by the parameter values.
+    </p>
+
+    <sample language="ruby">
+      def answer(x, y)
+        @answer ||= {}
+        @answer[x] ||= {}
+        @answer[x][y] ||= calculate_answer
+      end
+    </sample>
+
+    <p>
+      If a memoization method takes parameters but does not include them in the
+      memoization key, subsequent calls to the method with different parameter
+      values may incorrectly return the same result. This can lead to the method
+      returning stale data, or leaking sensitive information.
+    </p>
+  </overview>
+
+  <example>
+    <p>
+      In this example, the method does not include its parameters in the
+      memoization key. The first call to this method will cache the result, and
+      subsequent calls will return the same result regardless of what arguments
+      are given.
+    </p>
+
+    <sample language="ruby">
+      def answer(x, y)
+        @answer ||= calculate_answer(x, y)
+      end
+    </sample>
+
+    <p>
+      This can be fixed by storing the result of <code>calculate_answer</code>
+      in a hash, keyed by the parameter values.
+    </p>
+
+    <sample language="ruby">
+      def answer(x, y)
+        @answer ||= {}
+        @answer[x] ||= {}
+        @answer[x][y] ||= calculate_answer(x, y)
+      end
+    </sample>
+
+    <p>
+      Note that if the result of <code>calculate_answer</code> is
+      <code>false</code> or <code>nil</code>, then it will not be cached.
+      To cache these values you can use a different pattern:
+    </p>
+
+    <sample language="ruby">
+      def answer(x, y)
+        @answer ||= Hash.new do |h1, x|
+          h1[x] = Hash.new do |h2, y|
+            h2[y] = calculate_answer(x, y)
+          end
+        end
+        @answer[x][y]
+      end
+    </sample>
+  </example>
+</qhelp>

ruby/ql/src/experimental/improper-memoization/ImproperMemoization.ql

@@ -0,0 +1,17 @@
+/**
+ * @name Improper Memoization
+ * @description Omitting a parameter from the key of a memoization method can lead to reading stale or incorrect data.
+ * @kind problem
+ * @problem.severity warning
+ * @precision high
+ * @tags security
+ * @id rb/improper-memoization
+ */
+
+import ruby
+import codeql.ruby.security.ImproperMemoizationQuery
+
+from Method m, Parameter p, AssignLogicalOrExpr s
+where isImproperMemoizationMethod(m, p, s)
+select m, "A $@ in this memoization method does not form part of the memoization key.", p,
+  "parameter"

ruby/ql/test/query-tests/experimental/improper-memoization/ImproperMemoization.expected

@@ -0,0 +1,12 @@
+failures
+| improper_memoization.rb:100:1:104:3 | m14 | Unexpected result: result=BAD |
+#select
+| improper_memoization.rb:50:1:55:3 | m7 | improper_memoization.rb:50:8:50:10 | arg | improper_memoization.rb:51:3:53:5 | ... \|\|= ... |
+| improper_memoization.rb:58:1:63:3 | m8 | improper_memoization.rb:58:8:58:10 | arg | improper_memoization.rb:59:3:61:5 | ... \|\|= ... |
+| improper_memoization.rb:66:1:68:3 | m9 | improper_memoization.rb:66:8:66:10 | arg | improper_memoization.rb:67:3:67:34 | ... \|\|= ... |
+| improper_memoization.rb:71:1:73:3 | m10 | improper_memoization.rb:71:9:71:12 | arg1 | improper_memoization.rb:72:3:72:42 | ... \|\|= ... |
+| improper_memoization.rb:71:1:73:3 | m10 | improper_memoization.rb:71:15:71:18 | arg2 | improper_memoization.rb:72:3:72:42 | ... \|\|= ... |
+| improper_memoization.rb:76:1:79:3 | m11 | improper_memoization.rb:76:15:76:18 | arg2 | improper_memoization.rb:78:3:78:48 | ... \|\|= ... |
+| improper_memoization.rb:82:1:87:3 | m12 | improper_memoization.rb:82:15:82:18 | arg2 | improper_memoization.rb:83:3:85:5 | ... \|\|= ... |
+| improper_memoization.rb:90:1:97:3 | m13 | improper_memoization.rb:90:9:90:10 | id | improper_memoization.rb:91:3:95:5 | ... \|\|= ... |
+| improper_memoization.rb:100:1:104:3 | m14 | improper_memoization.rb:100:9:100:11 | arg | improper_memoization.rb:103:3:103:40 | ... \|\|= ... |

ruby/ql/test/query-tests/experimental/improper-memoization/ImproperMemoization.ql

@@ -0,0 +1,23 @@
+import ruby
+import TestUtilities.InlineExpectationsTest
+import codeql.ruby.security.ImproperMemoizationQuery
+
+class ImproperMemoizationTest extends InlineExpectationsTest {
+  ImproperMemoizationTest() { this = "ImproperMemoizationTest" }
+
+  override string getARelevantTag() { result = "BAD" }
+
+  override predicate hasActualResult(Location location, string element, string tag, string value) {
+    tag = "result" and
+    value = "BAD" and
+    exists(Expr e |
+      isImproperMemoizationMethod(e, _, _) and
+      location = e.getLocation() and
+      element = e.toString()
+    )
+  }
+}
+
+from Method m, Parameter p, AssignLogicalOrExpr s
+where isImproperMemoizationMethod(m, p, s)
+select m, p, s

ruby/ql/test/query-tests/experimental/improper-memoization/improper_memoization.rb

@@ -0,0 +1,104 @@
+# GOOD - Should not trigger CodeQL rule
+
+# No arguments passed to method
+def m1
+  @m1 ||= long_running_method
+end
+
+# No arguments passed to method
+def m2
+  @m2 ||= begin
+    long_running_method
+  end
+end
+
+# OK: argument used in key.
+# May be incorrect if arg is `false` or `nil`.
+def m3(arg)
+  @m3 ||= {}
+  @m3[arg] ||= long_running_method(arg)
+end
+
+# OK: both arguments used in key.
+# May be incorrect if either arg is `false` or `nil`.
+def m4(arg1, arg2)
+  @m4 ||= {}
+  @m4[[arg1, arg2]] ||= result(arg1, arg2)
+end
+
+# OK: argument used in key.
+# Still correct if arg is `false` or `nil`.
+def m5(arg)
+  @m5 ||= Hash.new do |h1, key|
+    h1[key] = long_running_method(key)
+  end
+  @m5[arg]
+end
+
+# OK: both arguments used in key.
+# Still correct if either arg is `false` or `nil`.
+def m6(arg1, arg2)
+  @m6 ||= Hash.new do |h1, arg1|
+    h1[arg1] = Hash.new do |h2, arg2|
+      h2[arg2] = result(arg1, arg2)
+    end
+  end
+  @m6[arg1][arg2]
+end
+
+# Bad: method has parameter but only one result is memoized.
+def m7(arg) # $result=BAD
+  @m7 ||= begin
+    arg += 3
+  end
+  @m7
+end
+
+# Bad: method has parameter but only one result is memoized.
+def m8(arg) # $result=BAD
+  @m8 ||= begin
+    long_running_method(arg)
+  end
+  @m8
+end
+
+# Bad: method has parameter but only one result is memoized.
+def m9(arg) # $result=BAD
+  @m9 ||= long_running_method(arg)
+end
+
+# Bad: method has parameter but only one result is memoized.
+def m10(arg1, arg2) # $result=BAD
+  @m10 ||= long_running_method(arg1, arg2)
+end
+
+# Bad: `arg2` not used in key.
+def m11(arg1, arg2) # $result=BAD
+  @m11 ||= {}
+  @m11[arg1] ||= long_running_method(arg1, arg2)
+end
+
+# Bad: `arg2` not used in key.
+def m12(arg1, arg2) # $result=BAD
+  @m12 ||= Hash.new do |h1, arg1|
+    h1[arg1] = result(arg1, arg2)
+  end
+  @m12[arg1]
+end
+
+# Bad: arg not used in key.
+def m13(id:) # $result=BAD
+  @m13 ||= Rails.cache.fetch("product_sku/#{id}", expires_in: 30.minutes) do
+    ActiveRecord::Base.transaction do
+      ProductSku.find_by(id: id)
+    end
+  end
+  @m13
+end
+
+# Good (FP): arg is used in key via string interpolation.
+def m14(arg)
+  @m14 ||= {}
+  key = "foo/#{arg}"
+  @m14[key] ||= long_running_method(arg)
+end