Use stdin for supplying auth to CodeQL

koesie10 · koesie10 · commit 20a06238a29d · 2023-01-13T15:02:51.000+02:00
This will supply the GitHub access token to certain CodeQL CLI commands
such that private packages can be resolved. It will only do so if the
user has an existing auth session. If they don't, they will now get a
prompt to login. However, this will only happen for commands which
actually use authentication, which is limited to packaging commands.
diff --git a/extensions/ql-vscode/src/authentication.ts b/extensions/ql-vscode/src/authentication.ts
@@ -4,10 +4,11 @@ import { retry } from "@octokit/plugin-retry";
 
 const GITHUB_AUTH_PROVIDER_ID = "github";
 
-// We need 'repo' scope for triggering workflows and 'gist' scope for exporting results to Gist.
+// We need 'repo' scope for triggering workflows, 'gist' scope for exporting results to Gist,
+// and 'read:packages' for reading private CodeQL packages.
 // For a comprehensive list of scopes, see:
 // https://docs.github.com/apps/building-oauth-apps/understanding-scopes-for-oauth-apps
-const SCOPES = ["repo", "gist"];
+const SCOPES = ["repo", "gist", "read:packages"];
 
 /**
  * Handles authentication to GitHub, using the VS Code [authentication API](https://code.visualstudio.com/api/references/vscode-api#authentication).
@@ -57,15 +58,31 @@ export class Credentials {
       return this.octokit;
     }
 
+    const accessToken = await this.getAccessToken();
+
+    return new Octokit.Octokit({
+      auth: accessToken,
+      retry,
+    });
+  }
+
+  async getAccessToken(): Promise<string> {
     const session = await vscode.authentication.getSession(
       GITHUB_AUTH_PROVIDER_ID,
       SCOPES,
       { createIfNone: true },
     );
 
-    return new Octokit.Octokit({
-      auth: session.accessToken,
-      retry,
-    });
+    return session.accessToken;
+  }
+
+  async getExistingAccessToken(): Promise<string | undefined> {
+    const session = await vscode.authentication.getSession(
+      GITHUB_AUTH_PROVIDER_ID,
+      SCOPES,
+      { createIfNone: false },
+    );
+
+    return session?.accessToken;
   }
 }
diff --git a/extensions/ql-vscode/src/cli.ts b/extensions/ql-vscode/src/cli.ts
@@ -1,3 +1,4 @@
+import { EOL } from "os";
 import { spawn } from "child-process-promise";
 import * as child_process from "child_process";
 import { readFile } from "fs-extra";
@@ -26,6 +27,7 @@ import { Logger, ProgressReporter } from "./common";
 import { CompilationMessage } from "./pure/legacy-messages";
 import { sarifParser } from "./sarif-parser";
 import { dbSchemeToLanguage, walkDirectory } from "./helpers";
+import { Credentials } from "./authentication";
 
 /**
  * The version of the SARIF format that we are using.
@@ -156,6 +158,10 @@ interface BqrsDecodeOptions {
   entities?: string[];
 }
 
+export type OnLineCallback = (
+  line: string,
+) => Promise<string | undefined> | string | undefined;
+
 /**
  * This class manages a cli server started by `codeql execute cli-server` to
  * run commands without the overhead of starting a new java
@@ -304,6 +310,7 @@ export class CodeQLCliServer implements Disposable {
     command: string[],
     commandArgs: string[],
     description: string,
+    onLine?: OnLineCallback,
   ): Promise<string> {
     const stderrBuffers: Buffer[] = [];
     if (this.commandInProcess) {
@@ -328,6 +335,22 @@ export class CodeQLCliServer implements Disposable {
         await new Promise<void>((resolve, reject) => {
           // Start listening to stdout
           process.stdout.addListener("data", (newData: Buffer) => {
+            if (onLine) {
+              void (async () => {
+                const response = await onLine(newData.toString("utf-8"));
+
+                if (!response) {
+                  return;
+                }
+
+                process.stdin.write(`${response}${EOL}`);
+
+                // Remove newData from stdoutBuffers because the data has been consumed
+                // by the onLine callback.
+                stdoutBuffers.splice(stdoutBuffers.indexOf(newData), 1);
+              })();
+            }
+
             stdoutBuffers.push(newData);
             // If the buffer ends in '0' then exit.
             // We don't have to check the middle as no output will be written after the null until
@@ -487,13 +510,15 @@ export class CodeQLCliServer implements Disposable {
    * @param commandArgs The arguments to pass to the `codeql` command.
    * @param description Description of the action being run, to be shown in log and error messages.
    * @param progressReporter Used to output progress messages, e.g. to the status bar.
+   * @param onLine Used for responding to interactive output on stdout/stdin.
    * @returns The contents of the command's stdout, if the command succeeded.
    */
   runCodeQlCliCommand(
     command: string[],
     commandArgs: string[],
     description: string,
     progressReporter?: ProgressReporter,
+    onLine?: OnLineCallback,
   ): Promise<string> {
     if (progressReporter) {
       progressReporter.report({ message: description });
@@ -503,10 +528,12 @@ export class CodeQLCliServer implements Disposable {
       // Construct the command that actually does the work
       const callback = (): void => {
         try {
-          this.runCodeQlCliInternal(command, commandArgs, description).then(
-            resolve,
-            reject,
-          );
+          this.runCodeQlCliInternal(
+            command,
+            commandArgs,
+            description,
+            onLine,
+          ).then(resolve, reject);
         } catch (err) {
           reject(err);
         }
@@ -528,6 +555,7 @@ export class CodeQLCliServer implements Disposable {
    * @param description Description of the action being run, to be shown in log and error messages.
    * @param addFormat Whether or not to add commandline arguments to specify the format as JSON.
    * @param progressReporter Used to output progress messages, e.g. to the status bar.
+   * @param onLine Used for responding to interactive output on stdout/stdin.
    * @returns The contents of the command's stdout, if the command succeeded.
    */
   async runJsonCodeQlCliCommand<OutputType>(
@@ -536,6 +564,7 @@ export class CodeQLCliServer implements Disposable {
     description: string,
     addFormat = true,
     progressReporter?: ProgressReporter,
+    onLine?: OnLineCallback,
   ): Promise<OutputType> {
     let args: string[] = [];
     if (addFormat)
@@ -547,6 +576,7 @@ export class CodeQLCliServer implements Disposable {
       args,
       description,
       progressReporter,
+      onLine,
     );
     try {
       return JSON.parse(result) as OutputType;
@@ -559,6 +589,44 @@ export class CodeQLCliServer implements Disposable {
     }
   }
 
+  /**
+   * Runs a CodeQL CLI command, returning the output as JSON.
+   * @param command The `codeql` command to be run, provided as an array of command/subcommand names.
+   * @param commandArgs The arguments to pass to the `codeql` command.
+   * @param description Description of the action being run, to be shown in log and error messages.
+   * @param addFormat Whether or not to add commandline arguments to specify the format as JSON.
+   * @param progressReporter Used to output progress messages, e.g. to the status bar.
+   * @returns The contents of the command's stdout, if the command succeeded.
+   */
+  async runJsonCodeQlCliCommandWithAuthentication<OutputType>(
+    command: string[],
+    commandArgs: string[],
+    description: string,
+    addFormat = true,
+    progressReporter?: ProgressReporter,
+  ): Promise<OutputType> {
+    const credentials = await Credentials.initialize();
+
+    const accessToken = await credentials.getExistingAccessToken();
+
+    const extraArgs = accessToken ? ["--github-auth-stdin"] : [];
+
+    return this.runJsonCodeQlCliCommand(
+      command,
+      [...extraArgs, ...commandArgs],
+      description,
+      addFormat,
+      progressReporter,
+      async (line) => {
+        if (line.startsWith("Enter value for --github-auth-stdin")) {
+          return credentials.getAccessToken();
+        }
+
+        return undefined;
+      },
+    );
+  }
+
   /**
    * Resolve the library path and dbscheme for a query.
    * @param workspaces The current open workspaces
@@ -1136,7 +1204,7 @@ export class CodeQLCliServer implements Disposable {
    * @param packs The `<package-scope/name[@version]>` of the packs to download.
    */
   async packDownload(packs: string[]) {
-    return this.runJsonCodeQlCliCommand(
+    return this.runJsonCodeQlCliCommandWithAuthentication(
       ["pack", "download"],
       packs,
       "Downloading packs",
@@ -1148,7 +1216,7 @@ export class CodeQLCliServer implements Disposable {
     if (forceUpdate) {
       args.push("--mode", "update");
     }
-    return this.runJsonCodeQlCliCommand(
+    return this.runJsonCodeQlCliCommandWithAuthentication(
       ["pack", "install"],
       args,
       "Installing pack dependencies",
@@ -1169,7 +1237,7 @@ export class CodeQLCliServer implements Disposable {
       ...this.getAdditionalPacksArg(workspaceFolders),
     ];
 
-    return this.runJsonCodeQlCliCommand(
+    return this.runJsonCodeQlCliCommandWithAuthentication(
       ["pack", "bundle"],
       args,
       "Bundling pack",
@@ -1200,7 +1268,7 @@ export class CodeQLCliServer implements Disposable {
   ): Promise<{ [pack: string]: string }> {
     // Uses the default `--mode use-lock`, which creates the lock file if it doesn't exist.
     const results: { [pack: string]: string } =
-      await this.runJsonCodeQlCliCommand(
+      await this.runJsonCodeQlCliCommandWithAuthentication(
         ["pack", "resolve-dependencies"],
         [dir],
         "Resolving pack dependencies",
diff --git a/extensions/ql-vscode/test/vscode-tests/cli-integration/run-cli.test.ts b/extensions/ql-vscode/test/vscode-tests/cli-integration/run-cli.test.ts
@@ -1,4 +1,4 @@
-import { extensions, Uri } from "vscode";
+import { authentication, extensions, Uri } from "vscode";
 import { join } from "path";
 import { SemVer } from "semver";
 
@@ -12,6 +12,7 @@ import {
 } from "../../../src/helpers";
 import { resolveQueries } from "../../../src/contextual/queryResolver";
 import { KeyType } from "../../../src/contextual/keyType";
+import { faker } from "@faker-js/faker";
 
 jest.setTimeout(60_000);
 
@@ -104,4 +105,59 @@ describe("Use cli", () => {
       }
     },
   );
+
+  describe("github authentication", () => {
+    itWithCodeQL()(
+      "should not use authentication if there are no credentials",
+      async () => {
+        const getSession = jest
+          .spyOn(authentication, "getSession")
+          .mockResolvedValue(undefined);
+
+        await cli.packDownload(["codeql/tutorial"]);
+        expect(getSession).toHaveBeenCalledTimes(1);
+        expect(getSession).toHaveBeenCalledWith(
+          "github",
+          expect.arrayContaining(["read:packages"]),
+          {
+            createIfNone: false,
+          },
+        );
+      },
+    );
+
+    itWithCodeQL()(
+      "should use authentication if there are credentials",
+      async () => {
+        const getSession = jest
+          .spyOn(authentication, "getSession")
+          .mockResolvedValue({
+            id: faker.datatype.uuid(),
+            accessToken: "gho_xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx",
+            account: {
+              id: faker.datatype.uuid(),
+              label: "Account",
+            },
+            scopes: ["read:packages"],
+          });
+
+        await cli.packDownload(["codeql/tutorial"]);
+        expect(getSession).toHaveBeenCalledTimes(2);
+        expect(getSession).toHaveBeenCalledWith(
+          "github",
+          expect.arrayContaining(["read:packages"]),
+          {
+            createIfNone: false,
+          },
+        );
+        expect(getSession).toHaveBeenCalledWith(
+          "github",
+          expect.arrayContaining(["read:packages"]),
+          {
+            createIfNone: true,
+          },
+        );
+      },
+    );
+  });
 });
