Auto relaunch

cspell.json

@@ -30,6 +30,7 @@
     "globstar",
     "groupstart",
     "HighlightJS",
+    "IMAGENAME",
     "installdir",
     "Jsonifiable",
     "konami",

scripts/inject/index.mts

@@ -3,7 +3,14 @@ import "./checks/elevate.mjs";
 import "./checks/env.mjs";
 
 import { join } from "path";
-import { AnsiEscapes, getCommand } from "./util.mjs";
+import {
+  AnsiEscapes,
+  PlatformNames,
+  getCommand,
+  getProcessInfoByName,
+  killProcessByPID,
+  openProcess,
+} from "./util.mjs";
 import { inject, uninject } from "./injector.mjs";
 
 import * as darwin from "./platforms/darwin.mjs";
@@ -20,6 +27,7 @@ const platformModules = {
 
 const exitCode = process.argv.includes("--no-exit-codes") ? 0 : 1;
 const prod = process.argv.includes("--production");
+const noRelaunch = process.argv.includes("--no-relaunch");
 const processArgs = process.argv.filter((v) => !v.startsWith("-"));
 
 if (!(process.platform in platformModules)) {
@@ -94,7 +102,20 @@ const run = async (cmd = processArgs[2], replug = false): Promise<void> => {
 
   if (cmd === "inject") {
     try {
-      result = await inject(platformModule, platform, prod);
+      if (noRelaunch) {
+        result = await inject(platformModule, platform, prod);
+      } else {
+        const processName = PlatformNames[platform].replace(" ", "");
+        const processInfo = getProcessInfoByName(processName)!;
+        await killProcessByPID(processInfo?.pid);
+        result = await inject(platformModule, platform, prod);
+        const appDir = await platformModule.getAppDir(platform);
+        openProcess(
+          join(appDir, "..", "..", "..", "Update"),
+          process.platform === "win32" ? ["--processStart", `${processName}.exe`] : [],
+          { detached: true, stdio: "ignore" },
+        );
+      }
     } catch (e) {
       console.error(
         `${AnsiEscapes.RED}An error occurred while trying to inject into Discord!${AnsiEscapes.RESET}`,
@@ -121,7 +142,22 @@ To plug into a different platform, use the following syntax: ${AnsiEscapes.BOLD}
     }
   } else if (cmd === "uninject") {
     try {
-      result = await uninject(platformModule, platform);
+      if (noRelaunch) {
+        result = await uninject(platformModule, platform);
+      } else {
+        const processName = PlatformNames[platform].replace(" ", "");
+        const processInfo = getProcessInfoByName(processName)!;
+        await killProcessByPID(processInfo?.pid);
+        result = await uninject(platformModule, platform);
+        if (!replug) {
+          const appDir = await platformModule.getAppDir(platform);
+          openProcess(
+            join(appDir, "..", "..", "..", "Update"),
+            process.platform === "win32" ? ["--processStart", "Discord.exe"] : [],
+            { detached: true, stdio: "ignore" },
+          );
+        }
+      }
     } catch (e) {
       console.error(
         `${AnsiEscapes.RED}An error occurred while trying to uninject from Discord!${AnsiEscapes.RESET}`,

scripts/inject/types.mts

@@ -5,3 +5,8 @@ export type DiscordPlatform = "stable" | "ptb" | "canary" | "dev";
 export interface PlatformModule {
   getAppDir: (platform: DiscordPlatform) => Promisable<string>;
 }
+
+export interface ProcessInfo {
+  pid: number;
+  cmd: string[];
+}

scripts/inject/util.mts

@@ -1,4 +1,5 @@
-import { DiscordPlatform } from "./types.mjs";
+import { SpawnOptions, execSync, spawn } from "child_process";
+import { DiscordPlatform, ProcessInfo } from "./types.mjs";
 
 export const AnsiEscapes = {
   RESET: "\x1b[0m",
@@ -29,3 +30,56 @@ export const getCommand = ({
   cmd += ` ${platform || `[${Object.keys(PlatformNames).join("|")}]`}`;
   return cmd;
 };
+
+export const getProcessInfoByName = (processName: string): ProcessInfo | null => {
+  if (process.platform === "win32") {
+    const command = `tasklist /FI "IMAGENAME eq ${processName}.exe" /FO CSV`;
+    const output = execSync(command).toString().trim();
+
+    const lines = output.split("\r\n");
+    if (lines.length <= 2) {
+      return null;
+    }
+
+    const [_header, data] = lines.slice(0, 2);
+    const [name, pid] = data.split('","');
+
+    return { pid: Number(pid), cmd: name.substring(1).split(/\s+/) };
+  }
+  const command = `ps -eo pid,command | grep -E "(^|/)${processName}(\\s|$)" | grep -v grep`;
+  const output = execSync(command).toString().trim();
+
+  if (output.length === 0) {
+    return null;
+  }
+
+  const [pid, ...cmd] = output.split(/\s+/);
+
+  return { pid: Number(pid), cmd };
+};
+
+export const killCheckProcessExists = (pid: number): boolean => {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch {
+    return false;
+  }
+};
+
+export const killProcessByPID = (pid: number): Promise<void> => {
+  return new Promise((resolve) => {
+    if (!pid) resolve();
+    process.kill(pid, "SIGTERM");
+    const checkInterval = setInterval(() => {
+      if (!killCheckProcessExists(pid)) {
+        clearInterval(checkInterval);
+        resolve();
+      }
+    }, 1000);
+  });
+};
+
+export const openProcess = (command: string, args: string[], options: SpawnOptions): void => {
+  spawn(command, args, options).unref();
+};