Validation: Ban function references to undeclared functions

Author: kateinoigakukun

Sources/WasmKit/Execution/ConstEvaluation.swift

@@ -8,10 +8,16 @@ protocol ConstEvaluationContextProtocol {
 struct ConstEvaluationContext: ConstEvaluationContextProtocol {
     let functions: ImmutableArray<InternalFunction>
     var globals: [Value]
+    let onFunctionReferenced: ((InternalFunction) -> Void)?
 
-    init(functions: ImmutableArray<InternalFunction>, globals: [Value]) {
+    init(
+        functions: ImmutableArray<InternalFunction>,
+        globals: [Value],
+        onFunctionReferenced: ((InternalFunction) -> Void)? = nil
+    ) {
         self.functions = functions
         self.globals = globals
+        self.onFunctionReferenced = onFunctionReferenced
     }
 
     init(instance: InternalInstance, moduleImports: ModuleImports) {
@@ -23,7 +29,9 @@ struct ConstEvaluationContext: ConstEvaluationContextProtocol {
     }
 
     func functionRef(_ index: FunctionIndex) throws -> Reference {
-        return try .function(from: self.functions[validating: Int(index)])
+        let function = try self.functions[validating: Int(index)]
+        self.onFunctionReferenced?(function)
+        return .function(from: function)
     }
     func globalValue(_ index: GlobalIndex) throws -> Value {
         guard index < globals.count else {

Sources/WasmKit/Execution/Instances.swift

@@ -76,6 +76,7 @@ struct InstanceEntity /* : ~Copyable */ {
     var elementSegments: ImmutableArray<InternalElementSegment>
     var dataSegments: ImmutableArray<InternalDataSegment>
     var exports: [String: InternalExternalValue]
+    var functionRefs: Set<InternalFunction>
     var features: WasmFeatureSet
     var hasDataCount: Bool
 
@@ -89,6 +90,7 @@ struct InstanceEntity /* : ~Copyable */ {
             elementSegments: ImmutableArray(),
             dataSegments: ImmutableArray(),
             exports: [:],
+            functionRefs: [],
             features: [],
             hasDataCount: false
         )

Sources/WasmKit/Execution/StoreAllocator.swift

@@ -360,11 +360,16 @@ extension StoreAllocator {
             allocateHandle: { m, _ in try allocate(memoryType: m, resourceLimiter: resourceLimiter) }
         )
 
+        var functionRefs: Set<InternalFunction> = []
         // Step 5.
         let constEvalContext = ConstEvaluationContext(
             functions: functions,
-            globals: importedGlobals.map(\.value)
+            globals: importedGlobals.map(\.value),
+            onFunctionReferenced: { function in
+                functionRefs.insert(function)
+            }
         )
+
         let globals = try allocateEntities(
             imports: importedGlobals,
             internals: module.globals,
@@ -379,13 +384,13 @@ extension StoreAllocator {
         // Step 6.
         let elements = try ImmutableArray<InternalElementSegment>(allocator: arrayAllocator, count: module.elements.count) { buffer in
             for (index, element) in module.elements.enumerated() {
-                let references: [Reference]
+                // TODO: Avoid evaluating element expr twice in `Module.instantiate` and here.
+                var references = try element.evaluateInits(context: constEvalContext)
                 switch element.mode {
                 case .active, .declarative:
                     // active & declarative segments are unavailable at runtime
                     references = []
-                case .passive:
-                    references = try element.evaluateInits(context: constEvalContext)
+                case .passive: break
                 }
                 let handle = allocate(elementType: element.type, references: references)
                 buffer.initializeElement(at: index, to: handle)
@@ -449,6 +454,7 @@ extension StoreAllocator {
             elementSegments: elements,
             dataSegments: dataSegments,
             exports: exports,
+            functionRefs: functionRefs,
             features: module.features,
             hasDataCount: module.hasDataCount
         )

Sources/WasmKit/Translator.swift

@@ -118,7 +118,10 @@ extension InternalInstance: TranslatorContext {
         _ = try self.dataSegments[validating: Int(index)]
     }
     func validateFunctionIndex(_ index: FunctionIndex) throws {
-        _ = try self.functions[validating: Int(index)]
+        let function = try self.functions[validating: Int(index)]
+        guard self.functionRefs.contains(function) else {
+            throw ValidationError("Function index \(index) is not declared but referenced as a function reference")
+        }
     }
 }
 
@@ -1752,7 +1755,7 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         emit(.refIsNull(Instruction.RefIsNullOperand(value: LVReg(ensureOnVReg(value)), result: LVReg(result))))
     }
     mutating func visitRefFunc(functionIndex: UInt32) throws -> Output {
-        try self.module.validateFunctionIndex(functionIndex)
+        try validator.validateRefFunc(functionIndex: functionIndex)
         pushEmit(.ref(.funcRef), { .refFunc(Instruction.RefFuncOperand(index: functionIndex, result: LVReg($0))) })
     }
 

Sources/WasmKit/Validator.swift

@@ -42,6 +42,10 @@ struct InstructionValidator<Context: TranslatorContext> {
             throw ValidationError("Table element type mismatch in table.init: \(tableType.elementType) != \(elementType)")
         }
     }
+
+    func validateRefFunc(functionIndex: UInt32) throws {
+        try context.validateFunctionIndex(functionIndex)
+    }
 }
 
 struct ModuleValidator {

Tests/WasmKitTests/SpectestTests.swift

@@ -22,7 +22,6 @@ final class SpectestTests: XCTestCase {
             path: Self.testPaths,
             include: [],
             exclude: [
-                "ref_func.wast",
                 "select.wast",
                 "start.wast",
                 "table-sub.wast",