Merge pull request #3020 from github/koesie10/refactor-predicates

Refactor model editor predicates

Author: koesie10

extensions/ql-vscode/src/model-editor/auto-modeler.ts

@@ -18,6 +18,7 @@ import { Mode } from "./shared/mode";
 import { CancellationTokenSource } from "vscode";
 import { ModelingStore } from "./modeling-store";
 import { ModelConfigListener } from "../config";
+import { QueryLanguage } from "../common/query-language";
 
 /**
  * The auto-modeler holds state around auto-modeling jobs and allows
@@ -36,6 +37,7 @@ export class AutoModeler {
     private readonly modelingStore: ModelingStore,
     private readonly queryStorageDir: string,
     private readonly databaseItem: DatabaseItem,
+    private readonly language: QueryLanguage,
     private readonly addModeledMethods: (
       modeledMethods: Record<string, ModeledMethod[]>,
     ) => Promise<void>,
@@ -202,7 +204,7 @@ export class AutoModeler {
       filename: "auto-model.yml",
     });
 
-    const loadedMethods = loadDataExtensionYaml(models);
+    const loadedMethods = loadDataExtensionYaml(models, this.language);
     if (!loadedMethods) {
       return;
     }

extensions/ql-vscode/src/model-editor/flow-model-queries.ts

@@ -5,20 +5,22 @@ import { QueryRunner } from "../query-server";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { showAndLogExceptionWithTelemetry } from "../common/logging";
 import { extLogger } from "../common/logging/vscode";
-import { extensiblePredicateDefinitions } from "./predicates";
+import { getModelsAsDataLanguage } from "./languages";
 import { ProgressCallback } from "../common/vscode/progress";
 import { getOnDiskWorkspaceFolders } from "../common/vscode/workspace-folders";
 import { ModeledMethod, ModeledMethodType } from "./modeled-method";
 import { redactableError } from "../common/errors";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { runQuery } from "../local-queries/run-query";
 import { resolveQueries } from "../local-queries";
+import { QueryLanguage } from "../common/query-language";
 
 type FlowModelOptions = {
   cliServer: CodeQLCliServer;
   queryRunner: QueryRunner;
   queryStorageDir: string;
   databaseItem: DatabaseItem;
+  language: QueryLanguage;
   progress: ProgressCallback;
   token: CancellationToken;
   onResults: (results: ModeledMethod[]) => void | Promise<void>;
@@ -104,6 +106,7 @@ async function runSingleFlowQuery(
     queryRunner,
     queryStorageDir,
     databaseItem,
+    language,
     progress,
     token,
   }: Omit<FlowModelOptions, "onResults">,
@@ -140,7 +143,9 @@ async function runSingleFlowQuery(
   }
 
   // Interpret the results
-  const definition = extensiblePredicateDefinitions[type];
+  const modelsAsDataLanguage = getModelsAsDataLanguage(language);
+
+  const definition = modelsAsDataLanguage[type];
 
   const bqrsPath = completedQuery.outputDir.bqrsPath;
 

extensions/ql-vscode/src/model-editor/languages/index.ts

@@ -0,0 +1,2 @@
+export * from "./languages";
+export * from "./models-as-data";

extensions/ql-vscode/src/model-editor/languages/languages.ts

@@ -0,0 +1,18 @@
+import { QueryLanguage } from "../../common/query-language";
+import { ModelsAsDataLanguage } from "./models-as-data";
+import { staticLanguage } from "./static";
+
+const languages: Partial<Record<QueryLanguage, ModelsAsDataLanguage>> = {
+  [QueryLanguage.CSharp]: staticLanguage,
+  [QueryLanguage.Java]: staticLanguage,
+};
+
+export function getModelsAsDataLanguage(
+  language: QueryLanguage,
+): ModelsAsDataLanguage {
+  const definition = languages[language];
+  if (!definition) {
+    throw new Error(`No models-as-data definition for ${language}`);
+  }
+  return definition;
+}

extensions/ql-vscode/src/model-editor/languages/models-as-data.ts

@@ -0,0 +1,19 @@
+import { ModeledMethod, ModeledMethodType } from "../modeled-method";
+import { DataTuple } from "../model-extension-file";
+
+type GenerateMethodDefinition = (method: ModeledMethod) => DataTuple[];
+type ReadModeledMethod = (row: DataTuple[]) => ModeledMethod;
+
+export type ModelsAsDataLanguageModelType = Exclude<ModeledMethodType, "none">;
+
+export type ModelsAsDataLanguageModel = {
+  extensiblePredicate: string;
+  supportedKinds: string[];
+  generateMethodDefinition: GenerateMethodDefinition;
+  readModeledMethod: ReadModeledMethod;
+};
+
+export type ModelsAsDataLanguage = Record<
+  ModelsAsDataLanguageModelType,
+  ModelsAsDataLanguageModel
+>;

extensions/ql-vscode/src/model-editor/languages/shared.ts

@@ -0,0 +1,25 @@
+export const sharedExtensiblePredicates = {
+  source: "sourceModel",
+  sink: "sinkModel",
+  summary: "summaryModel",
+  neutral: "neutralModel",
+};
+
+export const sharedKinds = {
+  source: ["local", "remote"],
+  sink: [
+    "code-injection",
+    "command-injection",
+    "file-content-store",
+    "html-injection",
+    "js-injection",
+    "ldap-injection",
+    "log-injection",
+    "path-injection",
+    "request-forgery",
+    "sql-injection",
+    "url-redirection",
+  ],
+  summary: ["taint", "value"],
+  neutral: ["summary", "source", "sink"],
+};

…ql-vscode/src/model-editor/predicates.ts …ode/src/model-editor/languages/static.tsextensions/ql-vscode/src/model-editor/predicates.ts renamed to extensions/ql-vscode/src/model-editor/languages/static.ts extensions/ql-vscode/src/model-editor/predicates.ts renamed to extensions/ql-vscode/src/model-editor/languages/static.ts

@@ -1,24 +1,16 @@
-import { ModeledMethod, ModeledMethodType, Provenance } from "./modeled-method";
-import { DataTuple } from "./model-extension-file";
-
-export type ExtensiblePredicateDefinition = {
-  extensiblePredicate: string;
-  generateMethodDefinition: (method: ModeledMethod) => DataTuple[];
-  readModeledMethod: (row: DataTuple[]) => ModeledMethod;
-
-  supportedKinds?: string[];
-};
+import { ModelsAsDataLanguage } from "./models-as-data";
+import { ModeledMethodType, Provenance } from "../modeled-method";
+import { DataTuple } from "../model-extension-file";
+import { sharedExtensiblePredicates, sharedKinds } from "./shared";
 
 function readRowToMethod(row: DataTuple[]): string {
   return `${row[0]}.${row[1]}#${row[3]}${row[4]}`;
 }
 
-export const extensiblePredicateDefinitions: Record<
-  Exclude<ModeledMethodType, "none">,
-  ExtensiblePredicateDefinition
-> = {
+export const staticLanguage: ModelsAsDataLanguage = {
   source: {
-    extensiblePredicate: "sourceModel",
+    extensiblePredicate: sharedExtensiblePredicates.source,
+    supportedKinds: sharedKinds.source,
     // extensible predicate sourceModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string output, string kind, string provenance
@@ -35,7 +27,7 @@ export const extensiblePredicateDefinitions: Record<
       method.provenance,
     ],
     readModeledMethod: (row) => ({
-      type: "source",
+      type: "source" as ModeledMethodType,
       input: "",
       output: row[6] as string,
       kind: row[7] as string,
@@ -46,10 +38,10 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: ["local", "remote"],
   },
   sink: {
-    extensiblePredicate: "sinkModel",
+    extensiblePredicate: sharedExtensiblePredicates.sink,
+    supportedKinds: sharedKinds.sink,
     // extensible predicate sinkModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string input, string kind, string provenance
@@ -77,22 +69,10 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: [
-      "code-injection",
-      "command-injection",
-      "file-content-store",
-      "html-injection",
-      "js-injection",
-      "ldap-injection",
-      "log-injection",
-      "path-injection",
-      "request-forgery",
-      "sql-injection",
-      "url-redirection",
-    ],
   },
   summary: {
-    extensiblePredicate: "summaryModel",
+    extensiblePredicate: sharedExtensiblePredicates.summary,
+    supportedKinds: sharedKinds.summary,
     // extensible predicate summaryModel(
     //   string package, string type, boolean subtypes, string name, string signature, string ext,
     //   string input, string output, string kind, string provenance
@@ -121,10 +101,10 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[3] as string,
       methodParameters: row[4] as string,
     }),
-    supportedKinds: ["taint", "value"],
   },
   neutral: {
-    extensiblePredicate: "neutralModel",
+    extensiblePredicate: sharedExtensiblePredicates.neutral,
+    supportedKinds: sharedKinds.neutral,
     // extensible predicate neutralModel(
     //   string package, string type, string name, string signature, string kind, string provenance
     // );
@@ -148,6 +128,5 @@ export const extensiblePredicateDefinitions: Record<
       methodName: row[2] as string,
       methodParameters: row[3] as string,
     }),
-    supportedKinds: ["summary", "source", "sink"],
   },
 };

extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts

@@ -15,6 +15,10 @@ import { ModelEditorViewTracker } from "../model-editor-view-tracker";
 import { ModelConfigListener } from "../../config";
 import { DatabaseItem } from "../../databases/local-databases";
 import { ModelingEvents } from "../modeling-events";
+import {
+  QueryLanguage,
+  tryGetQueryLanguage,
+} from "../../common/query-language";
 
 export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ToMethodModelingMessage,
@@ -24,6 +28,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
 
   private method: Method | undefined = undefined;
   private databaseItem: DatabaseItem | undefined = undefined;
+  private language: QueryLanguage | undefined = undefined;
 
   constructor(
     app: App,
@@ -45,6 +50,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
     await this.postMessage({
       t: "setMethodModelingPanelViewState",
       viewState: {
+        language: this.language,
         showMultipleModels: this.modelConfig.showMultipleModels,
       },
     });
@@ -56,6 +62,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
   ): Promise<void> {
     this.method = method;
     this.databaseItem = databaseItem;
+    this.language = databaseItem && tryGetQueryLanguage(databaseItem.language);
 
     if (this.isShowingView) {
       await this.postMessage({
@@ -70,6 +77,9 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
       const selectedMethod = this.modelingStore.getSelectedMethodDetails();
       if (selectedMethod) {
         this.databaseItem = selectedMethod.databaseItem;
+        this.language = tryGetQueryLanguage(
+          selectedMethod.databaseItem.language,
+        );
         this.method = selectedMethod.method;
 
         await this.postMessage({
@@ -185,6 +195,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
         if (this.webviewView) {
           this.method = e.method;
           this.databaseItem = e.databaseItem;
+          this.language = tryGetQueryLanguage(e.databaseItem.language);
 
           await this.postMessage({
             t: "setSelectedMethod",

extensions/ql-vscode/src/model-editor/model-editor-module.ts

@@ -233,6 +233,7 @@ export class ModelEditorModule extends DisposableObject {
             queryDir,
             db,
             modelFile,
+            language,
           );
 
           this.modelingEvents.onDbClosed(async (dbUri) => {

extensions/ql-vscode/src/model-editor/model-editor-view.ts

@@ -38,7 +38,10 @@ import { ModelConfigListener } from "../config";
 import { INITIAL_MODE, Mode } from "./shared/mode";
 import { loadModeledMethods, saveModeledMethods } from "./modeled-method-fs";
 import { pickExtensionPack } from "./extension-pack-picker";
-import { getLanguageDisplayName } from "../common/query-language";
+import {
+  getLanguageDisplayName,
+  QueryLanguage,
+} from "../common/query-language";
 import { AutoModeler } from "./auto-modeler";
 import { telemetryListener } from "../common/vscode/telemetry";
 import { ModelingStore } from "./modeling-store";
@@ -64,6 +67,8 @@ export class ModelEditorView extends AbstractWebview<
     private readonly queryDir: string,
     private readonly databaseItem: DatabaseItem,
     private readonly extensionPack: ExtensionPack,
+    // The language is equal to databaseItem.language but is properly typed as QueryLanguage
+    private readonly language: QueryLanguage,
     initialMode: Mode = INITIAL_MODE,
   ) {
     super(app);
@@ -82,6 +87,7 @@ export class ModelEditorView extends AbstractWebview<
       modelingStore,
       queryStorageDir,
       databaseItem,
+      language,
       async (modeledMethods) => {
         this.addModeledMethods(modeledMethods);
       },
@@ -218,7 +224,7 @@ export class ModelEditorView extends AbstractWebview<
               });
               await saveModeledMethods(
                 this.extensionPack,
-                this.databaseItem.language,
+                this.language,
                 methods,
                 modeledMethods,
                 mode,
@@ -367,6 +373,7 @@ export class ModelEditorView extends AbstractWebview<
       t: "setModelEditorViewState",
       viewState: {
         extensionPack: this.extensionPack,
+        language: this.language,
         showFlowGeneration: this.modelConfig.flowGeneration,
         showLlmButton,
         showMultipleModels: this.modelConfig.showMultipleModels,
@@ -394,6 +401,7 @@ export class ModelEditorView extends AbstractWebview<
     try {
       const modeledMethods = await loadModeledMethods(
         this.extensionPack,
+        this.language,
         this.cliServer,
         this.app.logger,
       );
@@ -458,6 +466,14 @@ export class ModelEditorView extends AbstractWebview<
           if (!addedDatabase) {
             return;
           }
+
+          if (addedDatabase.language !== this.language) {
+            void showAndLogErrorMessage(
+              this.app.logger,
+              `The selected database is for ${addedDatabase.language}, but the current database is for ${this.language}.`,
+            );
+            return;
+          }
         }
 
         progress({
@@ -472,6 +488,7 @@ export class ModelEditorView extends AbstractWebview<
             queryRunner: this.queryRunner,
             queryStorageDir: this.queryStorageDir,
             databaseItem: addedDatabase ?? this.databaseItem,
+            language: this.language,
             onResults: async (modeledMethods) => {
               const modeledMethodsByName: Record<string, ModeledMethod[]> = {};
 
@@ -579,6 +596,7 @@ export class ModelEditorView extends AbstractWebview<
         this.queryDir,
         addedDatabase,
         modelFile,
+        this.language,
         Mode.Framework,
       );
       await view.openView();