Additional security improvements: fix memory handling edge cases

Co-authored-by: wannaphong <8536487+wannaphong@users.noreply.github.com>

Author: Copilot

_codeql_detected_source_root

@@ -0,0 +1 @@
+.

src/trie.c

@@ -83,16 +83,23 @@ static TrieNode* trie_node_add_child(TrieNode* node, int codepoint) {
     if (node->num_children >= node->capacity) {
         int new_capacity = node->capacity == 0 ? INITIAL_CAPACITY : node->capacity * 2;
 
+        /* Allocate both arrays before updating pointers */
         TrieNode** new_children = (TrieNode**)realloc(node->children, 
                                                        new_capacity * sizeof(TrieNode*));
         if (!new_children) return NULL;
 
         int* new_chars = (int*)realloc(node->child_chars, new_capacity * sizeof(int));
         if (!new_chars) {
-            /* Restore old pointer to avoid leak, but still fail */
+            /* new_children was allocated but new_chars failed */
+            /* Since realloc succeeded for new_children, the old pointer is invalid */
+            /* We must use the new_children pointer, even though we can't proceed */
+            node->children = new_children;
+            /* Alternatively, we could try to restore by reallocating to old size */
+            /* But that could also fail, so we just update and return NULL */
             return NULL;
         }
 
+        /* Both allocations succeeded, update pointers */
         node->children = new_children;
         node->child_chars = new_chars;
         node->capacity = new_capacity;
@@ -201,8 +208,11 @@ int trie_prefixes(Trie* trie, const char* text, char*** prefixes, int** lengths)
     *prefixes = (char**)malloc(max_prefixes * sizeof(char*));
     *lengths = (int*)malloc(max_prefixes * sizeof(int));
     if (!*prefixes || !*lengths) {
-        free(*prefixes);
-        free(*lengths);
+        /* Clean up any successful allocation */
+        if (*prefixes) free(*prefixes);
+        if (*lengths) free(*lengths);
+        *prefixes = NULL;
+        *lengths = NULL;
         return 0;
     }
 

tests/test_newmm.c

@@ -17,6 +17,9 @@ typedef struct {
 static int test_count = 0;
 static int test_passed = 0;
 
+#define INITIAL_OUTPUT_SIZE 1024
+#define TOKEN_OVERHEAD 4  /* For "'', " around each token */
+
 void run_test(const char* text, const char* dict_path, const char* expected, const char* description) {
     test_count++;
     printf("\n[Test %d] %s\n", test_count, description);
@@ -44,7 +47,7 @@ void run_test(const char* text, const char* dict_path, const char* expected, con
     }
 
     /* Build output string with dynamic allocation */
-    size_t output_size = 1024;
+    size_t output_size = INITIAL_OUTPUT_SIZE;
     char* output = (char*)malloc(output_size);
     if (!output) {
         printf("❌ FAIL: Memory allocation failed\n");
@@ -54,7 +57,7 @@ void run_test(const char* text, const char* dict_path, const char* expected, con
 
     strcpy(output, "[");
     for (int i = 0; i < token_count; i++) {
-        size_t needed = strlen(output) + strlen(tokens[i]) + 10; /* "'', " + margins */
+        size_t needed = strlen(output) + strlen(tokens[i]) + TOKEN_OVERHEAD + 1; /* +1 for null */
         if (needed >= output_size) {
             output_size = needed * 2;
             char* new_output = (char*)realloc(output, output_size);