[wasm] Add support for try-catch blocks with result

Change-Id: Iddb7a57b15a25cd7c4eff9b1ac1293670ddf070e
Reviewed-on: https://chrome-internal-review.googlesource.com/c/v8/fuzzilli/+/7970030
Commit-Queue: Matthias Liedtke <[email protected]>
Reviewed-by: Carl Smith <[email protected]>

Author: Liedtke

Sources/Fuzzilli/Base/ProgramBuilder.swift

@@ -3165,11 +3165,13 @@ public class ProgramBuilder {
             b.emit(WasmBranchIf(labelTypes: b.type(of: label).wasmLabelType!.parameters), withInputs: [label] + args + [condition])
         }
 
-        public func wasmBuildIfElse(_ condition: Variable, ifBody: () -> Void, elseBody: () -> Void) {
+        public func wasmBuildIfElse(_ condition: Variable, ifBody: () -> Void, elseBody: (() -> Void)? = nil) {
             b.emit(WasmBeginIf(), withInputs: [condition])
             ifBody()
-            b.emit(WasmBeginElse())
-            elseBody()
+            if let elseBody = elseBody {
+                b.emit(WasmBeginElse())
+                elseBody()
+            }
             b.emit(WasmEndIf())
         }
 
@@ -3215,6 +3217,29 @@ public class ProgramBuilder {
             b.emit(WasmEndTry())
         }
 
+        // The catchClauses expect a list of (tag, block-generator lambda).
+        // The lambda's inputs are the block label, the exception label (for rethrowing) and the
+        // tag arguments.
+        @discardableResult
+        public func wasmBuildLegacyTryWithResult(with signature: Signature, args: [Variable],
+                body: (Variable, [Variable]) -> Variable,
+                catchClauses: [(tag: Variable, body: (Variable, Variable, [Variable]) -> Variable)],
+                catchAllBody: ((Variable) -> Variable)? = nil) -> Variable {
+            assert(signature.parameters.count == args.count)
+            assert(signature.outputType != .nothing)
+            let instr = b.emit(WasmBeginTry(with: signature), withInputs: args)
+            var result = body(instr.innerOutput(0), Array(instr.innerOutputs(1...)))
+            for (tag, generator) in catchClauses {
+                let instr = b.emit(WasmBeginCatch(with: b.type(of: tag).wasmTagType!.parameters => signature.outputType), withInputs: [tag, result])
+                result = generator(instr.innerOutput(0), instr.innerOutput(1), Array(instr.innerOutputs(2...)))
+            }
+            if let catchAllBody = catchAllBody {
+                let instr = b.emit(WasmBeginCatchAll(with: signature), withInputs: [result])
+                result = catchAllBody(instr.innerOutput(0))
+            }
+            return b.emit(WasmEndTry(outputType: signature.outputType), withInputs: [result]).output
+        }
+
         public func WasmBuildLegacyCatch(tag: Variable, body: ((Variable, Variable, [Variable]) -> Void)) {
             // TODO(mliedtke): A catch block can produce a result type, however that result type
             // has to be in sync with the try result type (afaict).

Sources/Fuzzilli/CodeGen/CodeGeneratorWeights.swift

@@ -275,6 +275,7 @@ public let codeGeneratorWeights = [
     "WasmLoopGenerator":                        8,
     "WasmLoopWithSignatureGenerator":           8,
     "WasmLegacyTryCatchGenerator":              8,
+    "WasmLegacyTryCatchWithResultGenerator":    8,
     "WasmLegacyTryDelegateGenerator":           8,
     "WasmThrowGenerator":                       2,
     "WasmRethrowGenerator":                     10,

Sources/Fuzzilli/CodeGen/WasmCodeGenerators.swift

@@ -633,6 +633,28 @@ public let WasmCodeGenerators: [CodeGenerator] = [
         }
     },
 
+    RecursiveCodeGenerator("WasmLegacyTryCatchWithResultGenerator", inContext: .wasmFunction) { b in
+        let function = b.currentWasmModule.currentWasmFunction
+        // Choose a few random wasm values as arguments if available.
+        let args = (0..<Int.random(in: 0...5)).map {_ in b.findVariable {b.type(of: $0).Is(.wasmPrimitive)}}.filter {$0 != nil}.map {$0!}
+        let parameters = args.map {arg in Parameter.plain(b.type(of: arg))}
+        let tags = (0..<Int.random(in: 0...5)).map {_ in b.findVariable { b.type(of: $0).isWasmTagType }}.filter {$0 != nil}.map {$0!}
+        // Disallowing void here to simplify the logic. The WasmLegacyTryCatchGenerator generates try-catch blocks without a result type.
+        let outputType = b.randomWasmBlockOutputType(allowVoid: false)
+        let signature = parameters => outputType
+        let recursiveCallCount = 2 + tags.count
+        function.wasmBuildLegacyTryWithResult(with: signature, args: args, body: { label, args in
+            b.buildRecursive(block: 1, of: recursiveCallCount, n: 4)
+            return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+        }, catchClauses: tags.enumerated().map {i, tag in (tag, {_, _, _ in
+                b.buildRecursive(block: 2 + i, of: recursiveCallCount, n: 4)
+                return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+        })}, catchAllBody: { label in
+            b.buildRecursive(block: 2 + tags.count, of: recursiveCallCount, n: 4)
+            return b.randomVariable(ofType: outputType) ?? function.generateRandomWasmVar(ofType: outputType)
+        })
+    },
+
     RecursiveCodeGenerator("WasmLegacyTryDelegateGenerator", inContext: .wasmFunction, inputs: .required(.anyLabel)) { b, label in
         let function = b.currentWasmModule.currentWasmFunction
         // Choose a few random wasm values as arguments if available.

Sources/Fuzzilli/FuzzIL/Instruction.swift

@@ -1221,8 +1221,10 @@ extension Instruction: ProtobufConvertible {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
                     $0.returnType = ILTypeToWasmTypeEnum(op.signature.outputType)
                 }
-            case .wasmEndTry(_):
-                $0.wasmEndTry = Fuzzilli_Protobuf_WasmEndTry()
+            case .wasmEndTry(let op):
+                $0.wasmEndTry = Fuzzilli_Protobuf_WasmEndTry.with {
+                    $0.returnType = ILTypeToWasmTypeEnum(op.outputType)
+                }
             case .wasmBeginTryDelegate(let op):
                 $0.wasmBeginTryDelegate = Fuzzilli_Protobuf_WasmBeginTryDelegate.with {
                     $0.parameters = convertParametersToWasmTypeEnums(op.signature.parameters)
@@ -2020,8 +2022,8 @@ extension Instruction: ProtobufConvertible {
                 Parameter.plain(WasmTypeEnumToILType(param))
             })
             op = WasmBeginCatch(with: parameters => WasmTypeEnumToILType(p.returnType))
-        case .wasmEndTry(_):
-            op = WasmEndTry()
+        case .wasmEndTry(let p):
+            op = WasmEndTry(outputType: WasmTypeEnumToILType(p.returnType))
         case .wasmBeginTryDelegate(let p):
             let parameters: [Parameter] = p.parameters.map({ param in
                 Parameter.plain(WasmTypeEnumToILType(param))

Sources/Fuzzilli/FuzzIL/WasmOperations.swift

@@ -1054,8 +1054,11 @@ final class WasmBeginCatchAll : WasmOperation {
     init(with signature: Signature) {
         self.signature = signature
 
+        let inputTypes = signature.outputType != .nothing ? [signature.outputType] : []
         super.init(
-            innerOutputTypes: [.label()],
+            inputTypes: inputTypes,
+            outputType: .nothing,
+            innerOutputTypes: [.label(inputTypes)],
             attributes: [
                 .isBlockEnd,
                 .isBlockStart,
@@ -1074,15 +1077,17 @@ final class WasmBeginCatch : WasmOperation {
 
     init(with signature: Signature) {
         self.signature = signature
+        let labelTypes: [ILType] = signature.outputType != .nothing ? [signature.outputType] : []
+        let inputTypes: [ILType] = [.object(ofGroup: "WasmTag")] + labelTypes
         // TODO: In an ideal world, the catch would only have one label that is used both for
         // branching as well as for rethrowing the exception. However, rethrows may only use labels
         // from catch blocks and branches may use any label but need to be very precise on the type
         // of the label parameters, so typing the label would require different subtyping based on
         // the usage. For now, we just emit a label for branching and the ".exceptionLabel" for
         // rethrows.
         super.init(
-            inputTypes: [.object(ofGroup: "WasmTag")],
-            innerOutputTypes: [.label([]), .exceptionLabel] + signature.parameters.convertPlainToILTypes(),
+            inputTypes: inputTypes,
+            innerOutputTypes: [.label(labelTypes), .exceptionLabel] + signature.parameters.convertPlainToILTypes(),
             attributes: [
                 .isBlockEnd,
                 .isBlockStart,
@@ -1095,8 +1100,9 @@ final class WasmBeginCatch : WasmOperation {
 final class WasmEndTry: WasmOperation {
     override var opcode: Opcode { .wasmEndTry(self) }
 
-    init() {
-        super.init(attributes: [.isBlockEnd], requiredContext: [.wasmFunction])
+    init(outputType: ILType = .nothing) {
+        let inputTypes = outputType != .nothing ? [outputType] : []
+        super.init(inputTypes: inputTypes, outputType: outputType, attributes: [.isBlockEnd], requiredContext: [.wasmFunction])
     }
 }
 

Sources/Fuzzilli/Lifting/FuzzILLifter.swift

@@ -1003,18 +1003,20 @@ public class FuzzILLifter: Lifter {
         case .wasmBeginCatchAll(_):
             assert(instr.numOutputs == 0)
             w.decreaseIndentionLevel()
-            w.emit("WasmBeginCatchAll -> L:\(instr.innerOutput(0))")
+            let inputs = instr.inputs.map(lift).joined(separator: ", ")
+            w.emit("WasmBeginCatchAll [\(inputs)] -> L:\(instr.innerOutput(0))")
             w.increaseIndentionLevel()
         case .wasmBeginCatch(_):
             assert(instr.numOutputs == 0)
             w.decreaseIndentionLevel()
-            w.emit("WasmBeginCatch \(input(0)) -> L:\(instr.innerOutput(0)) E:\(instr.innerOutput(1)) [\(liftCallArguments(instr.innerOutputs(2...)))]")
+            w.emit("WasmBeginCatch \(input(0)) [\(instr.numInputs > 1 ? input(1) : "")] -> L:\(instr.innerOutput(0)) E:\(instr.innerOutput(1)) [\(liftCallArguments(instr.innerOutputs(2...)))]")
             w.increaseIndentionLevel()
 
-        case .wasmEndTry(_):
-            assert(instr.numOutputs == 0)
+        case .wasmEndTry(let op):
             w.decreaseIndentionLevel()
-            w.emit("WasmEndTry")
+            let outputPrefix = op.numOutputs > 0 ? "\(output()) <- " : ""
+            let inputs = instr.inputs.map(lift).joined(separator: ", ")
+            w.emit("\(outputPrefix)WasmEndTry [\(inputs)]")
 
         case .wasmThrow(_):
             w.emit("WasmThrow \(instr.inputs.map(lift).joined(separator: ", "))")

Sources/Fuzzilli/Lifting/WasmLifter.swift

@@ -1092,9 +1092,11 @@ public class WasmLifter {
                 self.tags[instr.output] = op.parameters
 
             case .wasmBeginCatch(let op):
+                assert(typer.type(of: instr.input(0)).wasmTagType!.parameters == op.signature.parameters)
                 importTagIfNeeded(tag: instr.input(0), parameters: op.signature.parameters)
 
             case .wasmThrow(let op):
+                assert(typer.type(of: instr.input(0)).wasmTagType!.parameters == op.parameters)
                 importTagIfNeeded(tag: instr.input(0), parameters: op.parameters)
 
             default:

Sources/Fuzzilli/Protobuf/operations.pb.swift

@@ -4034,6 +4034,8 @@ public struct Fuzzilli_Protobuf_WasmEndTry: Sendable {
   // `Message` and `Message+*Additions` files in the SwiftProtobuf library for
   // methods supported on all messages.
 
+  public var returnType: Fuzzilli_Protobuf_WasmILType = .consti32
+
   public var unknownFields = SwiftProtobuf.UnknownStorage()
 
   public init() {}
@@ -11891,18 +11893,31 @@ extension Fuzzilli_Protobuf_WasmBeginCatch: SwiftProtobuf.Message, SwiftProtobuf
 
 extension Fuzzilli_Protobuf_WasmEndTry: SwiftProtobuf.Message, SwiftProtobuf._MessageImplementationBase, SwiftProtobuf._ProtoNameProviding {
   public static let protoMessageName: String = _protobuf_package + ".WasmEndTry"
-  public static let _protobuf_nameMap = SwiftProtobuf._NameMap()
+  public static let _protobuf_nameMap: SwiftProtobuf._NameMap = [
+    1: .same(proto: "returnType"),
+  ]
 
   public mutating func decodeMessage<D: SwiftProtobuf.Decoder>(decoder: inout D) throws {
-    // Load everything into unknown fields
-    while try decoder.nextFieldNumber() != nil {}
+    while let fieldNumber = try decoder.nextFieldNumber() {
+      // The use of inline closures is to circumvent an issue where the compiler
+      // allocates stack space for every case branch when no optimizations are
+      // enabled. https://github.com/apple/swift-protobuf/issues/1034
+      switch fieldNumber {
+      case 1: try { try decoder.decodeSingularEnumField(value: &self.returnType) }()
+      default: break
+      }
+    }
   }
 
   public func traverse<V: SwiftProtobuf.Visitor>(visitor: inout V) throws {
+    if self.returnType != .consti32 {
+      try visitor.visitSingularEnumField(value: self.returnType, fieldNumber: 1)
+    }
     try unknownFields.traverse(visitor: &visitor)
   }
 
   public static func ==(lhs: Fuzzilli_Protobuf_WasmEndTry, rhs: Fuzzilli_Protobuf_WasmEndTry) -> Bool {
+    if lhs.returnType != rhs.returnType {return false}
     if lhs.unknownFields != rhs.unknownFields {return false}
     return true
   }

Sources/Fuzzilli/Protobuf/operations.proto

@@ -1162,6 +1162,7 @@ message WasmBeginCatch {
 }
 
 message WasmEndTry {
+  WasmILType returnType = 1;
 }
 
 message WasmBeginTryDelegate {

Tests/FuzzilliTests/WasmTests.swift

@@ -1246,6 +1246,75 @@ class WasmFoundationTests: XCTestCase {
         testForOutput(program: jsProg, runner: runner, outputString: "444\n")
     }
 
+    func testTryWithBlockParametersAndResult() throws {
+        let runner = try GetJavaScriptExecutorOrSkipTest()
+        let liveTestConfig = Configuration(logLevel: .error, enableInspection: true)
+
+        // We have to use the proper JavaScriptEnvironment here.
+        // This ensures that we use the available builtins.
+        let fuzzer = makeMockFuzzer(config: liveTestConfig, environment: JavaScriptEnvironment())
+        let b = fuzzer.makeBuilder()
+
+        let outputFunc = b.createNamedVariable(forBuiltin: "output")
+        let tagVoid = b.createWasmTag(parameterTypes: [])
+        let tagi32 = b.createWasmTag(parameterTypes: [.wasmi32])
+        let tagi32Other = b.createWasmTag(parameterTypes: [.wasmi32])
+        let module = b.buildWasmModule { wasmModule in
+            wasmModule.addWasmFunction(with: [.wasmi32] => .wasmi32) { function, args in
+                let result = function.wasmBuildLegacyTryWithResult(with: [.wasmi32] => .wasmi32, args: args, body: { label, args in
+                    function.wasmBuildIfElse(function.wasmi32EqualZero(args[0])) {
+                        function.WasmBuildThrow(tag: tagVoid, inputs: [])
+                    }
+                    function.wasmBuildIfElse(function.wasmi32CompareOp(args[0], function.consti32(1), using: .Eq)) {
+                        function.WasmBuildThrow(tag: tagi32, inputs: [function.consti32(100)])
+                    }
+                    function.wasmBuildIfElse(function.wasmi32CompareOp(args[0], function.consti32(2), using: .Eq)) {
+                        function.WasmBuildThrow(tag: tagi32Other, inputs: [function.consti32(200)])
+                    }
+                    return args[0]
+                }, catchClauses: [
+                    (tagi32, {label, exception, args in
+                        return args[0]
+                    }),
+                    (tagi32Other, {label, exception, args in
+                        let value = function.wasmi32BinOp(args[0], function.consti32(2), binOpKind: .Add)
+                        function.wasmBranch(to: label, args: [value])
+                        return function.consti32(-1)
+                    }),
+                ], catchAllBody: { _ in
+                    return function.consti32(900)
+                })
+                function.wasmReturn(result)
+            }
+        }
+        let exports = module.loadExports()
+
+        var expectedString = ""
+        // Note that in the comments below "returns" means that this will be passed on as a result
+        // to the EndTry block. At the end of the wasm function it performs a wasm return of the
+        // result of the EndTry.
+        for (input, expected) in [
+                // input 0 throws tagVoid which is not caught, so the catchAllBody returns 900.
+                (0, 900),
+                // input 1 throws tagi32(100) which is caught by a catch clause that returns the
+                // tag argument.
+                (1, 100),
+                // input 2 throws tagi32Other(200) which is caught by a catch clause that branches
+                // to the end of its block adding 2 to the tag argument.
+                (2, 202),
+                // input 3 doesn't throw anything, the try body returns the value directly, meaning
+                // the value "falls-through" from the try body to the endTry operation.
+                (3, 3)] {
+            let wasmOut = b.callMethod(module.getExportedMethod(at: 0), on: exports, withArgs: [b.loadInt(Int64(input))])
+            b.callFunction(outputFunc, withArgs: [b.callMethod("toString", on: wasmOut)])
+            expectedString += "\(expected)\n"
+        }
+
+        let prog = b.finalize()
+        let jsProg = fuzzer.lifter.lift(prog)
+        testForOutput(program: jsProg, runner: runner, outputString: expectedString)
+    }
+
     func testTryDelegate() throws {
         let runner = try GetJavaScriptExecutorOrSkipTest()
         let liveTestConfig = Configuration(logLevel: .error, enableInspection: true)