diff --git a/scripts/inject/index.mts b/scripts/inject/index.mts index c80b49205..af99fcffb 100644 --- a/scripts/inject/index.mts +++ b/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 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 "\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}`, diff --git a/scripts/inject/injector.mts b/scripts/inject/injector.mts index d714c5f34..1403788d5 100644 --- a/scripts/inject/injector.mts +++ b/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 => { 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,79 @@ 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 => { + let result; + + const processName = 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 ${PlatformNames[platform]}`); + break; + } + } + } + + return result; +}; diff --git a/scripts/inject/types.mts b/scripts/inject/types.mts index 810797c12..ad7671975 100644 --- a/scripts/inject/types.mts +++ b/scripts/inject/types.mts @@ -5,3 +5,14 @@ export type DiscordPlatform = "stable" | "ptb" | "canary" | "dev"; export interface PlatformModule { getAppDir: (platform: DiscordPlatform) => Promisable; } + +export interface ProcessInfo { + pid: number; + ppid: number; +} + +export interface UserData { + env: NodeJS.ProcessEnv; + uid: number; + gid: number; +} diff --git a/scripts/inject/util.mts b/scripts/inject/util.mts index c7d8fa13c..0e97ff20a 100644 --- a/scripts/inject/util.mts +++ b/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 => { + 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) }; +};