From 279379c3bee1cd45d95b67bc8dcbff027db61195 Mon Sep 17 00:00:00 2001
From: baldurk <baldurk@baldurk.org>
Date: Thu, 10 Jul 2025 17:23:35 +0100
Subject: [PATCH] Fix potential trample in trie re-using leaf

---
 renderdoc/core/rdcbytetrie.h           | 20 +++++++++++++++-----
 renderdoc/replay/basic_types_tests.cpp |  2 ++
 2 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/renderdoc/core/rdcbytetrie.h b/renderdoc/core/rdcbytetrie.h
index 241bf2805..3f037a521 100644
--- a/renderdoc/core/rdcbytetrie.h
+++ b/renderdoc/core/rdcbytetrie.h
@@ -158,6 +158,9 @@ private:
 
     void SetPrefix(const Key &k)
     {
+      if(IsLeaf() && k.size > GetPrefixLength())
+        RDCERR("Expanding prefix on leaf - invalid");
+
       // prefix stored immediately after this structure for both node and leaf
       byte *prefix = (byte *)(this + 1);
       // use memmove to account for prefix shrinking in place
@@ -488,11 +491,18 @@ private:
         promoted->SetPrefix(leaf->GetPrefix());
         promoted->SetValue(std::move(leaf->GetValue()));
 
-        // re-use the leaf later. Since we just promoted this to a node we know it will have no
-        // children, so below we are going to hit the case of the next byte having no child and we
-        // can put this leaf there.
-        leaf->RemoveValue();
-        leaf->SetPrefix(searchAfter);
+        // re-use the leaf later if it has enough prefix space. Since we just promoted this to a
+        // node we know it will have no children, so below we are going to hit the case of the next
+        // byte having no child and we can put this leaf there.
+        if(leaf->GetPrefixLength() >= searchAfter.size)
+        {
+          leaf->RemoveValue();
+          leaf->SetPrefix(searchAfter);
+        }
+        else
+        {
+          leaf = MakeLeaf(searchAfter);
+        }
 
         root = promoted;
 
diff --git a/renderdoc/replay/basic_types_tests.cpp b/renderdoc/replay/basic_types_tests.cpp
index e84cf2736..96ea86896 100644
--- a/renderdoc/replay/basic_types_tests.cpp
+++ b/renderdoc/replay/basic_types_tests.cpp
@@ -2593,6 +2593,8 @@ TEST_CASE("Test rdcbytetrie type", "[basictypes][rdcbytetrie]")
         {{2, 1, 1}, TrieValue(5)},
         //
         {{2, 1, 6}, TrieValue(6)},
+        //
+        {{2, 1, 6, 9, 9}, TrieValue(7)},
     };
     bytebuf d = {1, 3};