support the NewEntity module in QL-for-QL

2023-01-30 10:06:08 +01:00 · 2023-01-30 10:06:08 +01:00 · 54c4c23b46
--- a/ql/ql/src/codeql_ql/ast/internal/Builtins.qll
+++ b/ql/ql/src/codeql_ql/ast/internal/Builtins.qll
@ -115,6 +115,12 @@ module QlBuiltinsMocks {
      or
      i = 2 and
      result instanceof EquivalenceRelation::EquivalenceRelationModule
+      or
+      i = 3 and
+      result instanceof NewEntity::EntityKeySigClass
+      or
+      i = 4 and
+      result instanceof NewEntity::NewEntityModule
    }
  }

@ -123,15 +129,16 @@ module QlBuiltinsMocks {
   * The equivalent to the following is implemented:
   * ```CodeQL
   * module QlBuiltins {
-   *  signature class T;
-   *  module EdgeSig<T MyT> { // This might not be needed.
-   *    signature predicate edgeSig(MyT a, MyT b);
-   *  }
-   *  module EquivalenceRelation<T MyT, EdgeSig<MyT>::edgeSig/2 edge> { // the `edge` parameter is not modeled
-   *    class EquivalenceClass;
-   *    EquivalenceClass getEquivalenceClass(MyT a);
-   *  }
-   *}
+   *   signature class T;
+   *   module EdgeSig<T MyT> { // This might not be needed.
+   *     signature predicate edgeSig(MyT a, MyT b);
+   *   }
+   *   module EquivalenceRelation<T MyT, EdgeSig<MyT>::edgeSig/2 edge> { // the `edge` parameter is not modeled
+   *     class EquivalenceClass;
+   *     EquivalenceClass getEquivalenceClass(MyT a);
+   *   }
+   * }
+   * ```
   */
  module EquivalenceRelation {
    class SigClass extends MockClass::Range {
@ -259,4 +266,92 @@ module QlBuiltinsMocks {
      override string getClassName() { result = "EquivalenceClass" }
    }
  }
+
+  /**
+   * A mock that implements the `NewEntity` module.
+   * The equivalent to the following is implemented:
+   * ```CodeQL
+   * class EntityKeySig;
+   * module NewEntity<EntityKeySig EntityKey> {
+   *   class EntityId;
+   *
+   *   EntityId map(EntityKey key) { none() }
+   * }
+   * ```
+   */
+  module NewEntity {
+    class EntityKeySigClass extends MockClass::Range {
+      EntityKeySigClass() { this = "Mock: QlBuiltins::NewEntity::EntityKeySig" }
+
+      override string getName() { result = "EntityKeySig" }
+    }
+
+    class NewEntityModule extends MockModule::Range {
+      NewEntityModule() { this = "Mock: QlBuiltins::NewEntity" }
+
+      override string getName() { result = "NewEntity" }
+
+      override string getMember(int i) {
+        i = 0 and result instanceof EntityIdClass
+        or
+        i = 1 and result instanceof NewEntityMapPredicate
+      }
+
+      ///  Holds if the `i`th type parameter has `type` (the ID of the mocked node) with `name`.
+      override predicate hasTypeParam(int i, string type, string name) {
+        i = 0 and
+        name = "EntityKey" and
+        type instanceof EntityKeySigTypeExpr
+      }
+    }
+
+    class EntityKeySigTypeExpr extends MockTypeExpr::Range {
+      EntityKeySigTypeExpr() { this = "Mock: QlBuiltins::NewEntity::EntityKey" }
+
+      override string getClassName() { result = "EntityKeySig" }
+    }
+
+    class EntityIdClass extends MockClass::Range {
+      EntityIdClass() { this = "Mock: QlBuiltins::NewEntity::EntityId" }
+
+      override string getName() { result = "EntityId" }
+    }
+
+    class NewEntityMapPredicate extends MockClasslessPredicate::Range {
+      NewEntityMapPredicate() { this = "Mock: QlBuiltins::NewEntity::map" }
+
+      override string getName() { result = "map" }
+
+      override string getParameter(int i) {
+        i = 0 and
+        result instanceof NewEntityMapPredicateParam
+      }
+
+      override MockTypeExpr::Range getReturnTypeExpr() {
+        result.(NewEntityMapPredicateTypes).getClassName() = "EntityId"
+      }
+    }
+
+    // both the TypeExprs used in the `map` predicate.
+    class NewEntityMapPredicateTypes extends MockTypeExpr::Range {
+      string type;
+
+      NewEntityMapPredicateTypes() {
+        type = ["EntityId", "EntityKey"] and
+        this = "Mock: QlBuiltins::NewEntity::map::T#" + type
+      }
+
+      override string getClassName() { result = type }
+    }
+
+    class NewEntityMapPredicateParam extends MockVarDecl::Range {
+      NewEntityMapPredicateParam() { this = "Mock: QlBuiltins::NewEntity::map::#0" }
+
+      override string getName() { result = "key" }
+
+      override MockTypeExpr::Range getType() {
+        result.(NewEntityMapPredicateTypes).getClassName() = "EntityKey"
+      }
+    }
+  }
 }
--- a/ql/ql/test/modules/Foo.qll
+++ b/ql/ql/test/modules/Foo.qll
@ -34,3 +34,27 @@ private class TypeFlowScc = Scc::EquivalenceClass;
 predicate sccRepr(Node n, TypeFlowScc scc) { scc = Scc::getEquivalenceClass(n) }

 predicate sccJoinStep(Node n, TypeFlowScc scc) { none() }
+
+module NewEntity {
+  newtype TFoo = TFoo1()
+
+  newtype EntityKey =
+    Key1() or
+    Key2()
+
+  // this errors out in normal QL, but QL-for-QL doesn't differentiate between upgrade scripts and "normal" code, and it also doesn't care if the number of type-parameters matches.
+  // so this should work fine in QL-for-QL
+  module NewEntityModule = QlBuiltins::NewEntity<EntityKey>;
+
+  class Union = TFoo or NewEntityModule::EntityId;
+
+  class Foo extends Union {
+    string toString() { none() }
+  }
+
+  predicate foo(Foo id, string message) {
+    id = NewEntityModule::map(Key1()) and message = "upgrade-1"
+    or
+    id = NewEntityModule::map(Key2()) and message = "upgrade-2"
+  }
+}
--- a/ql/ql/test/modules/test.expected
+++ b/ql/ql/test/modules/test.expected
@ -14,6 +14,14 @@ getTarget
 | Foo.qll:31:29:31:31 | Scc | file://:0:0:0:0 | EquivalenceRelation |
 | Foo.qll:34:52:34:54 | Scc | Foo.qll:29:16:29:18 | Scc |
 | Foo.qll:34:52:34:54 | Scc | file://:0:0:0:0 | EquivalenceRelation |
+| Foo.qll:47:28:47:37 | QlBuiltins | file://:0:0:0:0 | QlBuiltins |
+| Foo.qll:47:28:47:59 | NewEntity | file://:0:0:0:0 | NewEntity |
+| Foo.qll:49:25:49:39 | NewEntityModule | Foo.qll:47:10:47:24 | NewEntityModule |
+| Foo.qll:49:25:49:39 | NewEntityModule | file://:0:0:0:0 | NewEntity |
+| Foo.qll:56:10:56:24 | NewEntityModule | Foo.qll:47:10:47:24 | NewEntityModule |
+| Foo.qll:56:10:56:24 | NewEntityModule | file://:0:0:0:0 | NewEntity |
+| Foo.qll:58:10:58:24 | NewEntityModule | Foo.qll:47:10:47:24 | NewEntityModule |
+| Foo.qll:58:10:58:24 | NewEntityModule | file://:0:0:0:0 | NewEntity |
 getTargetType
 | ClassSig.qll:3:23:3:28 | TypeExpr | file://:0:0:0:0 | string |
 | ClassSig.qll:7:12:7:17 | TypeExpr | ClassSig.qll:1:17:1:22 | FooSig |
@ -44,6 +52,25 @@ getTargetType
 | Foo.qll:34:52:34:54 | Scc | file://:0:0:0:0 | EquivalenceRelation |
 | Foo.qll:36:23:36:26 | TypeExpr | Foo.qll:23:7:23:10 | Node |
 | Foo.qll:36:31:36:41 | TypeExpr | file://:0:0:0:0 | EquivalenceClass |
+| Foo.qll:47:28:47:37 | QlBuiltins | file://:0:0:0:0 | QlBuiltins |
+| Foo.qll:47:28:47:59 | NewEntity | file://:0:0:0:0 | NewEntity |
+| Foo.qll:47:50:47:58 | TypeExpr | Foo.qll:41:11:41:19 | EntityKey |
+| Foo.qll:49:17:49:20 | TypeExpr | Foo.qll:39:11:39:14 | TFoo |
+| Foo.qll:49:25:49:39 | NewEntityModule | Foo.qll:47:10:47:24 | NewEntityModule |
+| Foo.qll:49:25:49:39 | NewEntityModule | file://:0:0:0:0 | NewEntity |
+| Foo.qll:49:25:49:49 | TypeExpr | file://:0:0:0:0 | EntityId |
+| Foo.qll:51:21:51:25 | TypeExpr | Foo.qll:49:9:49:13 | Union |
+| Foo.qll:52:5:52:10 | TypeExpr | file://:0:0:0:0 | string |
+| Foo.qll:55:17:55:19 | TypeExpr | Foo.qll:51:9:51:11 | Foo |
+| Foo.qll:55:25:55:30 | TypeExpr | file://:0:0:0:0 | string |
+| Foo.qll:56:10:56:24 | NewEntityModule | Foo.qll:47:10:47:24 | NewEntityModule |
+| Foo.qll:56:10:56:24 | NewEntityModule | file://:0:0:0:0 | NewEntity |
+| Foo.qll:58:10:58:24 | NewEntityModule | Foo.qll:47:10:47:24 | NewEntityModule |
+| Foo.qll:58:10:58:24 | NewEntityModule | file://:0:0:0:0 | NewEntity |
+| file://:0:0:0:0 | TypeExpr | Foo.qll:41:11:41:19 | EntityKey |
+| file://:0:0:0:0 | TypeExpr | file://:0:0:0:0 | EntityId |
+| file://:0:0:0:0 | TypeExpr | file://:0:0:0:0 | EntityKeySig |
+| file://:0:0:0:0 | TypeExpr | file://:0:0:0:0 | EntityKeySig |
 | file://:0:0:0:0 | TypeExpr | file://:0:0:0:0 | EquivalenceClass |
 | file://:0:0:0:0 | TypeExpr | file://:0:0:0:0 | T |
 | file://:0:0:0:0 | TypeExpr | file://:0:0:0:0 | T |