Merge pull request #1964 from github/koesie10/packages-auth

Use stdin for supplying auth to CodeQL

Author: koesie10

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;
   }
 }

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);
         }
@@ -522,12 +549,13 @@ export class CodeQLCliServer implements Disposable {
   }
 
   /**
-   * Runs a CodeQL CLI command, returning the output as JSON.
+   * Runs a CodeQL CLI command, parsing 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.
+   * @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,67 @@ export class CodeQLCliServer implements Disposable {
     }
   }
 
+  /**
+   * Runs a CodeQL CLI command with authentication, parsing the output as JSON.
+   *
+   * This method is intended for use with commands that accept a `--github-auth-stdin` argument. This
+   * will be added to the command line arguments automatically if an access token is available.
+   *
+   * When the argument is given to the command, the CLI server will prompt for the access token on
+   * stdin. This method will automatically respond to the prompt with the access token.
+   *
+   * There are a few race conditions that can potentially happen:
+   * 1. The user logs in after the command has started. In this case, no access token will be given.
+   * 2. The user logs out after the command has started. In this case, the user will be prompted
+   *   to login again. If they cancel the login, the old access token that was present before the
+   *   command was started will be used.
+   *
+   * @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")) {
+          try {
+            return await credentials.getAccessToken();
+          } catch (e) {
+            // If the user cancels the authentication prompt, we still need to give a value to the CLI.
+            // By giving a potentially invalid value, the user will just get a 401/403 when they try to access a
+            // private package and the access token is invalid.
+            // This code path is very rare to hit. It would only be hit if the user is logged in when
+            // starting the command, then logging out before the getAccessToken() is called again and
+            // then cancelling the authentication prompt.
+            return accessToken;
+          }
+        }
+
+        return undefined;
+      },
+    );
+  }
+
   /**
    * Resolve the library path and dbscheme for a query.
    * @param workspaces The current open workspaces
@@ -1136,7 +1227,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 +1239,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 +1260,7 @@ export class CodeQLCliServer implements Disposable {
       ...this.getAdditionalPacksArg(workspaceFolders),
     ];
 
-    return this.runJsonCodeQlCliCommand(
+    return this.runJsonCodeQlCliCommandWithAuthentication(
       ["pack", "bundle"],
       args,
       "Bundling pack",
@@ -1200,7 +1291,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",

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,
+          },
+        );
+      },
+    );
+  });
 });