From efad44e7a6438a4c6bf2d5de45bf78da3de27219 Mon Sep 17 00:00:00 2001
From: David Hewitt <david.hewitt@pydantic.dev>
Date: Tue, 2 Apr 2024 18:47:22 +0100
Subject: [PATCH] fix hang when parsing egregiously long integer (#77)

---
 Cargo.lock            |  2 +-
 Cargo.toml            |  2 +-
 src/number_decoder.rs |  6 +++---
 tests/main.rs         | 29 ++++++++++++-----------------
 4 files changed, 17 insertions(+), 22 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 908feb45..c7094227 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -111,7 +111,7 @@ checksum = "b1a46d1a171d865aa5f83f92695765caa047a9b4cbae2cbf37dbd613a793fd4c"
 
 [[package]]
 name = "jiter"
-version = "0.2.0"
+version = "0.2.1"
 dependencies = [
  "ahash",
  "bencher",
diff --git a/Cargo.toml b/Cargo.toml
index 46fb868f..4c3f3a05 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "jiter"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2021"
 description = "Iterable JSON parser"
 readme = "README.md"
diff --git a/src/number_decoder.rs b/src/number_decoder.rs
index a3fc7882..c9668711 100644
--- a/src/number_decoder.rs
+++ b/src/number_decoder.rs
@@ -224,6 +224,9 @@ impl IntParse {
         // number is too big for i64, we need ot use a big int
         loop {
             let (chunk, new_index) = IntChunk::parse_big(data, index);
+            if (new_index - start) > 4300 {
+                return json_err!(NumberOutOfRange, start + 4301);
+            }
             match chunk {
                 IntChunk::Ongoing(value) => {
                     big_value *= ONGOING_CHUNK_MULTIPLIER;
@@ -231,9 +234,6 @@ impl IntParse {
                     index = new_index;
                 }
                 IntChunk::Done(value) => {
-                    if (new_index - start) > 4300 {
-                        return json_err!(NumberOutOfRange, start + 4301);
-                    }
                     big_value *= POW_10[new_index - index];
                     big_value += value;
                     if !positive {
diff --git a/tests/main.rs b/tests/main.rs
index 83ce3e00..636550e1 100644
--- a/tests/main.rs
+++ b/tests/main.rs
@@ -981,23 +981,18 @@ fn test_4300_int() {
 }
 
 #[test]
-fn test_4302_int_err() {
-    let json = (0..4302).map(|_| "9".to_string()).collect::<Vec<_>>().join("");
-    let bytes = json.as_bytes();
-    let e = JsonValue::parse(bytes, false).unwrap_err();
-    assert_eq!(e.error_type, JsonErrorType::NumberOutOfRange);
-    assert_eq!(e.index, 4301);
-    assert_eq!(e.description(bytes), "number out of range at line 1 column 4302");
-}
-
-#[test]
-fn test_5000_int_err() {
-    let json = ["9"; 5000].join("");
-    let bytes = json.as_bytes();
-    let e = JsonValue::parse(bytes, false).unwrap_err();
-    assert_eq!(e.error_type, JsonErrorType::NumberOutOfRange);
-    assert_eq!(e.index, 4301);
-    assert_eq!(e.description(bytes), "number out of range at line 1 column 4302");
+fn test_big_int_errs() {
+    for json in [
+        &[b'9'; 4302][..],
+        &[b'9'; 5900][..],
+        // If the check is only done at the end, this will hang
+        &[b'9'; 10usize.pow(7)][..],
+    ] {
+        let e = JsonValue::parse(json, false).unwrap_err();
+        assert_eq!(e.error_type, JsonErrorType::NumberOutOfRange);
+        assert_eq!(e.index, 4301);
+        assert_eq!(e.description(json), "number out of range at line 1 column 4302");
+    }
 }
 
 #[test]