Validation: Check stack height at the end of block frames

Author: kateinoigakukun

Sources/WasmKit/Translator.swift

@@ -310,14 +310,17 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
             var kind: Kind
             var reachable: Bool = true
 
-            var copyCount: UInt16 {
+            var copyTypes: [ValueType] {
                 switch self.kind {
                 case .block, .if:
-                    return UInt16(blockType.results.count)
+                    return blockType.results
                 case .loop:
-                    return UInt16(blockType.parameters.count)
+                    return blockType.parameters
                 }
             }
+            var copyCount: UInt16 {
+                return UInt16(copyTypes.count)
+            }
         }
 
         private var frames: [ControlFrame] = []
@@ -893,8 +896,7 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
     ///
     /// - Parameter typeHint: A type expected to be popped. Only used for diagnostic purpose.
     /// - Returns: `true` if check succeed. `false` if the pop operation is going to be performed in unreachable code path.
-    private func checkBeforePop(typeHint: ValueType?) throws -> Bool {
-        let controlFrame = try controlStack.currentFrame()
+    private func checkBeforePop(typeHint: ValueType?, controlFrame: ControlStack.ControlFrame) throws -> Bool {
         if _slowPath(valueStack.height <= controlFrame.stackHeight) {
             if controlFrame.reachable {
                 let message: String
@@ -910,6 +912,10 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         }
         return true
     }
+    private func checkBeforePop(typeHint: ValueType?) throws -> Bool {
+        let controlFrame = try controlStack.currentFrame()
+        return try self.checkBeforePop(typeHint: typeHint, controlFrame: controlFrame)
+    }
     private mutating func ensureOnVReg(_ source: ValueSource) -> VReg {
         // TODO: Copy to stack if source is on preg
         // let copyTo = valueStack.stackRegBase + VReg(valueStack.height)
@@ -959,6 +965,27 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         return try valueStack.pop()
     }
 
+    @discardableResult
+    private mutating func popPushValues(_ valueTypes: [ValueType]) throws -> Int {
+        var values: [ValueSource?] = []
+        for type in valueTypes.reversed() {
+            values.append(try popOperand(type))
+        }
+        let stackHeight = self.valueStack.height
+        for (type, value) in zip(valueTypes, values.reversed()) {
+            switch value {
+            case .local(let localIndex):
+                // Re-push local variables to the stack
+                _ = try valueStack.pushLocal(localIndex, locals: &locals)
+            case .vreg, nil:
+                _ = valueStack.push(type)
+            case .const(let index, let type):
+                valueStack.pushConst(index, type: type)
+            }
+        }
+        return stackHeight
+    }
+
     private mutating func visitReturnLike() throws {
         preserveOnStack(depth: self.type.results.count)
         for (index, resultType) in self.type.results.enumerated().reversed() {
@@ -1050,22 +1077,7 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
     mutating func visitBlock(blockType: WasmParser.BlockType) throws -> Output {
         let blockType = try module.resolveBlockType(blockType)
         let endLabel = iseqBuilder.allocLabel()
-        var parameters: [ValueSource?] = []
-        for param in blockType.parameters.reversed() {
-            parameters.append(try popOperand(param))
-        }
-        let stackHeight = self.valueStack.height
-        for (param, value) in zip(blockType.parameters, parameters.reversed()) {
-            switch value {
-            case .local(let localIndex):
-                // Re-push local variables to the stack
-                _ = try valueStack.pushLocal(localIndex, locals: &locals)
-            case .vreg, nil:
-                _ = valueStack.push(param)
-            case .const(let index, let type):
-                valueStack.pushConst(index, type: type)
-            }
-        }
+        let stackHeight = try popPushValues(blockType.parameters)
         controlStack.pushFrame(ControlStack.ControlFrame(blockType: blockType, stackHeight: stackHeight, continuation: endLabel, kind: .block))
     }
 
@@ -1158,6 +1170,10 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         if case .block(root: true) = poppedFrame.kind {
             if poppedFrame.reachable {
                 try translateReturn()
+                // TODO: Merge logic with regular block frame
+                guard valueStack.height == poppedFrame.stackHeight else {
+                    throw ValidationError("values remaining on stack at end of block")
+                }
             }
             try iseqBuilder.pinLabelHere(poppedFrame.continuation)
             return
@@ -1173,7 +1189,13 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         case .if:
             try iseqBuilder.pinLabelHere(poppedFrame.continuation)
         }
-        try valueStack.truncate(height: poppedFrame.stackHeight)
+        for result in poppedFrame.blockType.results.reversed() {
+            guard try checkBeforePop(typeHint: result, controlFrame: poppedFrame) else { continue }
+            _ = try valueStack.pop(result)
+        }
+        guard valueStack.height == poppedFrame.stackHeight else {
+            throw ValidationError("values remaining on stack at end of block")
+        }
         for result in poppedFrame.blockType.results {
             _ = valueStack.push(result)
         }
@@ -1231,6 +1253,9 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         try emitBranch(Instruction.br, relativeDepth: relativeDepth) { offset, copyCount, popCount in
             return offset
         }
+        for type in frame.copyTypes.reversed() {
+            _ = try popOperand(type)
+        }
         try markUnreachable()
     }
 
@@ -1279,12 +1304,12 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         try emitBranch(Instruction.br, relativeDepth: relativeDepth) { offset, copyCount, popCount in
             return offset
         }
+        try popPushValues(frame.copyTypes)
         try iseqBuilder.pinLabelHere(onBranchNotTaken)
     }
 
     mutating func visitBrTable(targets: WasmParser.BrTable) throws -> Output {
         guard let index = try popVRegOperand(.i32) else { return }
-        guard try controlStack.currentFrame().reachable else { return }
 
         let defaultFrame = try controlStack.branchTarget(relativeDepth: targets.defaultIndex)
 

Tests/WasmKitTests/ExtraSuite/const_slot.wast

@@ -17,6 +17,7 @@
       ;; emit its own instruction.
       (i32.const 0)
       (local.set 0)
+      (drop) ;; drop i32.add
       (local.get 0)
   )
 )

Tests/WasmKitTests/ExtraSuite/local_cow.wast

@@ -43,6 +43,7 @@
       ;; emit its own instruction.
       (local.get 1)
       (local.set 0)
+      (drop) ;; drop i32.add
       (local.get 0)
   )
 )

Tests/WasmKitTests/Spectest/TestCase.swift

@@ -143,8 +143,8 @@ extension TestCase {
                     if let result = try context.run(directive: directive) {
                         handler(self, location, result)
                     }
-                } catch let error as SpectestError {
-                    handler(self, location, .failed(error.description))
+                } catch let error {
+                    handler(self, location, .failed("\(error)"))
                 }
             }
         } catch let parseError as WatParserError {

Tests/WasmKitTests/SpectestTests.swift

@@ -22,7 +22,6 @@ final class SpectestTests: XCTestCase {
             path: Self.testPaths,
             include: [],
             exclude: [
-                "block.wast",
                 "br.wast",
                 "br_if.wast",
                 "call.wast",