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(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(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];