fix import-auto-injection cannot handle i64 (#13)

Author: XMadrid

package-lock.json

@@ -18,7 +18,7 @@
     "test:cpp": "cmake -B build -S . && cmake --build build --target wasm-instrumentation-test wasm-opt && build/bin/wasm-instrumentation-test",
     "test": "npm run test:as && npm run test:ts && npm run test:cpp",
     "lint": "eslint src assembly tests/ts/test --max-warnings=0 && prettier -c .",
-    "lint:fix": "eslint src assembly --fix && npx prettier --write .",
+    "lint:fix": "eslint src assembly --fix && prettier --write .",
     "example": "node bin/as-test.js --config example/as-test.config.cjs ; node bin/as-test.js --config example/as-test.config.js"
   },
   "dependencies": {
@@ -43,6 +43,7 @@
     "@types/glob": "^7.2.0",
     "@types/jest": "^29.5.2",
     "@types/node": "^20.9.1",
+    "wasmparser": "5.11.1",
     "assemblyscript-prettier": "^3.0.1",
     "cross-env": "^7.0.3",
     "jest": "^29.5.0",

package.json

@@ -7,7 +7,8 @@ import { AssertResult } from "../assertResult.js";
 import { Imports, ImportsArgument } from "../index.js";
 import { IAssertResult, InstrumentResult } from "../interface.js";
 import { mockInstruFunc, covInstruFunc } from "../utils/import.js";
-import { parseWasmImports, supplyDefaultFunction } from "../utils/index.js";
+import { supplyDefaultFunction } from "../utils/index.js";
+import { parseImportFunctionInfo } from "../utils/wasmparser.js";
 const readFile = promises.readFile;
 
 function nodeExecutor(wasms: string[], outFolder: string, imports: Imports) {
@@ -31,9 +32,8 @@ function nodeExecutor(wasms: string[], outFolder: string, imports: Imports) {
         ...userDefinedImportsObject,
       } as ASImports;
       const binary = await readFile(wasm);
-      const importList = await parseWasmImports(binary);
-      // supplying default function here, so no more need to define all of them in as-test.js
-      supplyDefaultFunction(importList, importObject);
+      const importFuncList = parseImportFunctionInfo(binary);
+      supplyDefaultFunction(importFuncList, importObject);
       const ins = await instantiate(binary, importObject);
       importsArg.module = ins.module;
       importsArg.instance = ins.instance;

src/core/execute.ts

@@ -3,6 +3,8 @@
  */
 // input
 
+import { Type } from "wasmparser";
+
 // instrumented file information
 export interface InstrumentResult {
   sourceWasm: string;
@@ -64,6 +66,13 @@ export interface IAssertResult {
   failed_info: AssertFailMessage;
 }
 
+export interface ImportFunctionInfo {
+  module: string;
+  name: string;
+  args: Type[];
+  return: Type | undefined;
+}
+
 // output
 export class Rate {
   used = 0;

src/interface.ts

@@ -1,4 +1,6 @@
 import { Imports as ASImports } from "@assemblyscript/loader";
+import { ImportFunctionInfo } from "../interface.js";
+import { TypeKind } from "wasmparser/dist/cjs/WasmParser.js";
 
 export function json2map<V>(json: Record<string, V>): Map<string, V> {
   const res = new Map<string, V>();
@@ -43,36 +45,18 @@ export function checkGenerics(functionName: string): string | undefined {
   return;
 }
 
-// list imports of a given wasm binary (buffer)
-// importList format should be as follows：
-// [
-//   { module: 'env', name: 'memory', kind: 'memory' },
-//   { module: 'env', name: 'myFunction', kind: 'function' },
-//   ...
-// ]
-export async function parseWasmImports(binary: Buffer) {
-  const mod = await WebAssembly.compile(binary);
-  const importList = WebAssembly.Module.imports(mod);
-
-  return importList;
-}
-
-export function supplyDefaultFunction(importList: WebAssembly.ModuleImportDescriptor[], importObject: ASImports) {
-  for (const imp of importList) {
-    if (imp.kind === "function") {
-      const moduleName = imp.module;
-      const funcName = imp.name;
-      if (importObject[moduleName]?.[funcName] === undefined) {
-        if (importObject[moduleName] === undefined) {
-          importObject[moduleName] = {};
-        }
-        // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any
-        (importObject[moduleName] as any)[funcName] = (...args: any[]): any => {
-          // notify that a default function has been called
-          console.log(`Default stub called for ${moduleName}.${funcName}, args:`, args);
-          return 0;
-        };
+export function supplyDefaultFunction(infos: ImportFunctionInfo[], importObject: ASImports) {
+  for (const info of infos) {
+    const module = info.module;
+    const name = info.name;
+    if (importObject[module]?.[name] === undefined) {
+      if (importObject[module] === undefined) {
+        importObject[module] = {};
       }
+      // eslint-disable-next-line @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-explicit-any, @typescript-eslint/no-unused-vars
+      (importObject[module] as any)[name] = (..._args: unknown[]): unknown => {
+        return info.return?.kind === TypeKind.i64 ? BigInt(0) : 0;
+      };
     }
   }
 }

src/utils/index.ts

@@ -0,0 +1,77 @@
+import {
+  BinaryReader,
+  BinaryReaderState,
+  IImportEntry,
+  ExternalKind,
+  ISectionInformation,
+  SectionCode,
+  ITypeEntry,
+} from "wasmparser";
+import { ImportFunctionInfo } from "../interface.js";
+import assert from "node:assert";
+
+export function parseImportFunctionInfo(buf: Uint8Array) {
+  const reader = new BinaryReader();
+  const types: ITypeEntry[] = [];
+  const result: ImportFunctionInfo[] = [];
+  reader.setData(buf, 0, buf.length);
+  // eslint-disable-next-line no-constant-condition
+  while (true) {
+    if (!reader.read()) {
+      return result;
+    }
+    switch (reader.state) {
+      case BinaryReaderState.END_WASM: {
+        break;
+      }
+      case BinaryReaderState.ERROR: {
+        throw reader.error;
+      }
+      case BinaryReaderState.END_SECTION: {
+        break;
+      }
+      case BinaryReaderState.BEGIN_SECTION: {
+        const sectionInfo = reader.result as ISectionInformation;
+        switch (sectionInfo.id) {
+          case SectionCode.Type:
+          case SectionCode.Import: {
+            break;
+          }
+          default: {
+            reader.skipSection();
+            break;
+          }
+        }
+        break;
+      }
+      case BinaryReaderState.TYPE_SECTION_ENTRY: {
+        const typeEntry = reader.result as ITypeEntry;
+        types.push(typeEntry);
+        break;
+      }
+      case BinaryReaderState.IMPORT_SECTION_ENTRY: {
+        const importInfo = reader.result as IImportEntry;
+        const decoder = new TextDecoder("utf8");
+        if (importInfo.kind === ExternalKind.Function) {
+          const typeIdx = importInfo.funcTypeIndex;
+          assert(typeIdx !== undefined, "ImportFunction must have a typeIndex");
+          const typeItem = types[typeIdx];
+          assert(typeItem !== undefined, "ImportFunction must have a typeItem");
+          assert(typeItem.params !== undefined);
+          assert(typeItem.returns !== undefined);
+          const returnValue = typeItem.returns.length === 0 ? undefined : typeItem.returns[0];
+          result.push({
+            module: decoder.decode(importInfo.module),
+            name: decoder.decode(importInfo.field),
+            args: typeItem.params,
+            return: returnValue,
+          });
+        }
+        break;
+      }
+      default: {
+        break;
+      }
+    }
+  }
+}

src/utils/wasmparser.ts

@@ -2,14 +2,15 @@ import fs from "fs-extra";
 import { join } from "node:path";
 import { Imports as ASImports } from "@assemblyscript/loader";
 import { fileURLToPath, URL } from "node:url";
-import { DebugInfo, CovDebugInfo } from "../../../../src/interface.js";
+import { DebugInfo, CovDebugInfo, ImportFunctionInfo } from "../../../../src/interface.js";
 import {
   isIncluded,
   json2map,
   checkFunctionName,
   checkGenerics,
   supplyDefaultFunction,
 } from "../../../../src/utils/index.js";
+import { Type } from "wasmparser";
 
 const __dirname = fileURLToPath(new URL(".", import.meta.url));
 
@@ -75,31 +76,28 @@ test("checkGenerics", () => {
 
 describe("supplyDefaultFunction", () => {
   test("supplyTest", () => {
-    const mockImportList: WebAssembly.ModuleImportDescriptor[] = [
-      { kind: "function", module: "myenv", name: "processEvent" },
-      { kind: "function", module: "externalMath", name: "add" },
-      { kind: "function", module: "system", name: "getStatus" },
-      { kind: "function", module: "logger", name: "logWarning" },
-      { kind: "function", module: "customOps", name: "combineValues" },
-      { kind: "global", module: "myenv", name: "globalVar" },
-      { kind: "memory", module: "other", name: "memChange" },
+    const mockInfos: ImportFunctionInfo[] = [
+      { module: "ns", name: "ut.i32", args: [new Type(-1)], return: new Type(-1) },
+      { module: "ns", name: "ut.i64", args: [new Type(-2)], return: new Type(-2) },
+      { module: "ns", name: "ut.f32", args: [new Type(-3)], return: new Type(-3) },
+      { module: "ns", name: "ut.f64", args: [new Type(-4)], return: new Type(-4) },
     ];
 
     const mockImportObject: ASImports = {
-      myenv: {},
-      externalMath: {},
-      system: {},
-      logger: {},
-      customOps: {},
+      env: {},
+      wasi_snapshot_preview1: {},
     };
+    supplyDefaultFunction(mockInfos, mockImportObject);
 
-    supplyDefaultFunction(mockImportList, mockImportObject);
-
-    expect(typeof mockImportObject["myenv"]?.["processEvent"]).toBe("function");
-    expect(typeof mockImportObject["system"]?.["getStatus"]).toBe("function");
-    expect(typeof mockImportObject["logger"]?.["logWarning"]).toBe("function");
-    expect(typeof mockImportObject["customOps"]?.["combineValues"]).toBe("function");
-    expect(mockImportObject["myenv"]?.["globalVar"]).toBeUndefined();
-    expect(mockImportObject["other"]?.["memChange"]).toBeUndefined();
+    expect(typeof mockImportObject["ns"]?.["ut.i32"]).toBe("function");
+    expect(typeof mockImportObject["ns"]?.["ut.i64"]).toBe("function");
+    expect(typeof mockImportObject["ns"]?.["ut.f32"]).toBe("function");
+    expect(typeof mockImportObject["ns"]?.["ut.f64"]).toBe("function");
+    /* eslint-disable @typescript-eslint/ban-types */
+    expect((mockImportObject["ns"]?.["ut.i32"] as Function)(0)).toEqual(0);
+    expect((mockImportObject["ns"]?.["ut.i64"] as Function)(0)).toEqual(BigInt(0));
+    expect((mockImportObject["ns"]?.["ut.f32"] as Function)(0)).toEqual(0);
+    expect((mockImportObject["ns"]?.["ut.f64"] as Function)(0)).toEqual(0);
+    /* eslint-enable @typescript-eslint/ban-types */
   });
 });

tests/ts/fixture/defaultImportTest.wasm

@@ -0,0 +1,52 @@
+import { parseImportFunctionInfo } from "../../../../src/utils/wasmparser.js";
+import { readFileSync } from "node:fs";
+import { join } from "node:path";
+import { fileURLToPath, URL } from "node:url";
+import { Type } from "wasmparser";
+
+const __dirname = fileURLToPath(new URL(".", import.meta.url));
+
+test("parseImportFunctionInfo", () => {
+  const fp = join(__dirname, "..", "..", "fixture", "defaultImportTest.wasm");
+  const buf = readFileSync(fp);
+  const expectedInfo = [
+    {
+      module: "env",
+      name: "abort",
+      args: [new Type(-1), new Type(-1), new Type(-1), new Type(-1)],
+      return: undefined,
+    },
+    {
+      module: "env",
+      name: "logInfo",
+      args: [new Type(-1), new Type(-1)],
+      return: undefined,
+    },
+    {
+      module: "env",
+      name: "getTimeSinceEpoch",
+      args: [],
+      return: new Type(-2),
+    },
+    {
+      module: "ns",
+      name: "ut.i32",
+      args: [],
+      return: new Type(-1),
+    },
+    {
+      module: "ns",
+      name: "ut.f32",
+      args: [new Type(-2)],
+      return: new Type(-3),
+    },
+    {
+      module: "ns",
+      name: "ut.f64",
+      args: [],
+      return: new Type(-4),
+    },
+  ];
+
+  expect(parseImportFunctionInfo(buf)).toEqual(expectedInfo);
+});