diff --git a/loader/src/index.cjs b/loader/src/index.cjs index 7eb2e5563..484ad70a8 100644 --- a/loader/src/index.cjs +++ b/loader/src/index.cjs @@ -93,8 +93,6 @@ const metering = require('@permaweb/wasm-metering') */ /* - * Custom WebAssembly.compileStreaming Implementation with WASM Metering - * * This implementation overrides WebAssembly.compileStreaming to add metering * to WebAssembly binaries. Metering enables tracking of resource usage (e.g., * CPU or memory) within the WebAssembly runtime, which is critical for monitoring @@ -106,28 +104,43 @@ const metering = require('@permaweb/wasm-metering') * instead of within each CU instance. By integrating metering directly into the * loader, we ensure that all WebAssembly binaries are metered consistently across * different CUs, reducing redundancy and enhancing modularity. - */ +*/ +// Save the original WebAssembly.compile and compileStreaming functions +const originalCompileStreaming = WebAssembly.compileStreaming +const originalCompile = WebAssembly.compile -// Override WebAssembly.compileStreaming to apply metering -WebAssembly.compileStreaming = async function (source, importObject = {}) { - // Read the response and convert it to an ArrayBuffer - const arrayBuffer = await source.arrayBuffer(); +const shouldApplyMetering = (importObject = {}) => { + return importObject.applyMetering && ['wasm32-unknown-emscripten4', 'wasm64-unknown-emscripten-draft_2024_02_15'].includes(importObject.format); +} - // Convert ArrayBuffer to Uint8Array for metering compatibility - const nodeBuffer = Buffer.from(arrayBuffer); +const applyMetering = (arrayBuffer, importObject) => { + const { format } = importObject - if(importObject.format === "wasm32-unknown-emscripten" || importObject.format === "wasm32-unknown-emscripten2" || importObject.format === "wasm32-unknown-emscripten3") { - return WebAssembly.compile(nodeBuffer); - } + const view = ArrayBuffer.isView(arrayBuffer) + ? arrayBuffer + : Buffer.from(arrayBuffer) + + // Determine meter type and apply metering + const meterType = format.startsWith('wasm32') ? 'i32' : 'i64' + const meteredView = metering.meterWASM(view, { meterType }) - let meterType = importObject.format.startsWith("wasm32") ? "i32" : "i64"; + return originalCompile(meteredView.buffer) +} + +// Override WebAssembly.compileStreaming to apply metering +WebAssembly.compileStreaming = async function (source, importObject = { applyMetering: false }) { + if (!shouldApplyMetering(importObject)) return originalCompileStreaming(source) - // Apply metering with the Uint8Array buffer - const meteredView = metering.meterWASM(nodeBuffer, { meterType }); + const arrayBuffer = await source.arrayBuffer() + return applyMetering(arrayBuffer, importObject) +} + +// Override WebAssembly.compile to apply metering +WebAssembly.compile = async function (source, importObject = { applyMetering: false }) { + if (!shouldApplyMetering(importObject)) return originalCompile(source) - // const meteredResponse = new Response(meteredBuffer, { headers: { 'Content-Type': 'application/wasm' } }); - return WebAssembly.compile(meteredView.buffer); -}; + return applyMetering(source, importObject) +} module.exports = async function (binary, options) { let instance = null diff --git a/loader/test/emscripten2.test.js b/loader/test/emscripten2.test.js index 2157be4f7..331d2c41a 100644 --- a/loader/test/emscripten2.test.js +++ b/loader/test/emscripten2.test.js @@ -56,7 +56,7 @@ describe('loader', async () => { Readable.toWeb(createReadStream('./test/process/process.wasm')), { headers: { 'Content-Type': 'application/wasm' } } ), - { format: 'wasm32-unknown-emscripten2' } + { format: 'wasm32-unknown-emscripten2', applyMetering: false } ) const handle = await AoLoader((info, receiveInstance) => { diff --git a/loader/test/index.test.js b/loader/test/index.test.js index 15d74eead..5359546ad 100644 --- a/loader/test/index.test.js +++ b/loader/test/index.test.js @@ -56,7 +56,7 @@ describe('loader', async () => { Readable.toWeb(createReadStream('./test/legacy/process.wasm')), { headers: { 'Content-Type': 'application/wasm' } } ), - { format: 'wasm32-unknown-emscripten' } + { format: 'wasm32-unknown-emscripten', applyMetering: false } ) const handle = await AoLoader((info, receiveInstance) => {