Merge pull request #9267 from hmac/hmac/improper-memoization

Ruby: Add Improper Memoization query

Author: hmac

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

@@ -0,0 +1,87 @@
+/**
+ * Provides predicates for reasoning about improper memoization methods.
+ */
+
+private import ruby
+private import codeql.ruby.DataFlow
+private 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() { this = any(MemoStmt m).getEnclosingMethod() }
+}
+
+/**
+ * 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.getEnclosingMethod() = 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.getEnclosingMethod() = 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 instanceof MemoCandidate and
+  m.getName() != "initialize" and
+  parameterUsedInMemoValue(m, p, s) and
+  not parameterUsedInMemoKey(m, p, s) and
+  memoReturnedFromMethod(m, s)
+}

ruby/ql/src/change-notes/2022-05-24-improper-memoization.md

@@ -0,0 +1,4 @@
+---
+category: newQuery
+---
+* Added a new query, `rb/improper-memoization`. The query finds cases where the parameter of a memoization method is not used in the memoization key.

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