From 38524984d5a15c433fe111b1367d74c910dbb677 Mon Sep 17 00:00:00 2001
From: "Soni L." <EnderMoneyMod@gmail.com>
Date: Mon, 23 Sep 2024 21:46:01 -0300
Subject: [PATCH] Fix handling of data count without data section (#2432)

Closes #2436
Fixes #2310
Fixes #2311
Fixes #2431
---
 src/binary-reader.cc                            | 17 +++++++++++------
 .../regress/data-count-without-data-section.txt | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 6 deletions(-)
 create mode 100644 test/regress/data-count-without-data-section.txt

diff --git a/src/binary-reader.cc b/src/binary-reader.cc
index 0fc9f9a99..b4ad5b0c4 100644
--- a/src/binary-reader.cc
+++ b/src/binary-reader.cc
@@ -202,6 +202,7 @@ class BinaryReader {
   Index num_tag_imports_ = 0;
   Index num_function_signatures_ = 0;
   Index num_function_bodies_ = 0;
+  Index num_data_segments_ = 0;
   Index data_count_ = kInvalidIndex;
 
   using ReadEndRestoreGuard =
@@ -2829,13 +2830,13 @@ Result BinaryReader::ReadCodeSection(Offset section_size) {
 
 Result BinaryReader::ReadDataSection(Offset section_size) {
   CALLBACK(BeginDataSection, section_size);
-  Index num_data_segments;
-  CHECK_RESULT(ReadCount(&num_data_segments, "data segment count"));
-  CALLBACK(OnDataSegmentCount, num_data_segments);
+  CHECK_RESULT(ReadCount(&num_data_segments_, "data segment count"));
+  CALLBACK(OnDataSegmentCount, num_data_segments_);
   // If the DataCount section is not present, then data_count_ will be invalid.
-  ERROR_UNLESS(data_count_ == kInvalidIndex || data_count_ == num_data_segments,
-               "data segment count does not equal count in DataCount section");
-  for (Index i = 0; i < num_data_segments; ++i) {
+  ERROR_UNLESS(
+      data_count_ == kInvalidIndex || data_count_ == num_data_segments_,
+      "data segment count does not equal count in DataCount section");
+  for (Index i = 0; i < num_data_segments_; ++i) {
     uint32_t flags;
     CHECK_RESULT(ReadU32Leb128(&flags, "data segment flags"));
     ERROR_IF(flags != 0 && !options_.features.bulk_memory_enabled(),
@@ -3037,6 +3038,10 @@ Result BinaryReader::ReadModule(const ReadModuleOptions& options) {
   // in case the code section was omitted.
   ERROR_UNLESS(num_function_signatures_ == num_function_bodies_,
                "function signature count != function body count");
+  // This is checked in ReadDataSection, but it must be checked at the end too,
+  // in case the data section was omitted.
+  ERROR_IF(num_data_segments_ == 0 && data_count_ != kInvalidIndex,
+           "Data section missing but DataCount non-zero");
   CALLBACK0(EndModule);
 
   return Result::Ok;
diff --git a/test/regress/data-count-without-data-section.txt b/test/regress/data-count-without-data-section.txt
new file mode 100644
index 000000000..8fa1ed69a
--- /dev/null
+++ b/test/regress/data-count-without-data-section.txt
@@ -0,0 +1,14 @@
+;;; TOOL: run-interp-spec
+(assert_malformed
+  (module binary
+    "\00asm" "\01\00\00\00"
+    "\05\03\01\00\01"          ;; Memory section with one entry
+    "\0c\01\01"                ;; Data count section with value 1
+  )
+  "data count and data section have inconsistent lengths"
+)
+(;; STDOUT ;;;
+out/test/regress/data-count-without-data-section.txt:3: assert_malformed passed:
+  0000010: error: Data section missing but DataCount non-zero
+1/1 tests passed.
+;;; STDOUT ;;)