Auto relaunch

scripts/inject/index.mts

@@ -3,15 +3,15 @@ import "./checks/elevate.mjs";
 import "./checks/env.mjs";
 
 import { join } from "path";
+import { smartInject } from "./injector.mjs";
 import { AnsiEscapes, getCommand } from "./util.mjs";
-import { inject, uninject } from "./injector.mjs";
 
+import { createContext, getPositionalArg } from "@marshift/argus";
+import { existsSync } from "fs";
 import * as darwin from "./platforms/darwin.mjs";
 import * as linux from "./platforms/linux.mjs";
 import * as win32 from "./platforms/win32.mjs";
-import { DiscordPlatform } from "./types.mjs";
-import { existsSync } from "fs";
-import { createContext, getPositionalArg } from "@marshift/argus";
+import type { DiscordPlatform } from "./types.mjs";
 
 const platformModules = {
   darwin,
@@ -23,6 +23,7 @@ const ctx = createContext(process.argv);
 
 export const exitCode = ctx.hasOptionalArg(/--no-exit-codes/) ? 0 : 1;
 const prod = ctx.hasOptionalArg(/--production/);
+const noRelaunch = ctx.hasOptionalArg(/--no-relaunch/);
 export const entryPoint = ctx.getOptionalArg(/--entryPoint/);
 
 if (!(process.platform in platformModules)) {
@@ -97,7 +98,7 @@ const run = async (cmd = ctx.getPositionalArg(2), replug = false): Promise<void>
 
   if (cmd === "inject") {
     try {
-      result = await inject(platformModule, platform, prod);
+      result = await smartInject(cmd, replug, platformModule, platform, prod, noRelaunch);
     } catch (e) {
       console.error(
         `${AnsiEscapes.RED}An error occurred while trying to inject into Discord!${AnsiEscapes.RESET}`,
@@ -114,7 +115,11 @@ const run = async (cmd = ctx.getPositionalArg(2), replug = false): Promise<void>
         "\n",
       );
       console.log(
-        `You now have to completely close the Discord client, from the system tray or through the task manager.\n
+        `${
+          noRelaunch
+            ? `You now have to completely close the Discord client, from the system tray or through the task manager.\n`
+            : "Your Discord client has been restarted automatically.\n"
+        }
 To plug into a different platform, use the following syntax: ${AnsiEscapes.BOLD}${
           AnsiEscapes.GREEN
         }${getCommand({ action: replug ? "replug" : "plug", prod })}${AnsiEscapes.RESET}`,
@@ -124,7 +129,7 @@ To plug into a different platform, use the following syntax: ${AnsiEscapes.BOLD}
     }
   } else if (cmd === "uninject") {
     try {
-      result = await uninject(platformModule, platform);
+      result = await smartInject(cmd, replug, platformModule, platform, prod, noRelaunch);
     } catch (e) {
       console.error(
         `${AnsiEscapes.RED}An error occurred while trying to uninject from Discord!${AnsiEscapes.RESET}`,
@@ -142,7 +147,11 @@ To plug into a different platform, use the following syntax: ${AnsiEscapes.BOLD}
           "\n",
         );
         console.log(
-          `You now have to completely close the Discord client, from the system tray or through the task manager.\n
+          `${
+            noRelaunch
+              ? `You now have to completely close the Discord client, from the system tray or through the task manager.\n`
+              : "Your Discord client has been restarted automatically.\n"
+          }
 To unplug from a different platform, use the following syntax: ${AnsiEscapes.BOLD}${
             AnsiEscapes.GREEN
           }${getCommand({ action: "unplug", prod })}${AnsiEscapes.RESET}`,

scripts/inject/injector.mts

@@ -1,14 +1,23 @@
+import { execSync } from "child_process";
+import { existsSync } from "fs";
 import { chown, copyFile, mkdir, rename, rm, stat, writeFile } from "fs/promises";
 import path, { join, sep } from "path";
 import { fileURLToPath } from "url";
-import { entryPoint as argEntryPoint, exitCode } from "./index.mjs";
-import { AnsiEscapes, getCommand } from "./util.mjs";
-import { execSync } from "child_process";
-import { DiscordPlatform, PlatformModule } from "./types.mjs";
 import { CONFIG_PATH } from "../../src/util.mjs";
-import { existsSync } from "fs";
+import { entryPoint as argEntryPoint, exitCode } from "./index.mjs";
+import type { DiscordPlatform, PlatformModule, ProcessInfo } from "./types.mjs";
+import {
+  AnsiEscapes,
+  PlatformNames,
+  getCommand,
+  getProcessInfoByName,
+  getUserData,
+  killProcessByPID,
+  openProcess,
+} from "./util.mjs";
 
 const dirname = path.dirname(fileURLToPath(import.meta.url));
+let processInfo: ProcessInfo | ProcessInfo[] | null;
 
 export const isDiscordInstalled = async (appDir: string, silent?: boolean): Promise<boolean> => {
   try {
@@ -112,6 +121,7 @@ export const inject = async (
   const entryPoint =
     argEntryPoint ??
     (prod ? join(CONFIG_PATH, "replugged.asar") : join(dirname, "..", "..", "dist/main.js"));
+
   const entryPointDir = path.dirname(entryPoint);
 
   if (appDir.includes("flatpak")) {
@@ -191,14 +201,82 @@ export const uninject = async (
     return false;
   }
 
-  await rm(appDir, { recursive: true, force: true });
-  await rename(join(appDir, "..", "app.orig.asar"), appDir);
-  // For discord_arch_electron
-  if (existsSync(join(appDir, "..", "app.orig.asar.unpacked"))) {
-    await rename(
-      join(appDir, "..", "app.orig.asar.unpacked"),
-      join(appDir, "..", "app.asar.unpacked"),
+  try {
+    await rm(appDir, { recursive: true, force: true });
+    await rename(join(appDir, "..", "app.orig.asar"), appDir);
+    // For discord_arch_electron
+    if (existsSync(join(appDir, "..", "app.orig.asar.unpacked"))) {
+      await rename(
+        join(appDir, "..", "app.orig.asar.unpacked"),
+        join(appDir, "..", "app.asar.unpacked"),
+      );
+    }
+  } catch {
+    console.error(
+      `${AnsiEscapes.RED}Failed to rename app.asar while unplugging. If Discord is open, make sure it is closed.${AnsiEscapes.RESET}`,
     );
+    process.exit(exitCode);
   }
+
   return true;
 };
+
+export const smartInject = async (
+  cmd: "uninject" | "inject",
+  replug: boolean,
+  platformModule: PlatformModule,
+  platform: DiscordPlatform,
+  production: boolean,
+  noRelaunch: boolean,
+): Promise<boolean> => {
+  let result;
+
+  const processName =
+    process.platform === "darwin"
+      ? PlatformNames[platform]
+      : PlatformNames[platform].replace(" ", "");
+  if (!noRelaunch) {
+    try {
+      if ((replug && cmd === "uninject") || !replug) {
+        processInfo = getProcessInfoByName(processName)!;
+        if (Array.isArray(processInfo)) {
+          await Promise.all(processInfo.map((info) => killProcessByPID(info.pid)));
+        } else {
+          await killProcessByPID(processInfo?.pid);
+        }
+      }
+    } catch {}
+  }
+
+  result =
+    cmd === "uninject"
+      ? await uninject(platformModule, platform)
+      : await inject(platformModule, platform, production);
+
+  if (!noRelaunch) {
+    if (((replug && cmd !== "uninject") || !replug) && processInfo) {
+      const appDir = await platformModule.getAppDir(platform);
+      switch (process.platform) {
+        case "win32":
+          openProcess(
+            join(appDir, "..", "..", "..", "Update"),
+            ["--processStart", `${processName}.exe`],
+            { detached: true, stdio: "ignore" },
+          );
+          break;
+        case "linux":
+          openProcess(join(appDir, "..", "..", processName), [], {
+            ...getUserData(),
+            detached: true,
+            stdio: "ignore",
+          });
+          break;
+        case "darwin":
+          openProcess(`open -a "${processName}.app"`);
+          break;
+      }
+    }
+  }
+
+  return result;
+};

scripts/inject/types.mts

@@ -5,3 +5,14 @@ export type DiscordPlatform = "stable" | "ptb" | "canary" | "dev";
 export interface PlatformModule {
   getAppDir: (platform: DiscordPlatform) => Promisable<string>;
 }
+
+export interface ProcessInfo {
+  pid: number;
+  ppid: number;
+}
+
+export interface UserData {
+  env: NodeJS.ProcessEnv;
+  uid: number;
+  gid: number;
+}

scripts/inject/util.mts

@@ -1,4 +1,5 @@
-import { DiscordPlatform } from "./types.mjs";
+import { type SpawnOptions, execSync, spawn } from "child_process";
+import type { DiscordPlatform, ProcessInfo, UserData } from "./types.mjs";
 
 export const AnsiEscapes = {
   RESET: "\x1b[0m",
@@ -29,3 +30,79 @@ export const getCommand = ({
   cmd += ` ${platform || `[${Object.keys(PlatformNames).join("|")}]`}`;
   return cmd;
 };
+
+export const getProcessInfoByName = (processName: string): ProcessInfo | ProcessInfo[] | null => {
+  try {
+    const isWindows = process.platform === "win32";
+    const command = isWindows
+      ? `wmic process where (Name="${processName}.exe") get ProcessId,ParentProcessId /FORMAT:CSV`
+      : `ps -eo cmd,ppid,pid | grep -E "(^|/)${processName}(\\s|$)" | grep -v grep`;
+    const output = execSync(command).toString();
+
+    if (!output.trim()) {
+      return null;
+    }
+
+    const lines = output
+      .trim()
+      .split(isWindows ? "\r\r\n" : "\n")
+      .slice(1);
+    const processInfo = lines.map((line) => {
+      const parts = isWindows ? line.split(",") : line.trim().split(/\s+/);
+      return {
+        ppid: parseInt(parts[1], 10),
+        pid: parseInt(parts[2], 10),
+      };
+    });
+
+    if (isWindows) {
+      const parentPIDs = processInfo.map((process) => process.ppid);
+      const mainProcess = processInfo.find((process) => parentPIDs.includes(process.pid));
+      return mainProcess || null;
+    } else {
+      return processInfo || null;
+    }
+  } catch {
+    return null;
+  }
+};
+
+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);
+    setTimeout(() => {
+      clearInterval(checkInterval);
+      resolve();
+    }, 6000);
+  });
+};
+
+export const openProcess = (command: string, args?: string[], options?: SpawnOptions): void => {
+  void (process.platform === "darwin"
+    ? execSync(command)
+    : spawn(command, args ?? [], options ?? {}).unref());
+};
+
+export const getUserData = (): UserData => {
+  const name = execSync("logname", { encoding: "utf8" }).toString().trim().replace(/\n$/, "");
+  const env = Object.assign({}, process.env, { HOME: `/home/${name}` });
+  const uid = execSync(`id -u ${name}`, { encoding: "utf8" }).toString().trim().replace(/\n$/, "");
+  const gid = execSync(`id -g ${name}`, { encoding: "utf8" }).toString().trim().replace(/\n$/, "");
+  return { env, uid: Number(uid), gid: Number(gid) };
+};