From bcf7604479236b16622fd0b01f4d4353ef21ce43 Mon Sep 17 00:00:00 2001
From: Azat Khuzhin <a3at.mail@gmail.com>
Date: Fri, 3 May 2024 09:23:41 +0200
Subject: [PATCH] Fix merge operator support (that can be visible by iterating
 through the node)

The problem is that const_map_to.get(*j->first) compares only the
shared_ptr's, while we need to compare the key itself.

v2: remove const iterator
v3: fix use-after-free due to const reference
---
 src/nodebuilder.cpp                 | 19 ++++++++++++++-----
 test/integration/load_node_test.cpp | 15 +++++++++++++++
 2 files changed, 29 insertions(+), 5 deletions(-)

diff --git a/src/nodebuilder.cpp b/src/nodebuilder.cpp
index a1daf8e..aed6581 100644
--- a/src/nodebuilder.cpp
+++ b/src/nodebuilder.cpp
@@ -1,3 +1,4 @@
+#include <algorithm>
 #include <cassert>
 
 #include "nodebuilder.h"
@@ -75,12 +76,20 @@ void NodeBuilder::OnMapStart(const Mark& mark, const std::string& tag,
 
 void MergeMapCollection(detail::node& map_to, detail::node& map_from,
                         detail::shared_memory_holder& pMemory) {
-  const detail::node& const_map_to = map_to;
   for (auto j = map_from.begin(); j != map_from.end(); j++) {
-    detail::node* s = const_map_to.get(*j->first, pMemory);
-    if (s == nullptr) {
-      map_to.insert(*j->first, *j->second, pMemory);
-    }
+    const auto from_key = j->first;
+    /// NOTE: const_map_to.get(*j->first) cannot be used here, since it
+    /// compares only the shared_ptr's, while we need to compare the key
+    /// itself.
+    ///
+    /// NOTE: get() also iterates over elements
+    bool found = std::any_of(map_to.begin(), map_to.end(), [&](const detail::node_iterator_value<detail::node> & kv)
+    {
+        const auto key_node = kv.first;
+        return key_node->scalar() == from_key->scalar();
+    });
+    if (!found)
+      map_to.insert(*from_key, *j->second, pMemory);
   }
 }
 
diff --git a/test/integration/load_node_test.cpp b/test/integration/load_node_test.cpp
index aa81faa..8ca2d35 100644
--- a/test/integration/load_node_test.cpp
+++ b/test/integration/load_node_test.cpp
@@ -1,6 +1,7 @@
 #include "yaml-cpp/yaml.h"  // IWYU pragma: keep
 
 #include "gtest/gtest.h"
+#include <algorithm>
 
 namespace YAML {
 namespace {
@@ -184,6 +185,20 @@ TEST(LoadNodeTest, MergeKeyA) {
   EXPECT_EQ(1, node["z"]["c"].as<int>());
 }
 
+TEST(LoadNodeTest, MergeKeyAIterator) {
+  Node node = Load(
+      "{x: &foo {a : 1,b : 1,c : 1}, y: &bar {d: 2, e : 2, f : 2, a : 2}, z: "
+      "&stuff { << : *foo, b : 3} }");
+  EXPECT_EQ(NodeType::Map, node["z"].Type());
+
+  const auto& z = node["z"];
+  size_t z_b_keys = std::count_if(z.begin(), z.end(), [&](const detail::iterator_value & kv)
+  {
+    return kv.first.as<std::string>() == "b";
+  });
+  ASSERT_EQ(z_b_keys, 1);
+}
+
 TEST(LoadNodeTest, MergeKeyB) {
   Node node = Load(
       "{x: &foo {a : 1,b : 1,c : 1}, y: &bar {d: 2, e : 2, f : 2, a : 2}, z: "