WasmParser: Move data count section validation to the validator

Author: kateinoigakukun

Sources/WasmKit/Execution/Instances.swift

@@ -82,7 +82,7 @@ struct InstanceEntity /* : ~Copyable */ {
     var exports: [String: InternalExternalValue]
     var functionRefs: Set<InternalFunction>
     var features: WasmFeatureSet
-    var hasDataCount: Bool
+    var dataCount: UInt32?
 
     static var empty: InstanceEntity {
         InstanceEntity(
@@ -96,7 +96,7 @@ struct InstanceEntity /* : ~Copyable */ {
             exports: [:],
             functionRefs: [],
             features: [],
-            hasDataCount: false
+            dataCount: nil
         )
     }
 
@@ -727,12 +727,6 @@ struct DataSegmentEntity {
     }
 }
 
-extension DataSegmentEntity: ValidatableEntity {
-    static func createOutOfBoundsError(index: Int, count: Int) -> Error {
-        Trap._raw("Data index out of bounds: \(index) (max: \(count))")
-    }
-}
-
 typealias InternalDataSegment = EntityHandle<DataSegmentEntity>
 
 /// > Note:

Sources/WasmKit/Execution/StoreAllocator.swift

@@ -452,7 +452,7 @@ extension StoreAllocator {
             exports: exports,
             functionRefs: functionRefs,
             features: module.features,
-            hasDataCount: module.hasDataCount
+            dataCount: module.dataCount
         )
         instancePointer.initialize(to: instanceEntity)
         instanceInitialized = true

Sources/WasmKit/Module.swift

@@ -63,7 +63,7 @@ public struct Module {
     let tableTypes: [TableType]
     let allocator: ISeqAllocator
     let features: WasmFeatureSet
-    let hasDataCount: Bool
+    let dataCount: UInt32?
 
     init(
         types: [FunctionType],
@@ -79,7 +79,7 @@ public struct Module {
         customSections: [CustomSection],
         allocator: ISeqAllocator,
         features: WasmFeatureSet,
-        hasDataCount: Bool
+        dataCount: UInt32?
     ) {
         self.functions = functions
         self.elements = elements
@@ -91,7 +91,7 @@ public struct Module {
         self.customSections = customSections
         self.allocator = allocator
         self.features = features
-        self.hasDataCount = hasDataCount
+        self.dataCount = dataCount
 
         var importedFunctionTypes: [TypeIndex] = []
         var globalTypes: [GlobalType] = []

Sources/WasmKit/ModuleParser.swift

@@ -123,6 +123,6 @@ func parseModule<Stream: ByteStream>(stream: Stream, features: WasmFeatureSet =
         customSections: customSections,
         allocator: allocator,
         features: features,
-        hasDataCount: dataCount != nil
+        dataCount: dataCount
     )
 }

Sources/WasmKit/Translator.swift

@@ -45,8 +45,8 @@ protocol TranslatorContext {
     func resolveCallee(_ index: FunctionIndex) -> InternalFunction?
     func isSameInstance(_ instance: InternalInstance) -> Bool
     func resolveGlobal(_ index: GlobalIndex) -> InternalGlobal?
-    func validateDataSegment(_ index: DataIndex) throws
     func validateFunctionIndex(_ index: FunctionIndex) throws
+    var dataCount: UInt32? { get }
 }
 
 extension TranslatorContext {
@@ -114,15 +114,15 @@ extension InternalInstance: TranslatorContext {
     func isSameInstance(_ instance: InternalInstance) -> Bool {
         return instance == self
     }
-    func validateDataSegment(_ index: DataIndex) throws {
-        _ = try self.dataSegments[validating: Int(index)]
-    }
     func validateFunctionIndex(_ index: FunctionIndex) throws {
         let function = try self.functions[validating: Int(index)]
         guard self.functionRefs.contains(function) else {
             throw ValidationError("Function index \(index) is not declared but referenced as a function reference")
         }
     }
+    var dataCount: UInt32? {
+        self.withValue { $0.dataCount }
+    }
 }
 
 private struct MetaProgramCounter {
@@ -2030,7 +2030,7 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
     }
 
     mutating func visitMemoryInit(dataIndex: UInt32) throws -> Output {
-        try self.module.validateDataSegment(dataIndex)
+        try self.validator.validateDataSegment(dataIndex)
         let addressType = try module.addressType(memoryIndex: 0)
         try pop3Emit((.i32, .i32, addressType)) { values, stack in
             let (size, sourceOffset, destOffset) = values
@@ -2045,7 +2045,7 @@ struct InstructionTranslator<Context: TranslatorContext>: InstructionVisitor {
         }
     }
     mutating func visitDataDrop(dataIndex: UInt32) throws -> Output {
-        try self.module.validateDataSegment(dataIndex)
+        try self.validator.validateDataSegment(dataIndex)
         emit(.memoryDataDrop(Instruction.MemoryDataDropOperand(segmentIndex: dataIndex)))
     }
     mutating func visitMemoryCopy(dstMem: UInt32, srcMem: UInt32) throws -> Output {

Sources/WasmKit/Validator.swift

@@ -54,6 +54,17 @@ struct InstructionValidator<Context: TranslatorContext> {
     func validateRefFunc(functionIndex: UInt32) throws {
         try context.validateFunctionIndex(functionIndex)
     }
+
+    func validateDataSegment(_ dataIndex: DataIndex) throws {
+        // instruction referring data segment requires data count section
+        // https://webassembly.github.io/spec/core/binary/modules.html#data-count-section
+        guard let dataCount = context.dataCount else {
+            throw ValidationError("Data count section is required but not found")
+        }
+        guard dataIndex < dataCount else {
+            throw ValidationError("Data index out of bounds: \(dataIndex) (max: \(dataCount))")
+        }
+    }
 }
 
 struct ModuleValidator {

Sources/WasmParser/WasmParser.swift

@@ -12,8 +12,6 @@ import WasmTypes
 public struct Parser<Stream: ByteStream> {
     @usableFromInline
     let stream: Stream
-    public private(set) var hasDataCount: Bool = false
-    public let features: WasmFeatureSet
     let limits: ParsingLimits
     var orderTracking = OrderTracking()
 
@@ -23,14 +21,14 @@ public struct Parser<Stream: ByteStream> {
     }
     private var nextParseTarget: NextParseTarget
 
+    public let features: WasmFeatureSet
     public var offset: Int {
         return stream.currentIndex
     }
 
-    public init(stream: Stream, features: WasmFeatureSet = .default, hasDataCount: Bool = false) {
+    public init(stream: Stream, features: WasmFeatureSet = .default) {
         self.stream = stream
         self.features = features
-        self.hasDataCount = hasDataCount
         self.nextParseTarget = .header
         self.limits = .default
     }
@@ -121,7 +119,7 @@ extension Code {
     /// ````
     @inlinable
     public func parseExpression<V: InstructionVisitor>(visitor: inout V) throws {
-        let parser = Parser(stream: StaticByteStream(bytes: self.expression), features: self.features, hasDataCount: self.hasDataCount)
+        let parser = Parser(stream: StaticByteStream(bytes: self.expression), features: self.features)
         var lastCode: InstructionCode?
         while try !parser.stream.hasReachedEnd() {
             lastCode = try parser.parseInstruction(visitor: &visitor)
@@ -153,8 +151,7 @@ public struct ExpressionParser {
     public init(code: Code) {
         self.parser = Parser(
             stream: StaticByteStream(bytes: code.expression),
-            features: code.features,
-            hasDataCount: code.hasDataCount
+            features: code.features
         )
         self.codeOffset = code.offset
         self.initialStreamOffset = self.parser.offset
@@ -166,7 +163,7 @@ public struct ExpressionParser {
         let shouldContinue = try !parser.stream.hasReachedEnd()
         if !shouldContinue {
             guard lastCode == .end else {
-                throw parser.makeError(.endOpcodeExpected)
+                throw WasmParserError(.endOpcodeExpected, offset: offset)
             }
         }
         return shouldContinue
@@ -212,6 +209,7 @@ public struct WasmParserError: Swift.Error {
     let message: Message
     let offset: Int
 
+    @usableFromInline
     init(_ message: Message, offset: Int) {
         self.message = message
         self.offset = offset
@@ -300,8 +298,6 @@ extension WasmParserError.Message {
 
     static let sectionOutOfOrder = Self("Sections in the module are out of order")
 
-    @usableFromInline static let dataCountSectionRequired = Self("Data count section is required but not found")
-
     static func malformedLimit(_ byte: UInt8) -> Self {
         Self("Malformed limit: \(byte)")
     }
@@ -874,24 +870,13 @@ extension Parser {
             case 7: return try v.visitConversion(.i64TruncSatF64U)
             case 8:
                 let dataIndex: DataIndex = try parseUnsigned()
-                // memory.init requires data count section
-                // https://webassembly.github.io/spec/core/binary/modules.html#data-count-section
-                guard hasDataCount else {
-                    throw makeError(.dataCountSectionRequired)
-                }
-
                 let zero = try stream.consumeAny()
                 guard zero == 0x00 else {
                     throw makeError(.zeroExpected(actual: zero))
                 }
 
                 return try v.visitMemoryInit(dataIndex: dataIndex)
             case 9:
-                // memory.drop requires data count section
-                // https://webassembly.github.io/spec/core/binary/modules.html#data-count-section
-                guard hasDataCount else {
-                    throw makeError(.dataCountSectionRequired)
-                }
                 return try v.visitDataDrop(dataIndex: try parseUnsigned())
             case 10:
                 let (zero1, zero2) = try (stream.consumeAny(), stream.consumeAny())
@@ -1140,7 +1125,7 @@ extension Parser {
             )
             return Code(
                 locals: locals, expression: expressionBytes,
-                offset: expressionStart, hasDataCount: hasDataCount, features: features
+                offset: expressionStart, features: features
             )
         }
     }
@@ -1327,7 +1312,6 @@ extension Parser {
                 payload = .dataSection(try parseDataSection())
             case 12:
                 order = .dataCount
-                hasDataCount = true
                 payload = .dataCount(try parseDataCountSection())
             default:
                 throw makeError(.malformedSectionID(sectionID))

Sources/WasmParser/WasmTypes.swift

@@ -13,8 +13,6 @@ public struct Code {
     @usableFromInline
     internal let offset: Int
     @usableFromInline
-    internal let hasDataCount: Bool
-    @usableFromInline
     internal let features: WasmFeatureSet
 }