WAT: Ban too large offset field in memarg position

kateinoigakukun · kateinoigakukun · commit 7bcf8103fcbf · 2024-10-03T12:02:28.000+09:00
diff --git a/Sources/WAT/Encoder.swift b/Sources/WAT/Encoder.swift
@@ -88,7 +88,7 @@ struct Encoder {
     }
 
     mutating func writeExpression(lexer: inout Lexer, wat: inout Wat) throws {
-        var parser = ExpressionParser<ExpressionEncoder>(lexer: lexer)
+        var parser = ExpressionParser<ExpressionEncoder>(lexer: lexer, features: wat.features)
         var exprEncoder = ExpressionEncoder()
         try parser.parse(visitor: &exprEncoder, wat: &wat)
         try exprEncoder.visitEnd()
@@ -156,7 +156,7 @@ struct ElementExprCollector: AnyInstructionVisitor {
     mutating func parse(indices: WatParser.ElementDecl.Indices, wat: inout Wat) throws {
         switch indices {
         case .elementExprList(let lexer):
-            var parser = ExpressionParser<ElementExprCollector>(lexer: lexer)
+            var parser = ExpressionParser<ElementExprCollector>(lexer: lexer, features: wat.features)
             try parser.parseElemExprList(visitor: &self, wat: &wat)
         case .functionList(let lexer):
             try self.parseFunctionList(lexer: lexer, wat: wat)
@@ -557,7 +557,7 @@ func encode(module: inout Wat) throws -> [UInt8] {
                         encoder.writeUnsignedLEB128(local.count)
                         local.type.encode(to: &encoder)
                     }
-                    let funcTypeIndex = try function.parse(visitor: &exprEncoder, wat: &module)
+                    let funcTypeIndex = try function.parse(visitor: &exprEncoder, wat: &module, features: module.features)
                     functionSection.append(UInt32(funcTypeIndex))
                     // TODO?
                     try exprEncoder.visitEnd()
diff --git a/Sources/WAT/Parser/ExpressionParser.swift b/Sources/WAT/Parser/ExpressionParser.swift
@@ -36,20 +36,24 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
     }
     var parser: Parser
     let locals: LocalsMap
+    let features: WasmFeatureSet
     private var labelStack = LabelStack()
 
     init(
         type: WatParser.FunctionType,
         locals: [WatParser.LocalDecl],
-        lexer: Lexer
+        lexer: Lexer,
+        features: WasmFeatureSet
     ) throws {
         self.parser = Parser(lexer)
         self.locals = try Self.computeLocals(type: type, locals: locals)
+        self.features = features
     }
 
-    init(lexer: Lexer) {
+    init(lexer: Lexer, features: WasmFeatureSet) {
         self.parser = Parser(lexer)
         self.locals = LocalsMap()
+        self.features = features
     }
 
     static func computeLocals(type: WatParser.FunctionType, locals: [WatParser.LocalDecl]) throws -> LocalsMap {
@@ -115,7 +119,7 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
     mutating func parseWastConstInstruction(
         visitor: inout Visitor
     ) throws -> Bool where Visitor: WastConstInstructionVisitor {
-        var wat = Wat.empty()
+        var wat = Wat.empty(features: features)
         // WAST allows extra const value instruction
         if try parser.takeParenBlockStart("ref.extern") {
             _ = try visitor.visitRefExtern(value: parser.expectUnsignedInt())
@@ -130,7 +134,7 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
     }
 
     mutating func parseConstInstruction(visitor: inout Visitor) throws -> Bool {
-        var wat = Wat.empty()
+        var wat = Wat.empty(features: features)
         if try foldedInstruction(visitor: &visitor, wat: &wat) {
             return true
         }
@@ -383,6 +387,10 @@ struct ExpressionParser<Visitor: InstructionVisitor> {
             try parser.consume()
             var subParser = Parser(String(maybeOffset.dropFirst(offsetPrefix.count)))
             offset = try subParser.expectUnsignedInt(UInt64.self)
+
+            if !features.contains(.memory64), offset > UInt32.max {
+                throw WatParserError("memory offset must be less than or equal to \(UInt32.max)", location: subParser.lexer.location())
+            }
         }
         var align: UInt32 = defaultAlign
         let alignPrefix = "align="
diff --git a/Sources/WAT/Parser/WastParser.swift b/Sources/WAT/Parser/WastParser.swift
@@ -9,9 +9,11 @@ protocol WastConstInstructionVisitor: InstructionVisitor {
 /// You can find its grammar definition in the [WebAssembly spec repository](https://github.com/WebAssembly/spec/blob/wg-1.0/interpreter/README.md#scripts)
 struct WastParser {
     var parser: Parser
+    let features: WasmFeatureSet
 
-    init(_ input: String) {
+    init(_ input: String, features: WasmFeatureSet) {
         self.parser = Parser(input)
+        self.features = features
     }
 
     mutating func nextDirective() throws -> WastDirective? {
@@ -24,7 +26,7 @@ struct WastParser {
                 let location = originalParser.lexer.location()
                 return .module(
                     ModuleDirective(
-                        source: .text(try parseWAT(&originalParser)), id: nil, location: location
+                        source: .text(try parseWAT(&originalParser, features: features)), id: nil, location: location
                     ))
             }
             throw WatParserError("unexpected wast directive token", location: parser.lexer.location())
@@ -78,7 +80,7 @@ struct WastParser {
     mutating func constExpression() throws -> [Value] {
         var values: [Value] = []
         var collector = ConstExpressionCollector(addValue: { values.append($0) })
-        var exprParser = ExpressionParser<ConstExpressionCollector>(lexer: parser.lexer)
+        var exprParser = ExpressionParser<ConstExpressionCollector>(lexer: parser.lexer, features: features)
         while try exprParser.parseWastConstInstruction(visitor: &collector) {}
         parser = exprParser.parser
         return values
@@ -87,7 +89,7 @@ struct WastParser {
     mutating func expectationValues() throws -> [WastExpectValue] {
         var values: [WastExpectValue] = []
         var collector = ConstExpressionCollector(addValue: { values.append(.value($0)) })
-        var exprParser = ExpressionParser<ConstExpressionCollector>(lexer: parser.lexer)
+        var exprParser = ExpressionParser<ConstExpressionCollector>(lexer: parser.lexer, features: features)
         while true {
             if let expectValue = try exprParser.parseWastExpectValue() {
                 values.append(expectValue)
@@ -115,7 +117,7 @@ public enum WastExecute {
             execute = .invoke(try WastInvoke.parse(wastParser: &wastParser))
         case "module":
             try wastParser.parser.consume()
-            execute = .wat(try parseWAT(&wastParser.parser))
+            execute = .wat(try parseWAT(&wastParser.parser, features: wastParser.features))
             try wastParser.parser.skipParenBlock()
         case "get":
             try wastParser.parser.consume()
@@ -224,9 +226,10 @@ public enum WastDirective {
             return .assertExhaustion(call: call, message: message)
         case "assert_unlinkable":
             try wastParser.parser.consume()
+            let features = wastParser.features
             let module = try wastParser.parens {
                 try $0.parser.expectKeyword("module")
-                let wat = try parseWAT(&$0.parser)
+                let wat = try parseWAT(&$0.parser, features: features)
                 try $0.parser.skipParenBlock()
                 return wat
             }
@@ -293,7 +296,7 @@ public enum ModuleSource {
             }
         }
 
-        let watModule = try parseWAT(&wastParser.parser)
+        let watModule = try parseWAT(&wastParser.parser, features: wastParser.features)
         try wastParser.parser.skipParenBlock()
         return .text(watModule)
     }
diff --git a/Sources/WAT/Parser/WatParser.swift b/Sources/WAT/Parser/WatParser.swift
@@ -80,12 +80,12 @@ struct WatParser {
         /// This method may modify TypesMap of the given WATModule
         ///
         /// - Returns: Type index of this function
-        func parse<V: InstructionVisitor>(visitor: inout V, wat: inout Wat) throws -> Int {
+        func parse<V: InstructionVisitor>(visitor: inout V, wat: inout Wat, features: WasmFeatureSet) throws -> Int {
             guard case let .definition(locals, body) = kind else {
                 fatalError("Imported functions cannot be parsed")
             }
             let (type, typeIndex) = try wat.types.resolve(use: typeUse)
-            var parser = try ExpressionParser<V>(type: type, locals: locals, lexer: body)
+            var parser = try ExpressionParser<V>(type: type, locals: locals, lexer: body, features: features)
             try parser.parse(visitor: &visitor, wat: &wat)
             return typeIndex
         }
diff --git a/Sources/WAT/WAT.swift b/Sources/WAT/WAT.swift
@@ -17,8 +17,8 @@ import WasmParser
 /// )
 /// """)
 /// ```
-public func wat2wasm(_ input: String) throws -> [UInt8] {
-    var wat = try parseWAT(input)
+public func wat2wasm(_ input: String, features: WasmFeatureSet = .default) throws -> [UInt8] {
+    var wat = try parseWAT(input, features: features)
     return try encode(module: &wat)
 }
 
@@ -36,10 +36,11 @@ public struct Wat {
     let imports: [Import]
     let exports: [Export]
     let customSections = [CustomSection]()
+    let features: WasmFeatureSet
 
     let parser: Parser
 
-    static func empty() -> Wat {
+    static func empty(features: WasmFeatureSet) -> Wat {
         Wat(
             types: TypesMap(),
             functionsMap: NameMapping<WatParser.FunctionDecl>(),
@@ -52,6 +53,7 @@ public struct Wat {
             start: nil,
             imports: [],
             exports: [],
+            features: features,
             parser: Parser("")
         )
     }
@@ -89,15 +91,15 @@ public struct Wat {
 ///
 /// let wasm = try wat.encode()
 /// ```
-public func parseWAT(_ input: String) throws -> Wat {
+public func parseWAT(_ input: String, features: WasmFeatureSet = .default) throws -> Wat {
     var parser = Parser(input)
     let wat: Wat
     if try parser.takeParenBlockStart("module") {
-        wat = try parseWAT(&parser)
+        wat = try parseWAT(&parser, features: features)
         try parser.skipParenBlock()
     } else {
         // The root (module) may be omitted
-        wat = try parseWAT(&parser)
+        wat = try parseWAT(&parser, features: features)
     }
     return wat
 }
@@ -106,8 +108,8 @@ public func parseWAT(_ input: String) throws -> Wat {
 public struct Wast {
     var parser: WastParser
 
-    init(_ input: String) {
-        self.parser = WastParser(input)
+    init(_ input: String, features: WasmFeatureSet) {
+        self.parser = WastParser(input, features: features)
     }
 
     /// Parses the next directive in the WAST script.
@@ -147,11 +149,11 @@ public struct Wast {
 ///     print("\(location): \(directive)")
 /// }
 /// ```
-public func parseWAST(_ input: String) throws -> Wast {
-    return Wast(input)
+public func parseWAST(_ input: String, features: WasmFeatureSet = .default) throws -> Wast {
+    return Wast(input, features: features)
 }
 
-func parseWAT(_ parser: inout Parser) throws -> Wat {
+func parseWAT(_ parser: inout Parser, features: WasmFeatureSet) throws -> Wat {
     // This parser is 2-pass: first it collects all module items and creates a mapping of names to indices.
 
     let initialParser = parser
@@ -285,6 +287,7 @@ func parseWAT(_ parser: inout Parser) throws -> Wat {
         start: startIndex,
         imports: imports,
         exports: exports,
+        features: features,
         parser: parser
     )
 }
diff --git a/Tests/WATTests/EncoderTests.swift b/Tests/WATTests/EncoderTests.swift
@@ -27,7 +27,7 @@ class EncoderTests: XCTestCase {
         }
 
         print("Checking\n  wast: \(wast.path)\n  json: \(json.path)")
-        var parser = WastParser(try String(contentsOf: wast))
+        var parser = WastParser(try String(contentsOf: wast), features: Spectest.deriveFeatureSet(wast: wast))
         var watModules: [ModuleDirective] = []
 
         while let directive = try parser.nextDirective() {
@@ -107,7 +107,6 @@ class EncoderTests: XCTestCase {
 
             var stats = CompatibilityTestStats()
             let excluded: [String] = [
-                "address.wast",
                 "align.wast",
                 "align64.wast",
                 "block.wast",
diff --git a/Tests/WATTests/ParserTests.swift b/Tests/WATTests/ParserTests.swift
@@ -5,8 +5,8 @@ import XCTest
 @testable import WAT
 
 class ParserTests: XCTestCase {
-    func parseWast(_ source: String) throws -> [WastDirective] {
-        var parser = WastParser(source)
+    func parseWast(_ source: String, features: WasmFeatureSet = .default) throws -> [WastDirective] {
+        var parser = WastParser(source, features: features)
         var directives: [WastDirective] = []
         while let directive = try parser.nextDirective() {
             directives.append(directive)
@@ -66,7 +66,7 @@ class ParserTests: XCTestCase {
                 (unknown expr)
               )
             )
-            """#)
+            """#, features: .default)
 
         while let directive = try parser.nextDirective() {
             switch directive {
@@ -185,7 +185,7 @@ class ParserTests: XCTestCase {
             totalCount += 1
             let source = try String(contentsOf: filePath)
             do {
-                _ = try parseWast(source)
+                _ = try parseWast(source, features: Spectest.deriveFeatureSet(wast: filePath))
             } catch {
                 failureCount += 1
                 XCTFail("Failed to parse \(filePath.path):\(error)")
diff --git a/Tests/WATTests/Spectest.swift b/Tests/WATTests/Spectest.swift
@@ -1,4 +1,5 @@
 import Foundation
+import WasmParser
 
 enum Spectest {
     static let rootDirectory = URL(fileURLWithPath: #filePath)
@@ -37,6 +38,14 @@ enum Spectest {
         }
     }
 
+    static func deriveFeatureSet(wast: URL) -> WasmFeatureSet {
+        var features = WasmFeatureSet.default
+        if wast.deletingLastPathComponent().path.hasSuffix("proposals/memory64") {
+            features.insert(.memory64)
+        }
+        return features
+    }
+
     struct Command: Decodable {
         enum CommandType: String, Decodable {
             case module
