Fix two fuzz bugs with ArrayNewSeg (#5242)

First, we forgot to note the type annotation on `ArrayNewSeg` instructions, so
in small modules where these are the only annotated instructions, the type
section would be incomplete.

Second, in the interpreter we were reserving space for the array before checking
that the segment access was valid. This could cause huge allocations that threw
bad_alloc exceptions before the interpreter could get around to trapping. Fix
the problem by reserving the array after validating the arguements.

Fixes #5236.

6 files changed, 63 insertions, 23 deletions

diff --git a/src/ir/module-utils.cpp b/src/ir/module-utils.cpp
index 744f62682..c10a45b15 100644
--- a/src/ir/module-utils.cpp
+++ b/src/ir/module-utils.cpp
@@ -66,6 +66,8 @@ struct CodeScanner
       counts.note(curr->type);
     } else if (curr->is<ArrayNew>()) {
       counts.note(curr->type);
+    } else if (curr->is<ArrayNewSeg>()) {
+      counts.note(curr->type);
     } else if (curr->is<ArrayInit>()) {
       counts.note(curr->type);
     } else if (auto* cast = curr->dynCast<RefCast>()) {
diff --git a/src/wasm-interpreter.h b/src/wasm-interpreter.h
index 06fee0317..165bfd18a 100644
--- a/src/wasm-interpreter.h
+++ b/src/wasm-interpreter.h
@@ -3519,7 +3519,6 @@ public:
     WASM_UNUSED(elemType);
 
     Literals contents;
-    contents.reserve(size);
 
     switch (curr->op) {
       case NewData: {
@@ -3532,6 +3531,7 @@ public:
             end > seg.data.size()) {
           trap("out of bounds segment access in array.new_data");
         }
+        contents.reserve(size);
         for (Index i = offset; i < end; i += elemBytes) {
           auto addr = (void*)&seg.data[i];
           contents.push_back(Literal::makeFromMemory(addr, element));
@@ -3546,6 +3546,7 @@ public:
         if (end > seg.data.size()) {
           trap("out of bounds segment access in array.new_elem");
         }
+        contents.reserve(size);
         for (Index i = offset; i < end; ++i) {
           auto val = self()->visit(seg.data[i]).getSingleValue();
           contents.push_back(val);
diff --git a/test/lit/array-new-seg-note-count.wast b/test/lit/array-new-seg-note-count.wast
new file mode 100644
index 000000000..45c08e313
--- /dev/null
+++ b/test/lit/array-new-seg-note-count.wast
@@ -0,0 +1,25 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-opt %s -all --roundtrip -S -o - | filecheck %s
+
+;; Test that the array type is emitted into the type section properly.
+(module
+ ;; CHECK:      (type $vec (array i32))
+ (type $vec (array i32))
+ ;; CHECK:      (type $none_=>_ref|$vec| (func (result (ref $vec))))
+
+ ;; CHECK:      (data "")
+ (data "")
+ ;; CHECK:      (func $test (result (ref $vec))
+ ;; CHECK-NEXT:  (array.new_data $vec 0
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:   (i32.const 0)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $test (result (ref $vec))
+  (array.new_data $vec 0
+   (i32.const 0)
+   (i32.const 0)
+  )
+ )
+)
diff --git a/test/lit/arrays.wast b/test/lit/arrays.wast
index 4a262ef13..bccbaeb09 100644
--- a/test/lit/arrays.wast
+++ b/test/lit/arrays.wast
@@ -12,6 +12,15 @@
 (module
  ;; CHECK:      (type $arrayref_=>_i32 (func (param arrayref) (result i32)))
 
+ ;; CHECK:      (type $byte-array (array (mut i8)))
+ ;; ROUNDTRIP:      (type $arrayref_=>_i32 (func (param arrayref) (result i32)))
+
+ ;; ROUNDTRIP:      (type $byte-array (array (mut i8)))
+ (type $byte-array (array (mut i8)))
+ ;; CHECK:      (type $func-array (array (mut funcref)))
+ ;; ROUNDTRIP:      (type $func-array (array (mut funcref)))
+ (type $func-array (array (mut funcref)))
+
  ;; CHECK:      (type $ref|array|_=>_i32 (func (param (ref array)) (result i32)))
 
  ;; CHECK:      (type $nullref_=>_i32 (func (param nullref) (result i32)))
@@ -20,9 +29,7 @@
 
  ;; CHECK:      (type $none_=>_ref|$func-array| (func (result (ref $func-array))))
 
- ;; CHECK:      (type $byte-array (array (mut i8)))
- ;; ROUNDTRIP:      (type $arrayref_=>_i32 (func (param arrayref) (result i32)))
-
+ ;; CHECK:      (data "hello")
  ;; ROUNDTRIP:      (type $ref|array|_=>_i32 (func (param (ref array)) (result i32)))
 
  ;; ROUNDTRIP:      (type $nullref_=>_i32 (func (param nullref) (result i32)))
@@ -31,13 +38,6 @@
 
  ;; ROUNDTRIP:      (type $none_=>_ref|$func-array| (func (result (ref $func-array))))
 
- ;; ROUNDTRIP:      (type $byte-array (array (mut i8)))
- (type $byte-array (array (mut i8)))
- ;; CHECK:      (type $func-array (array (mut funcref)))
- ;; ROUNDTRIP:      (type $func-array (array (mut funcref)))
- (type $func-array (array (mut funcref)))
-
- ;; CHECK:      (data "hello")
  ;; ROUNDTRIP:      (data "hello")
  (data "hello")
  ;; CHECK:      (elem func $len $impossible-len $unreachable-len)
diff --git a/test/lit/passes/memory-packing-gc.wast b/test/lit/passes/memory-packing-gc.wast
index 8a2e3a5da..d16fe3a54 100644
--- a/test/lit/passes/memory-packing-gc.wast
+++ b/test/lit/passes/memory-packing-gc.wast
@@ -5,11 +5,11 @@
 ;; RUN: foreach %s %t wasm-opt --memory-packing --all-features -S -o - | filecheck %s
 
 (module
-  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
-
   ;; CHECK:      (type $array (array i8))
   (type $array (array i8))
 
+  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
+
   ;; CHECK:      (data "hello")
   (data "hello")
 
@@ -29,11 +29,11 @@
 )
 
 (module
-  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
-
   ;; CHECK:      (type $array (array i8))
   (type $array (array i8))
 
+  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
+
   ;; CHECK:      (data "\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00hello\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00")
   (data "\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00hello\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00")
 
@@ -53,11 +53,11 @@
 )
 
 (module
-  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
-
   ;; CHECK:      (type $array (array i8))
   (type $array (array i8))
 
+  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
+
   ;; CHECK:      (memory $mem 1 1)
   (memory $mem 1 1)
 
@@ -80,11 +80,11 @@
 )
 
 (module
-  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
-
   ;; CHECK:      (type $array (array i8))
   (type $array (array i8))
 
+  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
+
   ;; CHECK:      (memory $mem 1 1)
   (memory $mem 1 1)
 
@@ -109,13 +109,13 @@
 )
 
 (module
-  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
-
   ;; CHECK:      (type $array (array i8))
   (type $array (array i8))
 
   (data (i32.const 0) "optimize\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00me")
 
+  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
+
   ;; CHECK:      (memory $mem 1 1)
 
   ;; CHECK:      (data (i32.const 0) "optimize")
@@ -143,13 +143,13 @@
 )
 
 (module
-  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
-
   ;; CHECK:      (type $array (array i8))
   (type $array (array i8))
 
   (data "dead\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00segment")
 
+  ;; CHECK:      (type $none_=>_ref|$array| (func (result (ref $array))))
+
   ;; CHECK:      (memory $mem 1 1)
 
   ;; CHECK:      (data "but not\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00\00me")
diff --git a/test/spec/array-new-data.wast b/test/spec/array-new-data.wast
index 8ef62cb87..5d49a0487 100644
--- a/test/spec/array-new-data.wast
+++ b/test/spec/array-new-data.wast
@@ -77,3 +77,15 @@
 (assert_return (invoke "get" (i32.const 0)) (i32.const 1))
 (assert_return (invoke "set_get" (i32.const 1) (i32.const 7)) (i32.const 7))
 (assert_return (invoke "len") (i32.const 3))
+
+(module
+  (type $vec (array i32))
+
+  (data "\01\00\00\00\02\00\00\00\03\00\00\00\04\00\00\00")
+
+  (func $new-huge (export "new-huge") (result (ref $vec))
+    (array.new_data $vec 0 (i32.const 1) (i32.const -1))
+  )
+)
+
+(assert_trap (invoke "new-huge") "out of bounds segment access in array.new_data")