From 5a60c55be74b377c850592b3387759d7261e57fd Mon Sep 17 00:00:00 2001
From: Benoit Blanchon <git@benoitblanchon.fr>
Date: Thu, 23 May 2024 18:36:24 +0200
Subject: [PATCH] Don't add partial objects when allocation fails

Fixes #2081
---
 CHANGELOG.md                                  |  1 +
 extras/tests/JsonArray/add.cpp                | 49 ++++++++++++++++
 extras/tests/JsonDocument/add.cpp             | 54 ++++++++++++++++--
 extras/tests/JsonObject/set.cpp               |  2 +-
 extras/tests/JsonVariant/add.cpp              | 57 +++++++++++++++++--
 src/ArduinoJson/Array/ArrayData.hpp           | 11 ++++
 src/ArduinoJson/Array/ArrayImpl.hpp           | 15 +++++
 src/ArduinoJson/Array/JsonArray.hpp           |  4 +-
 src/ArduinoJson/Collection/CollectionData.hpp |  2 +
 src/ArduinoJson/Collection/CollectionImpl.hpp | 12 ++++
 src/ArduinoJson/Document/JsonDocument.hpp     |  4 +-
 src/ArduinoJson/Variant/VariantData.hpp       | 15 +++++
 src/ArduinoJson/Variant/VariantRefBase.hpp    |  6 +-
 13 files changed, 213 insertions(+), 19 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8fa4d061..1f22103c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ HEAD
 * Allow using a `JsonVariant` as a key or index (issue #2080)
   Note: works only for reading, not for writing
 * Support `ElementProxy` and `MemberProxy` in `JsonDocument`'s constructor
+* Don't add partial objects when allocation fails (issue #2081)
 
 v7.0.4 (2024-03-12)
 ------
diff --git a/extras/tests/JsonArray/add.cpp b/extras/tests/JsonArray/add.cpp
index 22f99209..3a5a933c 100644
--- a/extras/tests/JsonArray/add.cpp
+++ b/extras/tests/JsonArray/add.cpp
@@ -179,3 +179,52 @@ TEST_CASE("JsonArray::add<T>()") {
     REQUIRE(doc.as<std::string>() == "[42]");
   }
 }
+
+TEST_CASE("JsonObject::add(JsonObject) ") {
+  JsonDocument doc1;
+  doc1[std::string("key1")] = std::string("value1");
+
+  TimebombAllocator allocator(10);
+  SpyingAllocator spy(&allocator);
+  JsonDocument doc2(&spy);
+  JsonArray array = doc2.to<JsonArray>();
+
+  SECTION("success") {
+    bool result = array.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == true);
+    REQUIRE(doc2.as<std::string>() == "[{\"key1\":\"value1\"}]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             Allocate(sizeofPool()),
+                             Allocate(sizeofString("key1")),
+                             Allocate(sizeofString("value1")),
+                         });
+  }
+
+  SECTION("partial failure") {  // issue #2081
+    allocator.setCountdown(2);
+
+    bool result = array.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == false);
+    REQUIRE(doc2.as<std::string>() == "[]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             Allocate(sizeofPool()),
+                             Allocate(sizeofString("key1")),
+                             AllocateFail(sizeofString("value1")),
+                             Deallocate(sizeofString("key1")),
+                         });
+  }
+
+  SECTION("complete failure") {
+    allocator.setCountdown(0);
+
+    bool result = array.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == false);
+    REQUIRE(doc2.as<std::string>() == "[]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             AllocateFail(sizeofPool()),
+                         });
+  }
+}
diff --git a/extras/tests/JsonDocument/add.cpp b/extras/tests/JsonDocument/add.cpp
index b7fe7581..c6ef3165 100644
--- a/extras/tests/JsonDocument/add.cpp
+++ b/extras/tests/JsonDocument/add.cpp
@@ -90,15 +90,57 @@ TEST_CASE("JsonDocument::add<T>()") {
     REQUIRE(doc.as<std::string>() == "[[1,2]]");
   }
 
-  SECTION("JsonObject") {
-    JsonObject object = doc.add<JsonObject>();
-    object["hello"] = "world";
-    REQUIRE(doc.as<std::string>() == "[{\"hello\":\"world\"}]");
-  }
-
   SECTION("JsonVariant") {
     JsonVariant variant = doc.add<JsonVariant>();
     variant.set(42);
     REQUIRE(doc.as<std::string>() == "[42]");
   }
 }
+
+TEST_CASE("JsonObject::add(JsonObject) ") {
+  JsonDocument doc1;
+  doc1[std::string("hello")] = std::string("world");
+
+  TimebombAllocator allocator(10);
+  SpyingAllocator spy(&allocator);
+  JsonDocument doc2(&spy);
+
+  SECTION("success") {
+    bool result = doc2.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == true);
+    REQUIRE(doc2.as<std::string>() == "[{\"hello\":\"world\"}]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             Allocate(sizeofPool()),
+                             Allocate(sizeofString("hello")),
+                             Allocate(sizeofString("world")),
+                         });
+  }
+
+  SECTION("partial failure") {  // issue #2081
+    allocator.setCountdown(2);
+
+    bool result = doc2.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == false);
+    REQUIRE(doc2.as<std::string>() == "[]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             Allocate(sizeofPool()),
+                             Allocate(sizeofString("hello")),
+                             AllocateFail(sizeofString("world")),
+                             Deallocate(sizeofString("hello")),
+                         });
+  }
+
+  SECTION("complete failure") {
+    allocator.setCountdown(0);
+
+    bool result = doc2.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == false);
+    REQUIRE(doc2.as<std::string>() == "[]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             AllocateFail(sizeofPool()),
+                         });
+  }
+}
diff --git a/extras/tests/JsonObject/set.cpp b/extras/tests/JsonObject/set.cpp
index 6c67d08a..3ea02939 100644
--- a/extras/tests/JsonObject/set.cpp
+++ b/extras/tests/JsonObject/set.cpp
@@ -118,7 +118,7 @@ TEST_CASE("JsonObject::set()") {
     bool success = obj3.set(obj1);
 
     REQUIRE(success == false);
-    REQUIRE(doc3.as<std::string>() == "{\"hello\":[null]}");
+    REQUIRE(doc3.as<std::string>() == "{\"hello\":[]}");
   }
 
   SECTION("destination is null") {
diff --git a/extras/tests/JsonVariant/add.cpp b/extras/tests/JsonVariant/add.cpp
index 71e53dc0..becd6645 100644
--- a/extras/tests/JsonVariant/add.cpp
+++ b/extras/tests/JsonVariant/add.cpp
@@ -6,6 +6,8 @@
 #include <stdint.h>
 #include <catch.hpp>
 
+#include "Allocators.hpp"
+
 TEST_CASE("JsonVariant::add(T)") {
   JsonDocument doc;
   JsonVariant var = doc.to<JsonVariant>();
@@ -56,15 +58,58 @@ TEST_CASE("JsonVariant::add<T>()") {
     REQUIRE(doc.as<std::string>() == "[[1,2]]");
   }
 
-  SECTION("JsonObject") {
-    JsonObject object = var.add<JsonObject>();
-    object["hello"] = "world";
-    REQUIRE(doc.as<std::string>() == "[{\"hello\":\"world\"}]");
-  }
-
   SECTION("JsonVariant") {
     JsonVariant variant = var.add<JsonVariant>();
     variant.set(42);
     REQUIRE(doc.as<std::string>() == "[42]");
   }
 }
+
+TEST_CASE("JsonObject::add(JsonObject) ") {
+  JsonDocument doc1;
+  doc1[std::string("hello")] = std::string("world");
+
+  TimebombAllocator allocator(10);
+  SpyingAllocator spy(&allocator);
+  JsonDocument doc2(&spy);
+  JsonVariant variant = doc2.to<JsonVariant>();
+
+  SECTION("success") {
+    bool result = variant.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == true);
+    REQUIRE(doc2.as<std::string>() == "[{\"hello\":\"world\"}]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             Allocate(sizeofPool()),
+                             Allocate(sizeofString("hello")),
+                             Allocate(sizeofString("world")),
+                         });
+  }
+
+  SECTION("partial failure") {  // issue #2081
+    allocator.setCountdown(2);
+
+    bool result = variant.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == false);
+    REQUIRE(doc2.as<std::string>() == "[]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             Allocate(sizeofPool()),
+                             Allocate(sizeofString("hello")),
+                             AllocateFail(sizeofString("world")),
+                             Deallocate(sizeofString("hello")),
+                         });
+  }
+
+  SECTION("complete failure") {
+    allocator.setCountdown(0);
+
+    bool result = variant.add(doc1.as<JsonObject>());
+
+    REQUIRE(result == false);
+    REQUIRE(doc2.as<std::string>() == "[]");
+    REQUIRE(spy.log() == AllocatorLog{
+                             AllocateFail(sizeofPool()),
+                         });
+  }
+}
diff --git a/src/ArduinoJson/Array/ArrayData.hpp b/src/ArduinoJson/Array/ArrayData.hpp
index 93a65d54..7e15ab19 100644
--- a/src/ArduinoJson/Array/ArrayData.hpp
+++ b/src/ArduinoJson/Array/ArrayData.hpp
@@ -20,6 +20,17 @@ class ArrayData : public CollectionData {
     return array->addElement(resources);
   }
 
+  template <typename T>
+  bool addValue(T&& value, ResourceManager* resources);
+
+  template <typename T>
+  static bool addValue(ArrayData* array, T&& value,
+                       ResourceManager* resources) {
+    if (!array)
+      return false;
+    return array->addValue(value, resources);
+  }
+
   VariantData* getOrAddElement(size_t index, ResourceManager* resources);
 
   VariantData* getElement(size_t index, const ResourceManager* resources) const;
diff --git a/src/ArduinoJson/Array/ArrayImpl.hpp b/src/ArduinoJson/Array/ArrayImpl.hpp
index 595e847b..9b96138a 100644
--- a/src/ArduinoJson/Array/ArrayImpl.hpp
+++ b/src/ArduinoJson/Array/ArrayImpl.hpp
@@ -47,4 +47,19 @@ inline void ArrayData::removeElement(size_t index, ResourceManager* resources) {
   remove(at(index, resources), resources);
 }
 
+template <typename T>
+inline bool ArrayData::addValue(T&& value, ResourceManager* resources) {
+  ARDUINOJSON_ASSERT(resources != nullptr);
+  auto slot = resources->allocSlot();
+  if (!slot)
+    return false;
+  JsonVariant variant(slot->data(), resources);
+  if (!variant.set(detail::forward<T>(value))) {
+    resources->freeSlot(slot);
+    return false;
+  }
+  addSlot(slot, resources);
+  return true;
+}
+
 ARDUINOJSON_END_PRIVATE_NAMESPACE
diff --git a/src/ArduinoJson/Array/JsonArray.hpp b/src/ArduinoJson/Array/JsonArray.hpp
index e0699a90..9a4e5f97 100644
--- a/src/ArduinoJson/Array/JsonArray.hpp
+++ b/src/ArduinoJson/Array/JsonArray.hpp
@@ -61,14 +61,14 @@ class JsonArray : public detail::VariantOperators<JsonArray> {
   // https://arduinojson.org/v7/api/jsonarray/add/
   template <typename T>
   bool add(const T& value) const {
-    return add<JsonVariant>().set(value);
+    return detail::ArrayData::addValue(data_, value, resources_);
   }
 
   // Appends a value to the array.
   // https://arduinojson.org/v7/api/jsonarray/add/
   template <typename T>
   bool add(T* value) const {
-    return add<JsonVariant>().set(value);
+    return detail::ArrayData::addValue(data_, value, resources_);
   }
 
   // Returns an iterator to the first element of the array.
diff --git a/src/ArduinoJson/Collection/CollectionData.hpp b/src/ArduinoJson/Collection/CollectionData.hpp
index f5706da0..07e69823 100644
--- a/src/ArduinoJson/Collection/CollectionData.hpp
+++ b/src/ArduinoJson/Collection/CollectionData.hpp
@@ -111,6 +111,8 @@ class CollectionData {
     return head_;
   }
 
+  void addSlot(SlotWithId slot, ResourceManager* resources);
+
  protected:
   iterator addSlot(ResourceManager*);
 
diff --git a/src/ArduinoJson/Collection/CollectionImpl.hpp b/src/ArduinoJson/Collection/CollectionImpl.hpp
index d077627b..c7e23a35 100644
--- a/src/ArduinoJson/Collection/CollectionImpl.hpp
+++ b/src/ArduinoJson/Collection/CollectionImpl.hpp
@@ -63,6 +63,18 @@ inline CollectionData::iterator CollectionData::addSlot(
   return iterator(slot, slot.id());
 }
 
+inline void CollectionData::addSlot(SlotWithId slot,
+                                    ResourceManager* resources) {
+  if (tail_ != NULL_SLOT) {
+    auto tail = resources->getSlot(tail_);
+    tail->setNext(slot.id());
+    tail_ = slot.id();
+  } else {
+    head_ = slot.id();
+    tail_ = slot.id();
+  }
+}
+
 inline void CollectionData::clear(ResourceManager* resources) {
   auto next = head_;
   while (next != NULL_SLOT) {
diff --git a/src/ArduinoJson/Document/JsonDocument.hpp b/src/ArduinoJson/Document/JsonDocument.hpp
index 1af36caf..acad2b2f 100644
--- a/src/ArduinoJson/Document/JsonDocument.hpp
+++ b/src/ArduinoJson/Document/JsonDocument.hpp
@@ -256,14 +256,14 @@ class JsonDocument : public detail::VariantOperators<const JsonDocument&> {
   // https://arduinojson.org/v7/api/jsondocument/add/
   template <typename TValue>
   bool add(const TValue& value) {
-    return add<JsonVariant>().set(value);
+    return data_.addValue(value, &resources_);
   }
 
   // Appends a value to the root array.
   // https://arduinojson.org/v7/api/jsondocument/add/
   template <typename TChar>
   bool add(TChar* value) {
-    return add<JsonVariant>().set(value);
+    return data_.addValue(value, &resources_);
   }
 
   // Removes an element of the root array.
diff --git a/src/ArduinoJson/Variant/VariantData.hpp b/src/ArduinoJson/Variant/VariantData.hpp
index 0c013e9d..5a2974ba 100644
--- a/src/ArduinoJson/Variant/VariantData.hpp
+++ b/src/ArduinoJson/Variant/VariantData.hpp
@@ -82,6 +82,21 @@ class VariantData {
     return var->addElement(resources);
   }
 
+  template <typename T>
+  bool addValue(T&& value, ResourceManager* resources) {
+    auto array = isNull() ? &toArray() : asArray();
+    return detail::ArrayData::addValue(array, detail::forward<T>(value),
+                                       resources);
+  }
+
+  template <typename T>
+  static bool addValue(VariantData* var, T&& value,
+                       ResourceManager* resources) {
+    if (!var)
+      return false;
+    return var->addValue(value, resources);
+  }
+
   bool asBoolean() const {
     switch (type()) {
       case VALUE_IS_BOOLEAN:
diff --git a/src/ArduinoJson/Variant/VariantRefBase.hpp b/src/ArduinoJson/Variant/VariantRefBase.hpp
index 00beacfe..3966f365 100644
--- a/src/ArduinoJson/Variant/VariantRefBase.hpp
+++ b/src/ArduinoJson/Variant/VariantRefBase.hpp
@@ -117,14 +117,16 @@ class VariantRefBase : public VariantTag {
   // https://arduinojson.org/v7/api/jsonvariant/add/
   template <typename T>
   bool add(const T& value) const {
-    return add<JsonVariant>().set(value);
+    return detail::VariantData::addValue(getOrCreateData(), value,
+                                         getResourceManager());
   }
 
   // Appends a value to the array.
   // https://arduinojson.org/v7/api/jsonvariant/add/
   template <typename T>
   bool add(T* value) const {
-    return add<JsonVariant>().set(value);
+    return detail::VariantData::addValue(getOrCreateData(), value,
+                                         getResourceManager());
   }
 
   // Removes an element of the array.