From f427706e06fa94067028e0319985418176be3179 Mon Sep 17 00:00:00 2001
From: Benoit Blanchon <git@benoitblanchon.fr>
Date: Thu, 20 Jul 2023 19:04:40 +0200
Subject: [PATCH] VariantPoolList: handle `SlotId` overflow

---
 extras/tests/ResourceManager/CMakeLists.txt   |  5 ++++
 extras/tests/ResourceManager/allocVariant.cpp | 26 +++++++++++++++++++
 src/ArduinoJson/Memory/VariantPoolList.hpp    | 11 +++++++-
 3 files changed, 41 insertions(+), 1 deletion(-)

diff --git a/extras/tests/ResourceManager/CMakeLists.txt b/extras/tests/ResourceManager/CMakeLists.txt
index 214d5b82..a8a7a4be 100644
--- a/extras/tests/ResourceManager/CMakeLists.txt
+++ b/extras/tests/ResourceManager/CMakeLists.txt
@@ -10,6 +10,11 @@ add_executable(ResourceManagerTests
 	StringBuilder.cpp
 )
 
+add_compile_definitions(ResourceManagerTests
+	ARDUINOJSON_SLOT_ID_SIZE=1 # require less RAM for overflow tests
+	ARDUINOJSON_POOL_CAPACITY=16
+)
+
 add_test(ResourceManager ResourceManagerTests)
 
 set_tests_properties(ResourceManager
diff --git a/extras/tests/ResourceManager/allocVariant.cpp b/extras/tests/ResourceManager/allocVariant.cpp
index fae36c6c..5bf734e5 100644
--- a/extras/tests/ResourceManager/allocVariant.cpp
+++ b/extras/tests/ResourceManager/allocVariant.cpp
@@ -66,4 +66,30 @@ TEST_CASE("ResourceManager::allocSlot()") {
     REQUIRE(variant.id() == NULL_SLOT);
     REQUIRE(static_cast<VariantSlot*>(variant) == nullptr);
   }
+
+  SECTION("Try overflow pool counter") {
+    ResourceManager resources;
+
+    // this test assumes SlotId is 8-bit; otherwise it consumes a lot of memory
+    // tyhe GitHub Workflow gets killed
+    REQUIRE(NULL_SLOT == 255);
+
+    // fill all the pools
+    for (SlotId i = 0; i < NULL_SLOT; i++) {
+      auto slot = resources.allocSlot();
+      REQUIRE(slot.id() == i);  // or != NULL_SLOT
+      REQUIRE(static_cast<VariantSlot*>(slot) != nullptr);
+    }
+
+    REQUIRE(resources.overflowed() == false);
+
+    // now all allocations should fail
+    for (int i = 0; i < 10; i++) {
+      auto slot = resources.allocSlot();
+      REQUIRE(slot.id() == NULL_SLOT);
+      REQUIRE(static_cast<VariantSlot*>(slot) == nullptr);
+    }
+
+    REQUIRE(resources.overflowed() == true);
+  }
 }
diff --git a/src/ArduinoJson/Memory/VariantPoolList.hpp b/src/ArduinoJson/Memory/VariantPoolList.hpp
index f7f145d9..71ed152f 100644
--- a/src/ArduinoJson/Memory/VariantPoolList.hpp
+++ b/src/ArduinoJson/Memory/VariantPoolList.hpp
@@ -108,11 +108,16 @@ class VariantPoolList {
     if (count_ == capacity_ && !increaseCapacity(allocator))
       return nullptr;
     auto pool = &pools_[count_++];
-    pool->create(ARDUINOJSON_POOL_CAPACITY, allocator);
+    SlotCount poolCapacity = ARDUINOJSON_POOL_CAPACITY;
+    if (count_ == maxPools)  // last pool is smaller because of NULL_SLOT
+      poolCapacity--;
+    pool->create(poolCapacity, allocator);
     return pool;
   }
 
   bool increaseCapacity(Allocator* allocator) {
+    if (count_ == maxPools)
+      return false;
     void* newPools;
     PoolCount newCapacity;
     if (pools_) {
@@ -134,6 +139,10 @@ class VariantPoolList {
   PoolCount count_ = 0;
   PoolCount capacity_ = 0;
   SlotId freeList_ = NULL_SLOT;
+
+ public:
+  static const PoolCount maxPools =
+      PoolCount(NULL_SLOT / ARDUINOJSON_POOL_CAPACITY + 1);
 };
 
 ARDUINOJSON_END_PRIVATE_NAMESPACE