diff --git a/src/Netscript/WorkerScript.ts b/src/Netscript/WorkerScript.ts index afd1f55e1..772dd4f96 100644 --- a/src/Netscript/WorkerScript.ts +++ b/src/Netscript/WorkerScript.ts @@ -131,7 +131,7 @@ export class WorkerScript { } shouldLog(fn: string): boolean { - return this.disableLogs[fn] == null; + return !(this.disableLogs.ALL || this.disableLogs[fn]); } log(func: string, txt: () => string): void { diff --git a/src/NetscriptFunctions.ts b/src/NetscriptFunctions.ts index d44b68520..7ec2da28b 100644 --- a/src/NetscriptFunctions.ts +++ b/src/NetscriptFunctions.ts @@ -495,37 +495,46 @@ export const ns: InternalAPI = { }, disableLog: (ctx) => (_fn) => { const fn = helpers.string(ctx, "fn", _fn); + if (possibleLogs[fn] === undefined) { + throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); + } if (fn === "ALL") { - for (const fn of Object.keys(possibleLogs)) { + ctx.workerScript.disableLogs = allDisabled; + // No need to log here, it's been disabled. + } else { + // We don't track individual log entries when all are disabled. + if (!ctx.workerScript.disableLogs["ALL"]) { ctx.workerScript.disableLogs[fn] = true; + helpers.log(ctx, () => `Disabled logging for ${fn}`); } - helpers.log(ctx, () => `Disabled logging for all functions`); - } else if (possibleLogs[fn] === undefined) { - throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); - } else { - ctx.workerScript.disableLogs[fn] = true; - helpers.log(ctx, () => `Disabled logging for ${fn}`); } }, enableLog: (ctx) => (_fn) => { const fn = helpers.string(ctx, "fn", _fn); + if (possibleLogs[fn] === undefined) { + throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); + } if (fn === "ALL") { - for (const fn of Object.keys(possibleLogs)) { - delete ctx.workerScript.disableLogs[fn]; - } + ctx.workerScript.disableLogs = {}; helpers.log(ctx, () => `Enabled logging for all functions`); - } else if (possibleLogs[fn] === undefined) { - throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); + } else { + if (ctx.workerScript.disableLogs["ALL"]) { + // As an optimization, we normally store only that key, but we have to + // expand it out to all keys at this point. + // Conveniently, possibleLogs serves as a model for "all keys disabled." + ctx.workerScript.disableLogs = Object.assign({}, possibleLogs, { ALL: false, [fn]: false }); + } else { + ctx.workerScript.disableLogs[fn] = false; + } + helpers.log(ctx, () => `Enabled logging for ${fn}`); } - delete ctx.workerScript.disableLogs[fn]; - helpers.log(ctx, () => `Enabled logging for ${fn}`); }, isLogEnabled: (ctx) => (_fn) => { const fn = helpers.string(ctx, "fn", _fn); if (possibleLogs[fn] === undefined) { throw helpers.errorMessage(ctx, `Invalid argument: ${fn}.`); } - return !ctx.workerScript.disableLogs[fn]; + return ctx.workerScript.shouldLog(fn); }, getScriptLogs: (ctx) => @@ -1840,7 +1849,13 @@ export function NetscriptFunctions(ws: WorkerScript): NSFull { return NSProxy(ws, ns, [], { args: ws.args.slice(), pid: ws.pid, enums }); } -const possibleLogs = Object.fromEntries([...getFunctionNames(ns, "")].map((a) => [a, true])); +const possibleLogs = Object.fromEntries(getFunctionNames(ns, "").map((a) => [a, true])); +possibleLogs.ALL = true; + +// We reuse this object for *all* scripts that disable all keys, to prevent memory growth. +// Any script that needs a custom set of values will use a fresh object. +const allDisabled = { ALL: true } as const; + /** Provides an array of all function names on a nested object */ function getFunctionNames(obj: object, prefix: string): string[] { const functionNames: string[] = []; diff --git a/test/jest/Netscript/DisableLog.test.ts b/test/jest/Netscript/DisableLog.test.ts new file mode 100644 index 000000000..c7ae21acf --- /dev/null +++ b/test/jest/Netscript/DisableLog.test.ts @@ -0,0 +1,87 @@ +import type { Script } from "../../../src/Script/Script"; +import type { ScriptFilePath } from "../../../src/Paths/ScriptFilePath"; +import { FormatsNeedToChange } from "../../../src/ui/formatNumber"; +import { Server } from "../../../src/Server/Server"; +import { RunningScript } from "../../../src/Script/RunningScript"; +import { CompletedProgramName } from "../../../src/Programs/Enums"; +import { WorkerScript } from "../../../src/Netscript/WorkerScript"; +import { NetscriptFunctions } from "../../../src/NetscriptFunctions"; +import { AddToAllServers, DeleteServer } from "../../../src/Server/AllServers"; + +test("Edge cases of disableLog", function () { + // Ensure that formatting functions work properly + FormatsNeedToChange.emit(); + let server; + try { + server = new Server({ hostname: "home", adminRights: true, maxRam: 8 }); + server.programs.push(CompletedProgramName.bruteSsh, CompletedProgramName.ftpCrack); + AddToAllServers(server); + // We don't need this script to be runnable, it just needs to exist so that + // we can create a RunningScript and WorkerScript object. + expect(server.writeToScriptFile("test.js" as ScriptFilePath, "")).toEqual({ overwritten: false }); + const script = server.scripts.get("test.js" as ScriptFilePath) as Script; + + const runningScript = new RunningScript(script, 2); + const ws = new WorkerScript(runningScript, 1, NetscriptFunctions); + + const ns = ws.env.vars; + + // Generate logs in a specific pattern that checks edge cases in + // disableLog. We want to check various combinations of things that + // are/aren't disabled, as well as previously disabled. + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("before disableLog ALL"); + + expect(() => ns.disableLog("all")).toThrow("Invalid argument: all."); + ns.disableLog("ALL"); + + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("after disableLog ALL"); + + ns.enableLog("brutessh"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("after enableLog brutessh"); + + ns.disableLog("brutessh"); + ns.enableLog("ftpcrack"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("after enableLog ftpcrack"); + + ns.enableLog("ftpcrack"); + ns.brutessh("home"); + ns.ftpcrack("home"); + + ns.print("after redundant enable"); + ns.disableLog("ALL"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.enableLog("ALL"); + ns.brutessh("home"); + ns.ftpcrack("home"); + ns.print("end"); + + expect(runningScript.logs).toEqual([ + "brutessh: Executed BruteSSH.exe on 'home' to open SSH port (22).", + "ftpcrack: Executed FTPCrack.exe on 'home' to open FTP port (21).", + "before disableLog ALL", + "disableLog: Invalid argument: all.", + "after disableLog ALL", + "brutessh: SSH Port (22) already opened on 'home'.", + "after enableLog brutessh", + "ftpcrack: FTP Port (21) already opened on 'home'.", + "after enableLog ftpcrack", + "ftpcrack: FTP Port (21) already opened on 'home'.", + "after redundant enable", + "enableLog: Enabled logging for all functions", + "brutessh: SSH Port (22) already opened on 'home'.", + "ftpcrack: FTP Port (21) already opened on 'home'.", + "end", + ]); + } finally { + DeleteServer(server.hostname); + } +});