From bf556a2b536f3ef7bc562f69a4adafe8a9cf8606 Mon Sep 17 00:00:00 2001
From: Arthur Baars <aibaars@github.com>
Date: Tue, 13 Apr 2021 16:48:06 +0200
Subject: [PATCH] Implement method lookup

---
 ql/src/codeql_ruby/ast/Method.qll          |  6 ++
 ql/src/codeql_ruby/ast/Module.qll          | 15 ++++
 ql/src/codeql_ruby/ast/internal/Module.qll | 84 ++++++++++++++++++++++
 3 files changed, 105 insertions(+)

diff --git a/ql/src/codeql_ruby/ast/Method.qll b/ql/src/codeql_ruby/ast/Method.qll
index 988bc8fde06..04eeaff9a1f 100644
--- a/ql/src/codeql_ruby/ast/Method.qll
+++ b/ql/src/codeql_ruby/ast/Method.qll
@@ -11,6 +11,12 @@ class Method extends TMethod {
   /** Get a declaration of this module, if any. */
   MethodBase getADeclaration() { result.getMethod() = this }
 
+  /** Get the name of this method */
+  string getName() { this = TInstanceMethod(_, result) }
+
+  /** Get the module in which this method is defined */
+  Module getModule() { this = TInstanceMethod(result, _) }
+
   /** Gets a textual representation of this method. */
   string toString() {
     exists(Module m, string name |
diff --git a/ql/src/codeql_ruby/ast/Module.qll b/ql/src/codeql_ruby/ast/Module.qll
index d449a576e4a..6e90a9f87db 100644
--- a/ql/src/codeql_ruby/ast/Module.qll
+++ b/ql/src/codeql_ruby/ast/Module.qll
@@ -11,6 +11,21 @@ class Module extends TModule {
   /** Get a declaration of this module, if any. */
   ModuleBase getADeclaration() { result.getModule() = this }
 
+  /** Get the super class of this module, if any. */
+  Module getSuperClass() { result = getSuperClass(this) }
+
+  /** Get a method defined in this module by name. */
+  Method getMethod(string name) { result.getName() = name and result.getModule() = this }
+
+  /** Look up a method in this module's ancestor chain. */
+  Method lookupMethod(string name) { result = lookupMethod(this, name) }
+
+  /** Get a `prepend`ed module. */
+  Module getAPrependedModule() { result = getAPrependedModule(this) }
+
+  /** Get an `include`d module. */
+  Module getAnIncludedModule() { result = getAnIncludedModule(this) }
+
   /** Gets a textual representation of this module. */
   string toString() {
     this = TResolved(result)
diff --git a/ql/src/codeql_ruby/ast/internal/Module.qll b/ql/src/codeql_ruby/ast/internal/Module.qll
index a1f128d0a08..01a77f583d0 100644
--- a/ql/src/codeql_ruby/ast/internal/Module.qll
+++ b/ql/src/codeql_ruby/ast/internal/Module.qll
@@ -34,6 +34,66 @@ private module Cached {
       result = scopeAppend(container, n.getName())
     )
   }
+
+  cached
+  Module getSuperClass(Module cls) {
+    cls = TResolved("Object") and result = TResolved("BasicObject")
+    or
+    cls = TResolved("Module") and result = TResolved("Object")
+    or
+    cls = TResolved("Class") and result = TResolved("Module")
+    or
+    not cls = TResolved(builtin()) and
+    (
+      exists(ClassDeclaration d |
+        d = cls.getADeclaration() and
+        result = resolveScopeExpr(d.getSuperclassExpr())
+      )
+      or
+      result = TResolved("Object") and
+      forex(ClassDeclaration d | d = cls.getADeclaration() |
+        not exists(resolveScopeExpr(d.getSuperclassExpr()))
+      )
+    )
+  }
+
+  cached
+  Module getAnIncludedModule(Module m) {
+    m = TResolved("Object") and result = TResolved("Kernel")
+    or
+    exists(IncludeOrPrependCall c |
+      c.getMethodName() = "include" and
+      (
+        m = resolveScopeExpr(c.getReceiver())
+        or
+        m = enclosingModule(c).getModule() and
+        (
+          c.getReceiver() instanceof Self
+          or
+          not exists(c.getReceiver())
+        )
+      ) and
+      result = resolveScopeExpr(c.getAnArgument())
+    )
+  }
+
+  cached
+  Module getAPrependedModule(Module m) {
+    exists(IncludeOrPrependCall c |
+      c.getMethodName() = "prepend" and
+      (
+        m = resolveScopeExpr(c.getReceiver())
+        or
+        m = enclosingModule(c).getModule() and
+        (
+          c.getReceiver() instanceof Self
+          or
+          not exists(c.getReceiver())
+        )
+      ) and
+      result = resolveScopeExpr(c.getAnArgument())
+    )
+  }
 }
 
 import Cached
@@ -242,3 +302,27 @@ private string qualifiedModuleName(string container, string name) {
     container = "Object" and name = result
   )
 }
+
+private Module getAncestors(Module m) {
+  result = m or
+  result = getAncestors(m.getAnIncludedModule()) or
+  result = getAncestors(m.getAPrependedModule())
+}
+
+private Method lookupMethod0(Module m, string name) {
+  result = lookupMethod0(m.getAPrependedModule(), name)
+  or
+  not exists(getAncestors(m.getAPrependedModule()).getMethod(name)) and
+  (
+    result = m.getMethod(name)
+    or
+    not exists(m.getMethod(name)) and result = lookupMethod0(m.getAnIncludedModule(), name)
+  )
+}
+
+Method lookupMethod(Module m, string name) {
+  result = lookupMethod0(m, name)
+  or
+  not exists(lookupMethod0(m, name)) and
+  result = lookupMethod(m.getSuperClass(), name)
+}