Check value stack height at the entry of function

And also fix the leak of the stack and frame buffers.

Author: kateinoigakukun

Sources/WasmKit/Execution/Instructions/Expression.swift

@@ -2,15 +2,19 @@ import WasmParser
 
 struct InstructionSequence: Equatable {
     let instructions: UnsafeBufferPointer<Instruction>
+    /// The maximum height of the value stack during execution of this function.
+    /// This height does not count the locals.
+    let maxStackHeight: Int
 
-    init(instructions: [Instruction]) {
+    init(instructions: [Instruction], maxStackHeight: Int) {
         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)
+        self.maxStackHeight = maxStackHeight
     }
 
     func deallocate() {
@@ -26,12 +30,6 @@ struct InstructionSequence: Equatable {
     }
 }
 
-extension InstructionSequence: ExpressibleByArrayLiteral {
-    init(arrayLiteral elements: Instruction...) {
-        self.init(instructions: elements)
-    }
-}
-
 struct ExpressionRef: Equatable {
     let _relativeOffset: UInt32
     var relativeOffset: Int {

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))

Sources/WasmKit/Execution/Runtime/Runtime.swift

@@ -88,12 +88,15 @@ extension Runtime {
                         .tableInit(tableIndex, elementIndex),
                         .tableElementDrop(elementIndex),
                     ])
-                    let initIseq = InstructionSequence(instructions: instructions)
+                    let initIseq = InstructionSequence(instructions: instructions, maxStackHeight: 2)
                     defer { initIseq.deallocate() }
                     try evaluateConstExpr(initIseq, instance: instance)
 
                 case .declarative:
-                    let initIseq: InstructionSequence = [.tableElementDrop(elementIndex)]
+                    let initIseq = InstructionSequence(
+                        instructions: [.tableElementDrop(elementIndex)],
+                        maxStackHeight: 0
+                    )
                     defer { initIseq.deallocate() }
                     try evaluateConstExpr(initIseq, instance: instance)
 
@@ -127,7 +130,7 @@ extension Runtime {
                     .numericConst(.i32(UInt32(data.initializer.count))),
                     .memoryInit(UInt32(dataIndex)),
                     .memoryDataDrop(UInt32(dataIndex)),
-                ])
+                ], maxStackHeight: 2)
                 defer { iseq.deallocate() }
                 try evaluateConstExpr(iseq, instance: instance)
             }
@@ -141,6 +144,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,7 +193,7 @@ extension Runtime {
                 default:
                     throw InstantiationError.unsupported("init expr in global section \(global.initializer)")
                 }
-                let iseq = InstructionSequence(instructions: instructions)
+                let iseq = InstructionSequence(instructions: instructions, maxStackHeight: 1)
                 defer { iseq.deallocate() }
                 return try evaluateConstExpr(iseq, instance: globalModuleInstance, arity: 1) { _, stack in
                     return stack.popValue()
@@ -212,6 +216,7 @@ extension Runtime {
     ) throws -> T {
         try withExecution { initExecution in
             var stack = Stack()
+            defer { stack.deallocate() }
             try stack.pushFrame(
                 iseq: iseq,
                 arity: arity,

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 {

Sources/WasmKit/ModuleParser.swift

@@ -155,15 +155,15 @@ func parseModule<Stream: ByteStream>(stream: Stream, features: WasmFeatureSet =
                         features: features, hasDataCount: hasDataCount,
                         visitor: &tracing
                     )
-                    let newISeq = InstructionSequence(instructions: tracing.visitor.finalize())
+                    let newISeq = 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

Sources/WasmKit/Translator.swift

@@ -145,12 +145,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 +397,14 @@ 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()
+        return InstructionSequence(instructions: instructions, maxStackHeight: valueStack.maxHeight)
     }
 
     // MARK: - Visitor