From e64ac8c17732bdb97689644b43d07f61d3f6cab2 Mon Sep 17 00:00:00 2001
From: Diego <96022404+dzfrias@users.noreply.github.com>
Date: Mon, 10 Jun 2024 12:42:13 -0700
Subject: [PATCH] LibWasm: Validate number of data sections

---
 Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp | 6 ++++++
 Userland/Libraries/LibWasm/AbstractMachine/Validator.h   | 1 +
 2 files changed, 7 insertions(+)

diff --git a/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp b/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp
index 58b72a0abcd..6102bd1a7a1 100644
--- a/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp
+++ b/Userland/Libraries/LibWasm/AbstractMachine/Validator.cpp
@@ -41,6 +41,10 @@ ErrorOr<void, ValidationError> Validator::validate(Module& module)
         m_context.types.extend(section.types());
     });
 
+    module.for_each_section_of_type<DataCountSection>([this](DataCountSection const& section) {
+        m_context.data_count = section.count();
+    });
+
     module.for_each_section_of_type<ImportSection>([&](ImportSection const& section) {
         for (auto& import_ : section.imports()) {
             import_.description().visit(
@@ -177,6 +181,8 @@ ErrorOr<void, ValidationError> Validator::validate(StartSection const& section)
 
 ErrorOr<void, ValidationError> Validator::validate(DataSection const& section)
 {
+    if (m_context.data_count.has_value() && section.data().size() != m_context.data_count)
+        return Errors::invalid("data count does not match segment count"sv);
     for (auto& entry : section.data()) {
         TRY(entry.value().visit(
             [](DataSection::Data::Passive const&) { return ErrorOr<void, ValidationError> {}; },
diff --git a/Userland/Libraries/LibWasm/AbstractMachine/Validator.h b/Userland/Libraries/LibWasm/AbstractMachine/Validator.h
index 2cd96730d6c..e1c7e3bc1a5 100644
--- a/Userland/Libraries/LibWasm/AbstractMachine/Validator.h
+++ b/Userland/Libraries/LibWasm/AbstractMachine/Validator.h
@@ -27,6 +27,7 @@ struct Context {
     COWVector<bool> datas;
     COWVector<ValueType> locals;
     COWVector<ResultType> labels;
+    Optional<u32> data_count;
     Optional<ResultType> return_;
     AK::HashTable<FunctionIndex> references;
     size_t imported_function_count { 0 };