diff --git a/src/Netscript/APIWrapper.ts b/src/Netscript/APIWrapper.ts index b6c2648a9..0cfd31138 100644 --- a/src/Netscript/APIWrapper.ts +++ b/src/Netscript/APIWrapper.ts @@ -2,7 +2,6 @@ import { getRamCost } from "./RamCostGenerator"; import type { WorkerScript } from "./WorkerScript"; import { helpers } from "./NetscriptHelpers"; import { ScriptArg } from "./ScriptArg"; -import { NSFull } from "src/NetscriptFunctions"; import { cloneDeep } from "lodash"; /** Generic type for an enums object */ @@ -37,42 +36,57 @@ export type InternalAPI = { /** Any of the possible values on a internal API layer */ type InternalValues = Enums | ScriptArg[] | InternalFn | InternalAPI; +export class StampedLayer { + #workerScript: WorkerScript; + constructor(ws: WorkerScript, obj: ExternalAPI) { + this.#workerScript = ws; + Object.setPrototypeOf(this, obj); + } + static wrapFunction(eLayer: ExternalAPI, func: InternalFn, tree: string[], key: Key) { + const arrayPath = [...tree, key]; + const functionPath = arrayPath.join("."); + function wrappedFunction(this: StampedLayer, ...args: unknown[]): unknown { + const ctx = { workerScript: this.#workerScript, function: key, functionPath }; + helpers.checkEnvFlags(ctx); + helpers.updateDynamicRam(ctx, getRamCost(...tree, key)); + return func(ctx)(...args); + } + Object.defineProperty(eLayer, key, { value: wrappedFunction, enumerable: true, writable: false }); + } +} +Object.defineProperty(StampedLayer.prototype, "constructor", { + value: Object, + enumerable: false, + writable: false, + configurable: false, +}); + export type NetscriptContext = { workerScript: WorkerScript; function: string; functionPath: string; }; -export function wrapAPI(ws: WorkerScript, internalAPI: InternalAPI, args: ScriptArg[]): ExternalAPI { - function wrapAPILayer(eLayer: ExternalAPI, iLayer: InternalAPI, tree: string[]): ExternalAPI { - for (const [key, value] of Object.entries(iLayer) as [Key, InternalValues][]) { - if (key === "enums") { - (eLayer[key] as Enums) = cloneDeep(value as Enums); - } else if (key === "args") continue; - // Args are added in wrapAPI function and should only exist at top level - else if (typeof value === "function") { - wrapFunction(eLayer, value as InternalFn, tree, key); - } else if (typeof value === "object") { - wrapAPILayer((eLayer[key] = {} as ExternalAPI[Key]), value, [...tree, key as string]); - } else { - console.warn(`Unexpected data while wrapping API.`, "tree:", tree, "key:", key, "value:", value); - throw new Error("Error while wrapping netscript API. See console."); - } +export function wrapAPILayer( + eLayer: ExternalAPI, + iLayer: InternalAPI, + tree: string[], +): ExternalAPI { + for (const [key, value] of Object.entries(iLayer) as [Key, InternalValues][]) { + if (key === "enums") { + const enumObj = Object.freeze(cloneDeep(value as Enums)); + for (const member of Object.values(enumObj)) Object.freeze(member); + (eLayer[key] as Enums) = enumObj; + } else if (key === "args") continue; + // Args only added on individual instances. + else if (typeof value === "function") { + StampedLayer.wrapFunction(eLayer, value as InternalFn, tree, key); + } else if (typeof value === "object") { + wrapAPILayer((eLayer[key] = {} as ExternalAPI[Key]), value, [...tree, key as string]); + } else { + console.warn(`Unexpected data while wrapping API.`, "tree:", tree, "key:", key, "value:", value); + throw new Error("Error while wrapping netscript API. See console."); } - return eLayer; } - function wrapFunction(eLayer: ExternalAPI, func: InternalFn, tree: string[], key: Key) { - const arrayPath = [...tree, key]; - const functionPath = arrayPath.join("."); - const ctx = { workerScript: ws, function: key, functionPath }; - function wrappedFunction(...args: unknown[]): unknown { - helpers.checkEnvFlags(ctx); - helpers.updateDynamicRam(ctx, getRamCost(...tree, key)); - return func(ctx)(...args); - } - (eLayer[key] as WrappedFn) = wrappedFunction; - } - - const wrappedAPI = wrapAPILayer({ args } as ExternalAPI, internalAPI, []); - return wrappedAPI; + return eLayer; } diff --git a/src/NetscriptFunctions.ts b/src/NetscriptFunctions.ts index 5d9164bdc..cfa18dbee 100644 --- a/src/NetscriptFunctions.ts +++ b/src/NetscriptFunctions.ts @@ -72,7 +72,7 @@ import { Flags } from "./NetscriptFunctions/Flags"; import { calculateIntelligenceBonus } from "./PersonObjects/formulas/intelligence"; import { CalculateShareMult, StartSharing } from "./NetworkShare/Share"; import { recentScripts } from "./Netscript/RecentScripts"; -import { ExternalAPI, InternalAPI, wrapAPI } from "./Netscript/APIWrapper"; +import { ExternalAPI, InternalAPI, StampedLayer, wrapAPILayer } from "./Netscript/APIWrapper"; import { INetscriptExtra } from "./NetscriptFunctions/Extra"; import { ScriptDeath } from "./Netscript/ScriptDeath"; import { getBitNodeMultipliers } from "./BitNode/BitNode"; @@ -93,10 +93,6 @@ export const enums: NSEnums = { export type NSFull = Readonly; -export function NetscriptFunctions(workerScript: WorkerScript): ExternalAPI { - return wrapAPI(workerScript, ns, workerScript.args.slice()); -} - const base: InternalAPI = { args: [], enums, @@ -1913,6 +1909,34 @@ export const ns = { ...NetscriptExtra(), }; +export const wrappedNS = wrapAPILayer({} as ExternalAPI, ns, []); + +// Figure out once which layers of ns have functions on them and will need to be stamped with a private workerscript field for API access +const layerLocations: string[][] = []; +function populateLayers(nsLayer: ExternalAPI, currentLayers: string[] = []) { + for (const [k, v] of Object.entries(nsLayer)) { + if (typeof v === "object" && k !== "enums") { + if (Object.values(v as object).some((member) => typeof member === "function")) + layerLocations.push([...currentLayers, k]); + populateLayers(v as ExternalAPI, [...currentLayers, k]); + } + } +} +populateLayers(wrappedNS); + +export function NetscriptFunctions(ws: WorkerScript): ExternalAPI { + //todo: better typing instead of relying on an any + const instance = new StampedLayer(ws, wrappedNS) as any; + for (const layerLocation of layerLocations) { + const key = layerLocation.pop() as string; + const obj = layerLocation.reduce((prev, curr) => prev[curr], instance); + layerLocation.push(key); + obj[key] = new StampedLayer(ws, obj[key]); + } + instance.args = ws.args.slice(); + return instance; +} + const possibleLogs = Object.fromEntries([...getFunctionNames(ns, "")].map((a) => [a, true])); /** Provides an array of all function names on a nested object */ function getFunctionNames(obj: object, prefix: string): string[] { diff --git a/src/NetscriptWorker.ts b/src/NetscriptWorker.ts index 9e9524edf..ded347790 100644 --- a/src/NetscriptWorker.ts +++ b/src/NetscriptWorker.ts @@ -11,7 +11,7 @@ import { generateNextPid } from "./Netscript/Pid"; import { CONSTANTS } from "./Constants"; import { Interpreter } from "./ThirdParty/JSInterpreter"; -import { NetscriptFunctions } from "./NetscriptFunctions"; +import { NetscriptFunctions, wrappedNS } from "./NetscriptFunctions"; import { compile, Node } from "./NetscriptJSEvaluator"; import { IPort } from "./NetscriptPort"; import { RunningScript } from "./Script/RunningScript"; @@ -81,14 +81,14 @@ async function startNetscript1Script(workerScript: WorkerScript): Promise //TODO: Make NS1 wrapping type safe instead of using BasicObject type BasicObject = Record; - function wrapNS1Layer(int: Interpreter, intLayer: unknown, nsLayer = workerScript.env.vars as BasicObject) { + function wrapNS1Layer(int: Interpreter, intLayer: unknown, nsLayer = wrappedNS as BasicObject) { for (const [name, entry] of Object.entries(nsLayer)) { if (typeof entry === "function") { const wrapper = async (...args: unknown[]) => { try { // Sent a resolver function as an extra arg. See createAsyncFunction JSInterpreter.js:3209 const callback = args.pop() as (value: unknown) => void; - const result = await entry(...args.map((arg) => int.pseudoToNative(arg))); + const result = await entry.bind(workerScript.env.vars)(...args.map((arg) => int.pseudoToNative(arg))); return callback(int.nativeToPseudo(result)); } catch (e: unknown) { errorToThrow = e; diff --git a/test/jest/Netscript/RamCalculation.test.ts b/test/jest/Netscript/RamCalculation.test.ts index 7ac34d9fd..8829e3a49 100644 --- a/test/jest/Netscript/RamCalculation.test.ts +++ b/test/jest/Netscript/RamCalculation.test.ts @@ -1,5 +1,5 @@ import { Player } from "../../../src/Player"; -import { NetscriptFunctions } from "../../../src/NetscriptFunctions"; +import { NetscriptFunctions, wrappedNS } from "../../../src/NetscriptFunctions"; import { RamCosts, getRamCost, RamCostConstants } from "../../../src/Netscript/RamCostGenerator"; import { Environment } from "../../../src/Netscript/Environment"; import { RunningScript } from "../../../src/Script/RunningScript"; @@ -99,11 +99,11 @@ describe("Netscript RAM Calculation/Generation Tests", function () { } describe("ns", () => { - Object.entries(ns as unknown as NSLayer).forEach(([key, val]) => { + Object.entries(wrappedNS as unknown as NSLayer).forEach(([key, val]) => { if (key === "args" || key === "enums") return; if (typeof val === "function") { const expectedRam = grabCost(RamCosts, [key]); - it(`${key}()`, () => combinedRamCheck(val, [key], expectedRam)); + it(`${key}()`, () => combinedRamCheck(val.bind(ns), [key], expectedRam)); } //The only other option should be an NSLayer const extraLayerCost = { corporation: 1022.4, hacknet: 4 }[key] ?? 0; @@ -112,13 +112,15 @@ describe("Netscript RAM Calculation/Generation Tests", function () { }); function testLayer(nsLayer: NSLayer, ramLayer: RamLayer, path: string[], extraLayerCost: number) { + // nsLayer is the layer on the main, unstamped wrappedNS object. The actualLayer is needed to check correct stamping. + const actualLayer = path.reduce((prev, curr) => prev[curr], ns as any); //todo: do this typesafely? describe(path[path.length - 1], () => { Object.entries(nsLayer).forEach(([key, val]) => { const newPath = [...path, key]; if (typeof val === "function") { const fnName = newPath.join("."); const expectedRam = grabCost(ramLayer, newPath); - it(`${fnName}()`, () => combinedRamCheck(val, newPath, expectedRam, extraLayerCost)); + it(`${fnName}()`, () => combinedRamCheck(val.bind(actualLayer), newPath, expectedRam, extraLayerCost)); } //Skip enums layers else if (key === "enums") return; @@ -130,11 +132,11 @@ describe("Netscript RAM Calculation/Generation Tests", function () { describe("Singularity multiplier checks", () => { sf4.lvl = 3; - const singFunctions = Object.entries(ns.singularity).filter(([__, val]) => typeof val === "function"); + const singFunctions = Object.entries(wrappedNS.singularity).filter(([__, val]) => typeof val === "function"); const singObjects = singFunctions.map(([key, val]) => { return { name: key, - fn: val, + fn: val.bind(ns.singularity), baseRam: grabCost(RamCosts.singularity, ["singularity", key]), }; });