Merge pull request #98 from swiftwasm/katei/fix-heap-overflow

kateinoigakukun · web-flow · commit d30a1de4fc0e · 2024-07-01T12:23:50.000+09:00
Fix value stack overflow
diff --git a/Makefile b/Makefile
@@ -1,18 +1,5 @@
 NAME := WasmKit
 
-MODULES = $(notdir $(wildcard Sources/*))
-
-.PHONY: all
-all: build
-
-.PHONY: build
-build:
-	@swift build
-
-.PHONY: test
-test:
-	@swift test
-
 .PHONY: docs
 docs:
 	swift package generate-documentation --target WasmKit
@@ -37,7 +24,7 @@ $(SPECTEST_ROOT)/%.json: $(TESTSUITE_DIR)/%.wast
 
 .PHONY: spectest
 spectest: spec
-	swift run Spectest $(SPECTEST_ROOT)
+	swift run --sanitize address Spectest $(SPECTEST_ROOT)
 
 
 ### WASI Test Suite
@@ -52,11 +39,3 @@ wasitest:
 generate:
 	swift ./Utilities/generate_inst_visitor.swift
 	swift ./Utilities/generate_inst_dispatch.swift
-
-GIT_STATUS = $(shell git status --porcelain)
-ensure_clean:
-	@[ -z "$(GIT_STATUS)" ] \
-    && echo Working directory is clean \
-	|| (printf "Uncommitted changes: \n $(GIT_STATUS)\n" && exit 1)
-
-FORCE:
diff --git a/Sources/WasmKit/Execution/Instructions/Expression.swift b/Sources/WasmKit/Execution/Instructions/Expression.swift
@@ -2,19 +2,14 @@ import WasmParser
 
 struct InstructionSequence: Equatable {
     let instructions: UnsafeBufferPointer<Instruction>
-
-    init(instructions: [Instruction]) {
-        assert(_isPOD(Instruction.self))
-        let buffer = UnsafeMutableBufferPointer<Instruction>.allocate(capacity: instructions.count + 1)
-        for (idx, instruction) in instructions.enumerated() {
-            buffer[idx] = instruction
-        }
-        buffer[instructions.count] = .endOfFunction
-        self.instructions = UnsafeBufferPointer(buffer)
-    }
-
-    func deallocate() {
-        instructions.deallocate()
+    /// The maximum height of the value stack during execution of this function.
+    /// This height does not count the locals.
+    let maxStackHeight: Int
+
+    init(instructions: UnsafeBufferPointer<Instruction>, maxStackHeight: Int) {
+        self.instructions = instructions
+        assert(self.instructions.last == .endOfFunction)
+        self.maxStackHeight = maxStackHeight
     }
 
     var baseAddress: UnsafePointer<Instruction> {
@@ -26,12 +21,6 @@ struct InstructionSequence: Equatable {
     }
 }
 
-extension InstructionSequence: ExpressibleByArrayLiteral {
-    init(arrayLiteral elements: Instruction...) {
-        self.init(instructions: elements)
-    }
-}
-
 struct ExpressionRef: Equatable {
     let _relativeOffset: UInt32
     var relativeOffset: Int {
diff --git a/Sources/WasmKit/Execution/Runtime/Function.swift b/Sources/WasmKit/Execution/Runtime/Function.swift
@@ -8,6 +8,7 @@ public struct Function: Equatable {
     public func invoke(_ arguments: [Value] = [], runtime: Runtime) throws -> [Value] {
         try withExecution { execution in
             var stack = Stack()
+            defer { stack.deallocate() }
             let numberOfResults = try invoke(execution: &execution, stack: &stack, with: arguments, runtime: runtime)
             try execution.run(runtime: runtime, stack: &stack)
             return Array(stack.popValues(count: numberOfResults))
diff --git a/Sources/WasmKit/Execution/Runtime/Runtime.swift b/Sources/WasmKit/Execution/Runtime/Runtime.swift
@@ -87,15 +87,21 @@ extension Runtime {
                         .numericConst(.i32(UInt32(element.initializer.count))),
                         .tableInit(tableIndex, elementIndex),
                         .tableElementDrop(elementIndex),
+                        .endOfFunction,
                     ])
-                    let initIseq = InstructionSequence(instructions: instructions)
-                    defer { initIseq.deallocate() }
-                    try evaluateConstExpr(initIseq, instance: instance)
+                    try instructions.withUnsafeBufferPointer {
+                        let initIseq = InstructionSequence(instructions: $0, maxStackHeight: 2)
+                        try evaluateConstExpr(initIseq, instance: instance)
+                    }
 
                 case .declarative:
-                    let initIseq: InstructionSequence = [.tableElementDrop(elementIndex)]
-                    defer { initIseq.deallocate() }
-                    try evaluateConstExpr(initIseq, instance: instance)
+                    let instructions: [Instruction] = [.tableElementDrop(elementIndex), .endOfFunction]
+                    try instructions.withUnsafeBufferPointer {
+                        let initIseq = InstructionSequence(
+                            instructions: $0, maxStackHeight: 0
+                        )
+                        try evaluateConstExpr(initIseq, instance: instance)
+                    }
 
                 case .passive:
                     continue
@@ -122,14 +128,17 @@ extension Runtime {
                 default:
                     throw InstantiationError.unsupported("init expr in data section \(data.offset)")
                 }
-                let iseq = InstructionSequence(instructions: instructions + [
+                instructions.append(contentsOf: [
                     .numericConst(.i32(0)),
                     .numericConst(.i32(UInt32(data.initializer.count))),
                     .memoryInit(UInt32(dataIndex)),
                     .memoryDataDrop(UInt32(dataIndex)),
+                    .endOfFunction,
                 ])
-                defer { iseq.deallocate() }
-                try evaluateConstExpr(iseq, instance: instance)
+                try instructions.withUnsafeBufferPointer {
+                    let iseq = InstructionSequence(instructions: $0, maxStackHeight: 2)
+                    try evaluateConstExpr(iseq, instance: instance)
+                }
             }
         } catch Trap.outOfBoundsMemoryAccess {
             throw InstantiationError.outOfBoundsMemoryAccess
@@ -141,6 +150,7 @@ extension Runtime {
         if let startIndex = module.start {
             try withExecution { initExecution in
                 var stack = Stack()
+                defer { stack.deallocate() }
                 try initExecution.invoke(functionAddress: instance.functionAddresses[Int(startIndex)], runtime: self, stack: &stack)
                 try initExecution.run(runtime: self, stack: &stack)
             }
@@ -189,10 +199,12 @@ extension Runtime {
                 default:
                     throw InstantiationError.unsupported("init expr in global section \(global.initializer)")
                 }
-                let iseq = InstructionSequence(instructions: instructions)
-                defer { iseq.deallocate() }
-                return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
-                    return stack.popValue()
+                instructions.append(.endOfFunction)
+                return try instructions.withUnsafeBufferPointer {
+                    let iseq = InstructionSequence(instructions: $0, maxStackHeight: 1)
+                    return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
+                        return stack.popValue()
+                    }
                 }
             }
             
@@ -212,6 +224,7 @@ extension Runtime {
     ) throws -> T {
         try withExecution { initExecution in
             var stack = Stack()
+            defer { stack.deallocate() }
             try stack.pushFrame(
                 iseq: iseq,
                 arity: arity,
diff --git a/Sources/WasmKit/Execution/Runtime/Stack.swift b/Sources/WasmKit/Execution/Runtime/Stack.swift
@@ -23,8 +23,7 @@ struct Stack {
         returnPC: ProgramCounter,
         address: FunctionAddress? = nil
     ) throws {
-        // TODO: Stack overflow check can be done at the entry of expression
-        guard (frames.count + numberOfValues) < limit else {
+        guard frames.count < limit, (numberOfValues + iseq.maxStackHeight + (defaultLocals?.count ?? 0)) < limit else {
             throw Trap.callStackExhausted
         }
         let valueFrameIndex = self.numberOfValues - argc
@@ -63,6 +62,11 @@ struct Stack {
         self.valueStack.copyValues(copyCount: copyCount, popCount: popCount)
     }
 
+    func deallocate() {
+        self.valueStack.deallocate()
+        self.frames.deallocate()
+    }
+
     var topValue: Value {
         self.valueStack.topValue
     }
@@ -204,6 +208,10 @@ struct FixedSizeStack<Element> {
     subscript(_ index: Int) -> Element {
         self.buffer[index]
     }
+
+    func deallocate() {
+        self.buffer.deallocate()
+    }
 }
 
 extension FixedSizeStack: Sequence {
diff --git a/Sources/WasmKit/ModuleParser.swift b/Sources/WasmKit/ModuleParser.swift
@@ -125,45 +125,25 @@ func parseModule<Stream: ByteStream>(stream: Stream, features: WasmFeatureSet =
         memoryTypes: module.memories.map { $0.type },
         tables: module.tables
     )
+    let allocator = module.allocator
     let functions = codes.enumerated().map { [hasDataCount = parser.hasDataCount, features] index, code in
         let funcTypeIndex = typeIndices[index]
         let funcType = module.types[Int(funcTypeIndex)]
         return GuestFunction(
-            type: typeIndices[index], locals: code.locals,
+            type: typeIndices[index], locals: code.locals, allocator: allocator,
             body: {
-                let enableAssert: Bool
-                #if ASSERT
-                enableAssert = true
-                #else
-                enableAssert = false
-                #endif
-                
                 var translator = InstructionTranslator(
-                    allocator: module.allocator,
+                    allocator: allocator,
                     module: translatorContext,
                     type: funcType, locals: code.locals
                 )
 
-                if enableAssert && !_isFastAssertConfiguration() {
-                    let globalFuncIndex = module.imports.count + index
-                    print("🚀 Starting Translation for code[\(globalFuncIndex)] (\(funcType))")
-                    var tracing = InstructionTracingVisitor(trace: {
-                        print("🍵 code[\(globalFuncIndex)] Translating \($0)")
-                    }, visitor: translator)
-                    try WasmParser.parseExpression(
-                        bytes: Array(code.expression),
-                        features: features, hasDataCount: hasDataCount,
-                        visitor: &tracing
-                    )
-                    let newISeq = InstructionSequence(instructions: tracing.visitor.finalize())
-                    return newISeq
-                }
                 try WasmParser.parseExpression(
                     bytes: Array(code.expression),
                     features: features, hasDataCount: hasDataCount,
                     visitor: &translator
                 )
-                return InstructionSequence(instructions: translator.finalize())
+                return translator.finalize()
             })
     }
     module.functions = functions
diff --git a/Sources/WasmKit/Translator.swift b/Sources/WasmKit/Translator.swift
@@ -11,6 +11,22 @@ class ISeqAllocator {
         return buffer
     }
 
+    func allocateDefaultLocals(_ locals: [ValueType]) -> UnsafeBufferPointer<Value> {
+        let buffer = UnsafeMutableBufferPointer<Value>.allocate(capacity: locals.count)
+        for (index, localType) in locals.enumerated() {
+            buffer[index] = localType.defaultValue
+        }
+        self.buffers.append(UnsafeMutableRawBufferPointer(buffer))
+        return UnsafeBufferPointer(buffer)
+    }
+
+    func allocateInstructions(capacity: Int) -> UnsafeMutableBufferPointer<Instruction> {
+        assert(_isPOD(Instruction.self), "Instruction must be POD")
+        let buffer = UnsafeMutableBufferPointer<Instruction>.allocate(capacity: capacity)
+        self.buffers.append(UnsafeMutableRawBufferPointer(buffer))
+        return buffer
+    }
+
     deinit {
         for buffer in buffers {
             buffer.deallocate()
@@ -145,12 +161,16 @@ struct InstructionTranslator: InstructionVisitor {
     }
     struct ValueStack {
         private var values: [MetaValue] = []
+        /// The maximum height of the stack within the function
+        private(set) var maxHeight: Int = 0
         var height: Int { values.count }
 
         mutating func push(_ value: ValueType) {
-            self.values.append(.some(value))
+            push(.some(value))
         }
         mutating func push(_ value: MetaValue) {
+            // Record the maximum height of the stack we have seen
+            maxHeight = max(maxHeight, height)
             self.values.append(value)
         }
 
@@ -393,13 +413,20 @@ struct InstructionTranslator: InstructionVisitor {
         valueStack.truncate(height: currentFrame.stackHeight)
     }
 
-    public mutating func finalize() -> [Instruction] {
+    public mutating func finalize() -> InstructionSequence {
         iseqBuilder.pinLabelHere(self.endOfFunctionLabel)
         #if DEBUG
         // Check dangling labels
         iseqBuilder.assertDanglingLabels()
         #endif
-        return iseqBuilder.finalize()
+        let instructions = iseqBuilder.finalize()
+        // TODO: Figure out a way to avoid the copy here while keeping the execution performance.
+        let buffer = allocator.allocateInstructions(capacity: instructions.count + 1)
+        for (idx, instruction) in instructions.enumerated() {
+            buffer[idx] = instruction
+        }
+        buffer[instructions.count] = .endOfFunction
+        return InstructionSequence(instructions: UnsafeBufferPointer(buffer), maxStackHeight: valueStack.maxHeight)
     }
 
     // MARK: - Visitor
diff --git a/Sources/WasmKit/Types/Module.swift b/Sources/WasmKit/Types/Module.swift
@@ -80,14 +80,14 @@ typealias LabelIndex = UInt32
 /// > Note:
 /// <https://webassembly.github.io/spec/core/syntax/modules.html#functions>
 struct GuestFunction {
-    init(type: TypeIndex, locals: [WasmParser.ValueType], body: @escaping () throws -> InstructionSequence) {
+    init(
+        type: TypeIndex,
+        locals: [WasmParser.ValueType],
+        allocator: ISeqAllocator,
+        body: @escaping () throws -> InstructionSequence
+    ) {
         self.type = type
-        // TODO: Deallocate const default locals after the module is deallocated
-        let defaultLocals = UnsafeMutableBufferPointer<Value>.allocate(capacity: locals.count)
-        for (index, localType) in locals.enumerated() {
-            defaultLocals[index] = localType.defaultValue
-        }
-        self.defaultLocals = UnsafeBufferPointer(defaultLocals)
+        self.defaultLocals = allocator.allocateDefaultLocals(locals)
         self.materializer = body
     }
 
diff --git a/Tests/WasmKitTests/Execution/HostModuleTests.swift b/Tests/WasmKitTests/Execution/HostModuleTests.swift
@@ -23,28 +23,49 @@ final class HostModuleTests: XCTestCase {
     func testReentrancy() throws {
         let runtime = Runtime()
         let voidSignature = WasmParser.FunctionType(parameters: [], results: [])
+        let allocator = ISeqAllocator()
+        func compile(_ instructions: [WasmKit.Instruction], maxStackHeight: Int) -> InstructionSequence {
+            let buffer = allocator.allocateInstructions(capacity: instructions.count + 1)
+            for (i, instruction) in instructions.enumerated() {
+                buffer[i] = instruction
+            }
+            buffer[instructions.count] = .endOfFunction
+            return InstructionSequence(
+                instructions: UnsafeBufferPointer(buffer), maxStackHeight: maxStackHeight
+            )
+        }
         let module = Module(
             types: [voidSignature],
             functions: [
                 // [0] (import "env" "bar" func)
                 // [1] (import "env" "qux" func)
                 // [2] "foo"
                 GuestFunction(
-                    type: 0, locals: [],
+                    type: 0,
+                    locals: [],
+                    allocator: allocator,
                     body: {
-                        [
-                            .call(functionIndex: 0),
-                            .call(functionIndex: 0),
-                            .call(functionIndex: 0),
-                        ]
+                        compile(
+                            [
+                                .call(functionIndex: 0),
+                                .call(functionIndex: 0),
+                                .call(functionIndex: 0),
+                            ],
+                            maxStackHeight: 0
+                        )
                     }),
                 // [3] "bar"
                 GuestFunction(
-                    type: 0, locals: [],
+                    type: 0,
+                    locals: [],
+                    allocator: allocator,
                     body: {
-                        [
-                            .call(functionIndex: 1)
-                        ]
+                        compile(
+                            [
+                                .call(functionIndex: 1)
+                            ],
+                            maxStackHeight: 0
+                        )
                     }),
             ],
             imports: [

Original file line number	Diff line number	Diff line change
`@@ -23,8 +23,7 @@ struct Stack {`
`23`	`23`	`returnPC: ProgramCounter,`
`24`	`24`	`address: FunctionAddress? = nil`
`25`	`25`	`) throws {`
`26`		`- // TODO: Stack overflow check can be done at the entry of expression`
`27`		`- guard (frames.count + numberOfValues) < limit else {`
	`26`	`+ guard frames.count < limit, (numberOfValues + iseq.maxStackHeight + (defaultLocals?.count ?? 0)) < limit else {`
`28`	`27`	`throw Trap.callStackExhausted`
`29`	`28`	`}`
`30`	`29`	`let valueFrameIndex = self.numberOfValues - argc`
`@@ -63,6 +62,11 @@ struct Stack {`
`63`	`62`	`self.valueStack.copyValues(copyCount: copyCount, popCount: popCount)`
`64`	`63`	`}`
`65`	`64`
	`65`	`+ func deallocate() {`
	`66`	`+ self.valueStack.deallocate()`
	`67`	`+ self.frames.deallocate()`
	`68`	`+ }`
	`69`	`+`
`66`	`70`	`var topValue: Value {`
`67`	`71`	`self.valueStack.topValue`
`68`	`72`	`}`
`@@ -204,6 +208,10 @@ struct FixedSizeStack<Element> {`
`204`	`208`	`subscript(_ index: Int) -> Element {`
`205`	`209`	`self.buffer[index]`
`206`	`210`	`}`
	`211`	`+`
	`212`	`+ func deallocate() {`
	`213`	`+ self.buffer.deallocate()`
	`214`	`+ }`
`207`	`215`	`}`
`208`	`216`
`209`	`217`	`extension FixedSizeStack: Sequence {`