From 15dafb1c1ca0db2b06715e30cdee30d21c54572d Mon Sep 17 00:00:00 2001
From: Shubham Lagwankar <shubhu105@gmail.com>
Date: Thu, 20 Jun 2019 02:58:08 -0400
Subject: [PATCH] Problem: radix tree's apply function uses incorrect resizing
 logic (#3548)

* Problem: radix tree's apply function uses incorrect resizing logic

Solution: rewrite the function using a vector

The buffer is extended once by 256 bytes, which will not be enough if
the node holds data larger than this number.
---
 src/radix_tree.cpp | 35 +++++++++++++++++------------------
 1 file changed, 17 insertions(+), 18 deletions(-)

diff --git a/src/radix_tree.cpp b/src/radix_tree.cpp
index 208856ad..7c40e09f 100644
--- a/src/radix_tree.cpp
+++ b/src/radix_tree.cpp
@@ -34,6 +34,7 @@
 
 #include <stdlib.h>
 #include <string.h>
+#include <vector>
 
 node::node (unsigned char *data) : data_ (data)
 {
@@ -517,35 +518,33 @@ bool zmq::radix_tree::check (const unsigned char *key, size_t size)
 
 static void
 visit_keys (node n,
-            unsigned char **buffer,
-            size_t buffer_size,
-            size_t maxbuffer_size,
+            std::vector<unsigned char> &buffer,
             void (*func) (unsigned char *data, size_t size, void *arg),
             void *arg)
 {
-    if (buffer_size >= maxbuffer_size) {
-        maxbuffer_size += 256;
-        *buffer =
-          static_cast<unsigned char *> (realloc (*buffer, maxbuffer_size));
-        zmq_assert (*buffer);
+    for (size_t i = 0; i < n.prefix_length (); ++i)
+        buffer.push_back (n.prefix ()[i]);
+
+    if (n.refcount () > 0) {
+        zmq_assert (!buffer.empty ());
+        func (&buffer[0], buffer.size (), arg);
     }
 
-    for (size_t i = 0; i < n.prefix_length (); ++i)
-        (*buffer)[buffer_size++] = n.prefix ()[i];
-    if (n.refcount () > 0)
-        func (*buffer, buffer_size, arg);
     for (size_t i = 0; i < n.edgecount (); ++i)
-        visit_keys (n.node_at (i), buffer, buffer_size, maxbuffer_size, func,
-                    arg);
-    buffer_size -= n.prefix_length ();
+        visit_keys (n.node_at (i), buffer, func, arg);
+    for (size_t i = 0; i < n.prefix_length (); ++i)
+        buffer.pop_back ();
 }
 
 void zmq::radix_tree::apply (
   void (*func) (unsigned char *data, size_t size, void *arg), void *arg)
 {
-    unsigned char *buffer = NULL;
-    visit_keys (root_, &buffer, 0, 0, func, arg);
-    free (buffer);
+    if (root_.refcount () > 0)
+        func (NULL, 0, arg); // Root node is always empty.
+
+    std::vector<unsigned char> buffer;
+    for (size_t i = 0; i < root_.edgecount (); ++i)
+        visit_keys (root_.node_at (i), buffer, func, arg);
 }
 
 size_t zmq::radix_tree::size () const