vscode-codeql/extensions/ql-vscode/src/codeql-cli/cli.ts

import { EOL } from "os";
import { spawn } from "child-process-promise";
import type { ChildProcessWithoutNullStreams } from "child_process";
import { spawn as spawnChildProcess } from "child_process";
import { readFile } from "fs-extra";
import { delimiter, dirname, join } from "path";
import type { Log } from "sarif";
import { SemVer } from "semver";
import type { Readable } from "stream";
import tk from "tree-kill";
import type { CancellationToken, Disposable, Uri } from "vscode";

import type {
  BqrsInfo,
  DecodedBqrs,
  DecodedBqrsChunk,
} from "../common/bqrs-cli-types";
import type { CliConfig } from "../config";
import type { DistributionProvider } from "./distribution";
import { FindDistributionResultKind } from "./distribution";
import {
  assertNever,
  getErrorMessage,
  getErrorStack,
} from "../common/helpers-pure";
import { walkDirectory } from "../common/files";
import type { QueryMetadata } from "../common/interface-types";
import { SortDirection } from "../common/interface-types";
import type { BaseLogger, Logger } from "../common/logging";
import type { ProgressReporter } from "../common/logging/vscode";
import { sarifParser } from "../common/sarif-parser";
import type { App } from "../common/app";
import { QueryLanguage } from "../common/query-language";
import { LINE_ENDINGS, splitStreamAtSeparators } from "../common/split-stream";
import type { Position } from "../query-server/messages";
import { LOGGING_FLAGS } from "./cli-command";
import type { CliFeatures, VersionAndFeatures } from "./cli-version";
import { ExitCodeError, getCliError } from "./cli-errors";
import { UserCancellationException } from "../common/vscode/progress";

const SARIF_FORMAT = "sarifv2.1.0";
const CSV_FORMAT = "csv";

export interface DbInfo {
  sourceLocationPrefix: string;
  columnKind: string;
  unicodeNewlines: boolean;
  sourceArchiveZip: string;
  sourceArchiveRoot: string;
  datasetFolder: string;
  logsFolder: string;
  languages: string[];
}

export type QueryInfoByLanguage = Record<string, never>; // If it's always empty

export type QlpacksInfo = Record<string, string[]>; // `codeql resolve qlpacks`
type LanguagesInfo = Record<string, string[]>; // `codeql resolve languages`

type MlModelInfo = { checksum: string; path: string };
type MlModelsInfo = { models: MlModelInfo[] }; // `codeql resolve ml-models`

type DataExtensionResult = { predicate: string; file: string; index: number };
type ResolveExtensionsResult = {
  models: MlModelInfo[];
  data: Record<string, DataExtensionResult[]>; // `codeql resolve extensions`
};

type GenerateExtensiblePredicateMetadataResult = {
  extensible_predicates: Array<{ path: string }>; // Pack-relative path
};

type PackDownloadResult = {
  packs: Array<{ name: string; version: string }>;
  packDir: string;
};

type QlrefInfo = { resolvedPath: string }; // `codeql resolve qlref`

export interface SourceInfo {
  sourceArchive: string;
  sourceLocationPrefix: string;
}

type ResolvedQueries = string[]; // `codeql resolve queries`
type ResolvedTests = string[]; // `codeql resolve tests`

export enum CompilationMessageSeverity {
  Error = "ERROR",
  Warning = "WARNING",
}

export interface CompilationMessage {
  message: string; // Text of the message
  position: Position; // Source position
  severity: CompilationMessageSeverity; // Error or Warning
}

export interface TestCompleted {
  // Event fired by `codeql test run`
  test: string;
  pass: boolean;
  messages: CompilationMessage[];
  compilationMs: number;
  evaluationMs: number;
  expected: string;
  actual?: string;
  diff?: string[];
  failureDescription?: string;
  failureStage?: string;
}

interface BqrsDecodeOptions {
  // Optional arguments for `bqrsDecode`
  pageSize?: number; // Number of results to retrieve
  offset?: number; // 0-based index of the first result
  entities?: string[]; // Default: ["url", "string"]
}

type OnLineCallback = (
  line: string,
) => Promise<string | undefined> | string | undefined;
type VersionChangedListener = (
  newVersionAndFeatures?: VersionAndFeatures,
) => void;

type RunOptions = {
  progressReporter?: ProgressReporter; // Outputs progress messages (e.g., status bar)
  onLine?: OnLineCallback; // Handles interactive output on stdout/stdin
  silent?: boolean; // Suppresses logs in the CodeQL extension log
  runInNewProcess?: boolean; // Runs command in a new process instead of CLI server
  token?: CancellationToken; // Enables cancellation (only if `runInNewProcess` is true)
};

type JsonRunOptions = RunOptions & {
  addFormat?: boolean; // Adds command-line arguments to specify JSON format
};

export class CodeQLCliServer implements Disposable {
  process?: ChildProcessWithoutNullStreams;
  commandQueue: Array<() => void> = [];
  commandInProcess = false;
  nullBuffer = Buffer.alloc(1);
  private _versionAndFeatures?: VersionAndFeatures;
  private _versionChangedListeners: VersionChangedListener[] = [];
  private _supportedLanguages?: string[];
  codeQlPath?: string;
  cliConstraints = new CliVersionConstraint(this);
  public quiet = false;

  constructor(
    private readonly app: App,
    private distributionProvider: DistributionProvider,
    private cliConfig: CliConfig,
    public readonly logger: Logger,
  ) {
    this.setupListeners();
  }

  async bqrsDecode(
    bqrsPath: string,
    resultSet: string,
    { pageSize, offset, entities = ["url", "string"] }: BqrsDecodeOptions = {},
  ): Promise<DecodedBqrsChunk> {
    const args = [
      `--entities=${entities.join(",")}`,
      "--result-set",
      resultSet,
      ...(pageSize ? ["--rows", pageSize.toString()] : []),
      ...(offset ? ["--start-at", offset.toString()] : []),
      bqrsPath,
    ];

    return this.runJsonCodeQlCliCommand<DecodedBqrsChunk>(
      ["bqrs", "decode"],
      args,
      "Reading bqrs data",
    );
  }

  async bqrsInfo(bqrsPath: string, pageSize?: number): Promise<BqrsInfo> {
    const args = pageSize
      ? ["--paginate-rows", pageSize.toString(), bqrsPath]
      : [bqrsPath];
    return this.runJsonCodeQlCliCommand<BqrsInfo>(
      ["bqrs", "info"],
      args,
      "Reading bqrs header",
    );
  }

  private setupListeners() {
    this.distributionProvider.onDidChangeDistribution?.(() =>
      this.restartCliServer(),
    );
    this.cliConfig.onDidChangeConfiguration?.(() => this.restartCliServer());
  }

  dispose() {
    this.killProcessIfRunning();
  }

  private killProcessIfRunning() {
    if (!this.process) return;

    void this.logger.log("Sending shutdown request");
    try {
      this.process.stdin.write(JSON.stringify(["shutdown"]), "utf8");
      this.process.stdin.write(this.nullBuffer);
    } catch (e) {
      void this.logger.log(`Shutdown failed: ${e}`);
    }

    this.process.stdin.end();
    this.process.kill();
    this.process.stdout.destroy();
    this.process.stderr.destroy();
    this.process = undefined;
  }

  private restartCliServer() {
    const callback = () => {
      this.killProcessIfRunning();
      this._versionAndFeatures = undefined;
      this._supportedLanguages = undefined;
      this.runNext();
    };

    this.commandInProcess ? this.commandQueue.unshift(callback) : callback();
  }

  private async getCodeQlPath(): Promise<string> {
    const codeqlPath =
      await this.distributionProvider.getCodeQlPathWithoutVersionCheck();
    if (!codeqlPath) throw new Error("Failed to find CodeQL distribution.");
    return codeqlPath;
  }

  private async launchProcess(): Promise<ChildProcessWithoutNullStreams> {
    const codeQlPath = await this.getCodeQlPath();
    const args = shouldDebugCliServer()
      ? [
          "-J=-agentlib:jdwp=transport=dt_socket,address=localhost:9012,server=n,suspend=y,quiet=y",
        ]
      : [];
    return spawnServer(
      codeQlPath,
      "CodeQL CLI Server",
      ["execute", "cli-server"],
      args,
      this.logger,
      () => {},
    );
  }

  private async runCodeQlCliInternal(
    command: string[],
    commandArgs: string[],
    description: string,
    onLine?: OnLineCallback,
    silent = false,
  ): Promise<string> {
    if (this.commandInProcess) throw new Error("CLI is already running");
    this.commandInProcess = true;

    try {
      if (!this.process) this.process = await this.launchProcess();
      const process = this.process;

      const args = [...command, ...LOGGING_FLAGS, ...commandArgs];
      if (!silent) void this.logger.log(`${description}: ${args.join(" ")}...`);

      return await this.handleProcessOutput(process, {
        args,
        description,
        onLine,
        silent,
        handleNullTerminator: true,
      });
    } catch (err) {
      this.killProcessIfRunning();
      throw err;
    } finally {
      this.commandInProcess = false;
      this.runNext();
    }
  }

  private async runCodeQlCliInNewProcess(
    command: string[],
    commandArgs: string[],
    description: string,
    onLine?: OnLineCallback,
    silent = false,
    token?: CancellationToken,
  ): Promise<string> {
    const codeqlPath = await this.getCodeQlPath();
    const args = [...command, ...LOGGING_FLAGS, ...commandArgs];
    if (!silent) void this.logger.log(`${description}: ${args.join(" ")}...`);

    const abortController = new AbortController();
    const process = spawnChildProcess(codeqlPath, args, {
      signal: abortController.signal,
    });

    if (!process?.pid)
      throw new Error(
        `Failed to start ${description}: ${codeqlPath} ${args.join(" ")}`,
      );

    let exited = false;
    process.on("exit", () => (exited = true));

    const cancellationRegistration = token?.onCancellationRequested(() => {
      abortController.abort("Token was cancelled.");
      if (process.pid && !exited) tk(process.pid);
    });

    try {
      return await this.handleProcessOutput(process, {
        args,
        description,
        onLine,
        silent,
        handleNullTerminator: false,
      });
    } catch (e) {
      if (token?.isCancellationRequested) {
        void this.logger.log(`Process cancelled: ${getErrorMessage(e)}`);
        throw new UserCancellationException(`${args.join(" ")} was cancelled.`);
      }
      throw e;
    } finally {
      process.stdin.end();
      if (!exited) tk(process.pid);
      process.stdout.destroy();
      process.stderr.destroy();
      cancellationRegistration?.dispose();
    }
  }

  async resolveDatabase(databasePath: string): Promise<DbInfo> {
    return this.runJsonCodeQlCliCommand(
      ["resolve", "database"],
      [databasePath],
      "Resolving database",
    );
  }

  async resolveUpgrades(
    dbScheme: string,
    searchPath: string[],
    allowDowngradesIfPossible: boolean,
    targetDbScheme?: string,
  ) {
    const args = [
      ...this.getAdditionalPacksArg(searchPath),
      "--dbscheme",
      dbScheme,
    ];

    if (targetDbScheme) {
      args.push("--target-dbscheme", targetDbScheme);
      if (allowDowngradesIfPossible) args.push("--allow-downgrades");
    }

    return this.runJsonCodeQlCliCommand(
      ["resolve", "upgrades"],
      args,
      "Resolving database upgrade scripts",
    );
  }

  async runJsonCodeQlCliCommand<OutputType>(
    command: string[],
    commandArgs: string[],
    description: string,
    { addFormat = true, ...runOptions }: JsonRunOptions = {},
  ): Promise<OutputType> {
    const args = addFormat ? ["--format", "json", ...commandArgs] : commandArgs;
    const result = await this.runCodeQlCliCommand(
      command,
      args,
      description,
      runOptions,
    );

    try {
      return JSON.parse(result) as OutputType;
    } catch (err) {
      throw new Error(
        `Parsing output of ${description} failed: ${getErrorMessage(err)}`,
      );
    }
  }

  private getAdditionalPacksArg(paths: string[]): string[] {
    return paths.length ? ["--additional-packs", paths.join(delimiter)] : [];
  }
  private async handleProcessOutput(
    process: ChildProcessWithoutNullStreams,
    {
      handleNullTerminator,
      args,
      description,
      onLine,
      silent,
    }: {
      handleNullTerminator: boolean;
      args: string[];
      description: string;
      onLine?: OnLineCallback;
      silent?: boolean;
    },
  ): Promise<string> {
    const stderrBuffers: Buffer[] = [];
    let stdoutBuffers: Buffer[] = [];

    return new Promise<void>((resolve, reject) => {
      const stdoutListener = (newData: Buffer) => {
        stdoutBuffers.push(newData);
        if (handleNullTerminator && newData.at(-1) === 0) resolve();
      };

      const stderrListener = (newData: Buffer) => {
        stderrBuffers.push(newData);
        if (!silent) this.logger.log(newData.toString("utf-8"));
      };

      const closeListener = (code: number | null) => {
        code === 0 ? resolve() : reject(new ExitCodeError(code));
      };

      process.stdout.on("data", stdoutListener);
      process.stderr.on("data", stderrListener);
      process.on("close", closeListener);
      process.on("error", reject);
    })
      .then(() => {
        const output = Buffer.concat(stdoutBuffers).toString("utf8").trim();
        if (!silent) this.logger.log("CLI command succeeded.");
        return output;
      })
      .catch((err) => {
        throw getCliError(
          err,
          stderrBuffers.length
            ? Buffer.concat(stderrBuffers).toString("utf8")
            : undefined,
          description,
          args,
        );
      });
  }

  private runNext(): void {
    this.commandQueue.shift()?.();
  }

  async resolve(databasePath: string): Promise<DbInfo> {
    return this.runJsonCodeQlCliCommand(
      ["resolve", "database"],
      [databasePath],
      "Resolving database",
    );
  }

  runCodeQlCliCommand(
    command: string[],
    commandArgs: string[],
    description: string,
    options: RunOptions = {},
  ): Promise<string> {
    const {
      progressReporter,
      onLine,
      silent = false,
      runInNewProcess = false,
      token,
    } = options;

    progressReporter?.report({ message: description });

    return runInNewProcess
      ? this.runCodeQlCliInNewProcess(
          command,
          commandArgs,
          description,
          onLine,
          silent,
          token,
        )
      : new Promise((resolve, reject) => {
          const callback = () => {
            this.runCodeQlCliInternal(
              command,
              commandArgs,
              description,
              onLine,
              silent,
            )
              .then(resolve)
              .catch(reject);
          };
          this.commandInProcess ? this.commandQueue.push(callback) : callback();
        });
  }

  async resolveQlpacks(
    additionalPacks: string[],
    extensionPacksOnly = false,
    kind?: "query" | "library" | "all",
  ): Promise<QlpacksInfo> {
    const args = this.getAdditionalPacksArg(additionalPacks);
    if (extensionPacksOnly) args.push("--kind", "extension", "--no-recursive");
    else if (kind) args.push("--kind", kind);

    return this.runJsonCodeQlCliCommand(
      ["resolve", "qlpacks"],
      args,
      "Resolving qlpack information",
    );
  }
  public async getVersion(): Promise<SemVer> {
    return (await this.getVersionAndFeatures()).version;
  }

  public async getFeatures(): Promise<CliFeatures> {
    return (await this.getVersionAndFeatures()).features;
  }

  private async refreshVersion(): Promise<VersionAndFeatures> {
    const distribution = await this.distributionProvider.getDistribution();

    if (
      distribution.kind === FindDistributionResultKind.CompatibleDistribution ||
      distribution.kind === FindDistributionResultKind.IncompatibleDistribution
    ) {
      return distribution.versionAndFeatures;
    }

    throw new Error("No distribution found");
  }

  private async getVersionAndFeatures(): Promise<VersionAndFeatures> {
    if (!this._versionAndFeatures) {
      try {
        const newVersionAndFeatures = await this.refreshVersion();
        this._versionAndFeatures = newVersionAndFeatures;
        this._versionChangedListeners.forEach((listener) =>
          listener(newVersionAndFeatures),
        );

        await this.app.commands.execute(
          "setContext",
          "codeql.supportsTrimCache",
          newVersionAndFeatures.version.compare(
            CliVersionConstraint.CLI_VERSION_WITH_TRIM_CACHE,
          ) >= 0,
        );
      } catch (e) {
        this._versionChangedListeners.forEach((listener) =>
          listener(undefined),
        );
        throw e;
      }
    }
    return this._versionAndFeatures;
  }
}

export function spawnServer(
  codeqlPath: string,
  name: string,
  command: string[],
  commandArgs: string[],
  logger: Logger,
  stderrListener: (data: string | Buffer) => void,
  stdoutListener?: (data: string | Buffer) => void,
  progressReporter?: ProgressReporter,
): ChildProcessWithoutNullStreams {
  const args = [...command, ...commandArgs, ...LOGGING_FLAGS];
  progressReporter?.report({ message: `Starting ${name}` });

  logger.log(
    `Starting ${name} using CodeQL CLI: ${codeqlPath} ${args.join(" ")}`,
  );
  const child = spawnChildProcess(codeqlPath, args);

  if (!child?.pid) {
    throw new Error(
      `Failed to start ${name} using command ${codeqlPath} ${args.join(" ")}`,
    );
  }

  let lastStdout: string | Buffer | undefined;
  child.stdout!.on("data", (data) => (lastStdout = data));

  child.on("close", (code, signal) => {
    if (code !== null) logger.log(`Child process exited with code ${code}`);
    if (signal) logger.log(`Child process exited due to signal ${signal}`);
    if (code !== 0 && lastStdout)
      logger.log(`Last stdout was "${lastStdout.toString()}"`);
  });

  child.stderr!.on("data", stderrListener);
  stdoutListener && child.stdout!.on("data", stdoutListener);

  progressReporter?.report({ message: `Started ${name}` });
  logger.log(`${name} started on PID: ${child.pid}`);

  return child;
}

/**
 * Logs a text stream to a `Logger` interface.
 */
async function logStream(stream: Readable, logger: BaseLogger): Promise<void> {
  for await (const line of splitStreamAtSeparators(stream, LINE_ENDINGS)) {
    await logger.log(line); // Ensures correct log order.
  }
}

/**
 * Checks if an environment variable is set to a truthy value.
 */
function isEnvTrue(name: string): boolean {
  const value = process.env[name]?.toLowerCase();
  return value !== undefined && value !== "0" && value !== "false";
}

/**
 * Determines whether to debug the language server.
 */
export function shouldDebugLanguageServer(): boolean {
  return isEnvTrue("IDE_SERVER_JAVA_DEBUG");
}

/**
 * Determines whether to debug the query server.
 */
export function shouldDebugQueryServer(): boolean {
  return isEnvTrue("QUERY_SERVER_JAVA_DEBUG");
}

/**
 * Determines whether to debug the CLI server.
 */
export function shouldDebugCliServer(): boolean {
  return isEnvTrue("CLI_SERVER_JAVA_DEBUG");
}
export class CliVersionConstraint {
  public static readonly OLDEST_SUPPORTED_CLI_VERSION = new SemVer("2.14.6");
  public static readonly CLI_VERSION_WITH_TRIM_CACHE = new SemVer("2.15.1");
  public static readonly CLI_VERSION_WITHOUT_MRVA_EXTENSIBLE_PREDICATE_HACK =
    new SemVer("2.16.1");
  public static readonly CLI_VERSION_WITH_MULTI_QUERY_PACK_CREATE = new SemVer(
    "2.16.1",
  );

  constructor(private readonly cli: CodeQLCliServer) {}

  private async isVersionAtLeast(v: SemVer): Promise<boolean> {
    return (await this.cli.getVersion()).compare(v) >= 0;
  }

  async preservesExtensiblePredicatesInMrvaPack(): Promise<boolean> {
    return !(await this.isVersionAtLeast(
      CliVersionConstraint.CLI_VERSION_WITHOUT_MRVA_EXTENSIBLE_PREDICATE_HACK,
    ));
  }

  async supportsPackCreateWithMultipleQueries(): Promise<boolean> {
    return this.isVersionAtLeast(
      CliVersionConstraint.CLI_VERSION_WITH_MULTI_QUERY_PACK_CREATE,
    );
  }

  async supportsMrvaPackCreate(): Promise<boolean> {
    return (await this.cli.getFeatures()).mrvaPackCreate === true;
  }
}