From dc42619c663fb13185047452741481055d3e5761 Mon Sep 17 00:00:00 2001
From: Burkov Egor <eburkov@rvision.ru>
Date: Tue, 29 Jul 2025 15:24:59 +0300
Subject: [PATCH] fix: malloc checks at msc_tree

---
 src/utils/msc_tree.cc | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/utils/msc_tree.cc b/src/utils/msc_tree.cc
index 27c7461c..977d7147 100644
--- a/src/utils/msc_tree.cc
+++ b/src/utils/msc_tree.cc
@@ -419,6 +419,9 @@ TreeNode *CPTAddElement(unsigned char *ipdata, unsigned int ip_bitmask, CPTTree
                 node->count++;
                 new_node = node;
                 node->netmasks = reinterpret_cast<unsigned char *>(malloc(node->count *  sizeof(unsigned char)));
+                if (node->netmasks == NULL) {
+                    return node;
+                }
                 memset(node->netmasks, 0, (node->count *  sizeof(unsigned char)));
 
                 if ((node->count -1) == 0) {
@@ -491,8 +494,6 @@ TreeNode *CPTAddElement(unsigned char *ipdata, unsigned int ip_bitmask, CPTTree
             }
 
             i_node->netmasks = reinterpret_cast<unsigned char *>(malloc((node->count - i) * sizeof(unsigned char)));
-            memset(i_node->netmasks, 0, ((node->count - i) * sizeof(unsigned char)));
-
             if(i_node->netmasks == NULL) {
                 free(new_node->prefix);
                 free(new_node);
@@ -500,6 +501,8 @@ TreeNode *CPTAddElement(unsigned char *ipdata, unsigned int ip_bitmask, CPTTree
                 return NULL;
             }
 
+            memset(i_node->netmasks, 0, ((node->count - i) * sizeof(unsigned char)));
+
             j = 0;
             while (j < (node->count - i))   {
                 i_node->netmasks[j] = node->netmasks[i + j];