From d1749131c5795524753161d06244a196909068fd Mon Sep 17 00:00:00 2001
From: Stephen Chung <schungx@live.com>
Date: Sat, 8 Jan 2022 23:23:43 +0800
Subject: [PATCH] Add missing data size check.

---
 src/eval/stmt.rs   |  3 ++-
 tests/data_size.rs | 59 ++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 59 insertions(+), 3 deletions(-)

diff --git a/src/eval/stmt.rs b/src/eval/stmt.rs
index 0502c996..5994236e 100644
--- a/src/eval/stmt.rs
+++ b/src/eval/stmt.rs
@@ -141,6 +141,7 @@ impl Engine {
             let args = &mut [lhs_ptr_inner, &mut new_val];
 
             match self.call_native_fn(global, state, lib, op, hash, args, true, true, op_pos) {
+                Ok(_) => self.check_data_size(&mut args[0], root.1)?,
                 Err(err) if matches!(*err, ERR::ErrorFunctionNotFound(ref f, _) if f.starts_with(op)) =>
                 {
                     // Expand to `var = var op rhs`
@@ -153,7 +154,7 @@ impl Engine {
 
                     *args[0] = value.flatten();
                 }
-                err => return err.map(|_| ()),
+                Err(err) => return Err(err),
             }
         } else {
             // Normal assignment
diff --git a/tests/data_size.rs b/tests/data_size.rs
index a4dc1818..efe6e12c 100644
--- a/tests/data_size.rs
+++ b/tests/data_size.rs
@@ -1,5 +1,5 @@
 #![cfg(not(feature = "unchecked"))]
-use rhai::{Engine, EvalAltResult, ParseErrorType};
+use rhai::{Engine, EvalAltResult, ParseErrorType, INT};
 
 #[cfg(not(feature = "no_index"))]
 use rhai::Array;
@@ -101,6 +101,39 @@ fn test_max_array_size() -> Result<(), Box<EvalAltResult>> {
         EvalAltResult::ErrorDataTooLarge(_, _)
     ));
 
+    #[cfg(not(feature = "no_closure"))]
+    assert_eq!(
+        engine.eval::<INT>(
+            "
+                    let x = 42;
+                    let y = [];
+                    let f = || x;
+                    for n in 0..10 {
+                        y += x;
+                    }
+                    len(y)
+                "
+        )?,
+        10
+    );
+
+    #[cfg(not(feature = "no_closure"))]
+    assert!(matches!(
+        *engine
+            .run(
+                "
+                    let x = 42;
+                    let y = [];
+                    let f = || x;
+                    for n in 0..11 {
+                        y += x;
+                    }
+                "
+            )
+            .expect_err("should error"),
+        EvalAltResult::ErrorDataTooLarge(_, _)
+    ));
+
     assert!(matches!(
         *engine
             .run(
@@ -113,13 +146,25 @@ fn test_max_array_size() -> Result<(), Box<EvalAltResult>> {
         EvalAltResult::ErrorDataTooLarge(_, _)
     ));
 
+    #[cfg(not(feature = "no_object"))]
+    assert_eq!(
+        engine.eval::<INT>(
+            "
+                let x = [1,2,3,4,5,6];
+                x.pad(10, 42);
+                len(x)
+            "
+        )?,
+        10
+    );
+
     #[cfg(not(feature = "no_object"))]
     assert!(matches!(
         *engine
             .run(
                 "
                     let x = [1,2,3,4,5,6];
-                    x.pad(100, 42);
+                    x.pad(11, 42);
                     x
                 "
             )
@@ -127,6 +172,16 @@ fn test_max_array_size() -> Result<(), Box<EvalAltResult>> {
         EvalAltResult::ErrorDataTooLarge(_, _)
     ));
 
+    assert_eq!(
+        engine.eval::<INT>(
+            "
+                let x = [1,2,3];
+                len([x, x, x])
+            "
+        )?,
+        3
+    );
+
     assert!(matches!(
         *engine
             .run(