From 4bac002ec33f3593ef03ff01494b524234c40484 Mon Sep 17 00:00:00 2001
From: renoseven <dev@renoseven.net>
Date: Tue, 26 Dec 2023 13:07:53 +0800
Subject: [PATCH 15/15] upatch: fix memory leak

Signed-off-by: renoseven <dev@renoseven.net>
---
 upatch/upatch-manage/upatch-patch.c  |   4 +-
 upatch/upatch-tool/upatch-meta.c     |   8 ++-
 upatch/upatch-tool/upatch-resolve.c  |  63 +++++++---------
 upatch/upatch-tool/upatch-resolve.h  |   3 +-
 upatch/upatch-tool/upatch-tool-lib.c | 104 +++++++++++++++++++++------
 5 files changed, 117 insertions(+), 65 deletions(-)

diff --git a/upatch/upatch-manage/upatch-patch.c b/upatch/upatch-manage/upatch-patch.c
index 762641c..4763778 100644
--- a/upatch/upatch-manage/upatch-patch.c
+++ b/upatch/upatch-manage/upatch-patch.c
@@ -590,8 +590,8 @@ static int upatch_apply_patches(struct upatch_process *proc,
 
 	if (!found) {
 		ret = -1;
-		log_debug("can't found inode %lu in pid %d\n",
-			  uelf->relf->info.inode, proc->pid);
+		log_debug("Cannot find inode %lu in pid %d, file is not loaded\n",
+			uelf->relf->info.inode, proc->pid);
 		goto out;
 	}
 
diff --git a/upatch/upatch-tool/upatch-meta.c b/upatch/upatch-tool/upatch-meta.c
index c625115..2bd39b4 100644
--- a/upatch/upatch-tool/upatch-meta.c
+++ b/upatch/upatch-tool/upatch-meta.c
@@ -197,9 +197,11 @@ static int patch_deactive_in_cover(struct upatch_meta_patch *patch)
 static int list_add_symbol(struct list_head *head, struct upatch_meta_symbol *sym)
 {
 	struct upatch_meta_symbol *newsym = (struct upatch_meta_symbol *)malloc(sizeof(struct upatch_meta_symbol));
-	if (newsym == NULL)
+	if (newsym == NULL) {
 		return ENOMEM;
+	}
 	memset(newsym, 0, sizeof(struct upatch_meta_symbol));
+
 	strncpy(newsym->name, sym->name, sizeof(newsym->name));
 	newsym->offset = sym->offset;
 	INIT_LIST_HEAD(&newsym->self);
@@ -212,9 +214,11 @@ static int list_add_symbol(struct list_head *head, struct upatch_meta_symbol *sy
 static int list_add_symbol_for_patch(struct upatch_meta_patch *patch, struct list_head *head, struct upatch_meta_symbol *sym)
 {
 	struct upatch_meta_symbol *newsym = (struct upatch_meta_symbol *)malloc(sizeof(struct upatch_meta_symbol));
-	if (newsym == NULL)
+	if (newsym == NULL) {
 		return ENOMEM;
+	}
 	memset(newsym, 0, sizeof(struct upatch_meta_symbol));
+
 	strncpy(newsym->name, sym->name, sizeof(newsym->name));
 	newsym->offset = sym->offset;
 	INIT_LIST_HEAD(&newsym->self);
diff --git a/upatch/upatch-tool/upatch-resolve.c b/upatch/upatch-tool/upatch-resolve.c
index aeb655c..e9d75e0 100644
--- a/upatch/upatch-tool/upatch-resolve.c
+++ b/upatch/upatch-tool/upatch-resolve.c
@@ -42,79 +42,66 @@ out:
 static int list_add_symbol(struct list_head *head, patch_symbols_t *sym)
 {
 	patch_symbols_t *newsym = (patch_symbols_t *)malloc(sizeof(patch_symbols_t));
-	if (newsym == NULL)
+	if (newsym == NULL) {
 		return ENOMEM;
+	}
 
 	memset(newsym, 0, sizeof(patch_symbols_t));
 	strncpy(newsym->name, sym->name, sizeof(newsym->name));
 	newsym->offset = sym->offset;
 	INIT_LIST_HEAD(&newsym->self);
 	list_add(&newsym->self, head);
+
 	return 0;
 }
 
-struct list_head* patch_symbols_resolve(const char *target_elf, const char *patch_file) {
-	struct upatch_elf uelf;
-	struct running_elf relf;
-	GElf_Shdr *upatch_shdr = NULL;
-	struct upatch_patch_func *upatch_funcs = NULL;
-	GElf_Off min_addr; // binary base
-	int num;
-	struct list_head *head = malloc(sizeof(struct list_head));
+struct list_head* patch_symbols_resolve(struct upatch_elf *uelf, struct running_elf *relf) {
+	struct list_head *head = NULL;
 
-	INIT_LIST_HEAD(head);
-
-	int ret = upatch_init(&uelf, patch_file);
-	if (ret < 0) {
-		log_warn("upatch-resolve: upatch_init failed\n");
-		goto out;
-	}
-
-	ret = binary_init(&relf, target_elf);
-	if (ret < 0) {
-		log_warn("upatch-resolve: binary_init failed %d \n", ret);
-		goto out;
-	}
-
-	if (check_build_id(&uelf.info, &relf.info) == false) {
+	if (check_build_id(&uelf->info, &relf->info) == false) {
 		log_error("upatch-resolve: Build id mismatched!\n");
 		goto out;
 	}
 
-	uelf.relf = &relf;
-	upatch_shdr = &uelf.info.shdrs[uelf.index.upatch_funcs];
-	upatch_funcs = uelf.info.patch_buff + upatch_shdr->sh_offset;
-	min_addr = calculate_load_address(uelf.relf, false);
+	GElf_Shdr *upatch_shdr = &uelf->info.shdrs[uelf->index.upatch_funcs];
+	GElf_Off min_addr = calculate_load_address(uelf->relf, false);
 	if (min_addr == (GElf_Off)-1) {
 		goto out;
 	}
 
-	num = upatch_shdr->sh_size / sizeof(*upatch_funcs);
+	struct upatch_patch_func *upatch_funcs = uelf->info.patch_buff + upatch_shdr->sh_offset;
+	int num = upatch_shdr->sh_size / sizeof(*upatch_funcs);
 
 	log_debug("upatch-resolve: sh_size %lu, sizeof %lu \n", upatch_shdr->sh_size, sizeof(*upatch_funcs));
 	log_debug("upatch-resolve: elf base addr is 0x%lx, num is %d\n", min_addr, num);
 
+	head = malloc(sizeof(struct list_head));
+	INIT_LIST_HEAD(head);
+
 	for (int i = 0; i < num; i++) {
-		patch_symbols_t *sym = malloc(sizeof(patch_symbols_t));
-		sprintf(sym->name, "sym_%d", i);
-		sym->offset = upatch_funcs[i].old_addr - min_addr;;
-		log_debug("+upatch-resolve: sym->offset addr is 0x%lx\n", sym->offset);
-		list_add_symbol(head, sym);
+		patch_symbols_t sym;
+
+		sprintf(sym.name, "sym_%d", i);
+		sym.offset = upatch_funcs[i].old_addr - min_addr;
+		log_debug("upatch-resolve: sym->offset addr is 0x%lx\n", sym.offset);
+
+		list_add_symbol(head, &sym); // This would copy the symbol
 	}
 
 	return head;
+
 out:
-	free(head);
+	if (head != NULL) {
+		free(head);
+	}
 	return NULL;
 }
 
 void patch_symbols_free(struct list_head *symbols) {
 	patch_symbols_t *sym, *next;
-
-	if (!symbols)
-		return;
 	list_for_each_entry_safe (sym, next, symbols, self) {
 		list_del(&sym->self);
 		free(sym);
 	}
+	free(symbols);
 }
diff --git a/upatch/upatch-tool/upatch-resolve.h b/upatch/upatch-tool/upatch-resolve.h
index 7d0bb47..7c06345 100644
--- a/upatch/upatch-tool/upatch-resolve.h
+++ b/upatch/upatch-tool/upatch-resolve.h
@@ -2,8 +2,9 @@
 #define __UPATCH_RESOLVE_H_
 
 #include "list.h"
+#include "upatch-elf.h"
 
-struct list_head* patch_symbols_resolve(const char *target_elf, const char *patch_file);
+struct list_head* patch_symbols_resolve(struct upatch_elf *uelf, struct running_elf *relf);
 void patch_symbols_free(struct list_head *symbols);
 
 #endif
diff --git a/upatch/upatch-tool/upatch-tool-lib.c b/upatch/upatch-tool/upatch-tool-lib.c
index 4e5bda1..3731abe 100644
--- a/upatch/upatch-tool/upatch-tool-lib.c
+++ b/upatch/upatch-tool/upatch-tool-lib.c
@@ -14,21 +14,43 @@
 
 #include "log.h"
 #include "list.h"
+#include "upatch-elf.h"
 #include "upatch-meta.h"
 #include "upatch-resolve.h"
 #include "upatch-ioctl.h"
 
 int upatch_check(const char *target_elf, const char *patch_file, char *err_msg, size_t max_len)
 {
-    struct list_head *patch_syms = patch_symbols_resolve(target_elf, patch_file);
+    int ret = 0;
+    struct list_head *patch_syms = NULL;
+    struct list_head *collision_list = NULL;
+
+    struct upatch_elf uelf;
+    ret = upatch_init(&uelf, patch_file);
+	if (ret < 0) {
+        snprintf(err_msg, max_len, "Failed to read patch");
+		goto out;
+	}
+
+    struct running_elf relf;
+	ret = binary_init(&relf, target_elf);
+	if (ret < 0) {
+        snprintf(err_msg, max_len, "Failed to read target elf");
+		goto out;
+	}
+	uelf.relf = &relf;
+
+    patch_syms = patch_symbols_resolve(&uelf, &relf);
     if (patch_syms == NULL) {
         snprintf(err_msg, max_len, "Patch format error");
-        return ENOEXEC;
+        ret = ENOEXEC;
+        goto out;
     }
 
-    struct list_head *collision_list = meta_get_symbol_collision(target_elf, patch_syms);
+    collision_list = meta_get_symbol_collision(target_elf, patch_syms);
     if (collision_list == NULL) {
-        return 0;
+        ret = 0;
+        goto out;
     }
 
     int offset = snprintf(err_msg, max_len, "Patch is conflicted with ");
@@ -39,50 +61,78 @@ int upatch_check(const char *target_elf, const char *patch_file, char *err_msg,
         offset = snprintf(err_msg, max_len, "\"%s\" ", collision->uuid);
     }
 
-    patch_symbols_free(patch_syms);
-    meta_put_symbol_collision(collision_list);
+out:
+    if (patch_syms != NULL) {
+        patch_symbols_free(patch_syms);
+    }
+    if (collision_list != NULL) {
+        meta_put_symbol_collision(collision_list);
+    }
+	binary_close(&relf);
+	upatch_close(&uelf);
 
-    return EEXIST;
+    return ret;
 }
 
 int upatch_load(const char *uuid, const char *target, const char *patch, bool force)
 {
+    int ret = 0;
+    struct list_head *patch_syms = NULL;
+    patch_entity_t *patch_entity = NULL;
+    struct list_head *collision_syms = NULL;
+
     // Pointer check
     if (uuid == NULL || target == NULL || patch == NULL) {
         return EINVAL;
     }
     log_normal("Loading patch {%s} (\"%s\") for \"%s\"\n", uuid, patch, target);
 
+    struct upatch_elf uelf;
+    ret = upatch_init(&uelf, patch);
+	if (ret < 0) {
+        log_warn("Failed to read patch\n");
+		goto out;
+	}
+
+    struct running_elf relf;
+	ret = binary_init(&relf, target);
+	if (ret < 0) {
+        log_warn("Failed to read target elf\n");
+		goto out;
+	}
+	uelf.relf = &relf;
+
     // Fails if patch is already exist
     if (meta_get_patch_status(uuid) != UPATCH_PATCH_STATUS_NOT_APPLIED) {
         log_warn("{%s}: Patch status is invalid\n", uuid);
-        return EPERM;
+        ret = EPERM;
+        goto out;
     }
 
     // Resolve patch symbols
-    struct list_head *patch_syms = patch_symbols_resolve(target, patch);
+    patch_syms = patch_symbols_resolve(&uelf, &relf);
     if (patch_syms == NULL) {
         log_warn("{%s}: Patch format error\n", uuid);
-        return ENOEXEC;
+        ret = ENOEXEC;
+        goto out;
     }
 
     // Check patch symbol collision
     if (!force) {
-        struct list_head *collision_syms = meta_get_symbol_collision(target, patch_syms);
+        collision_syms = meta_get_symbol_collision(target, patch_syms);
         if (collision_syms != NULL) {
             log_warn("{%s}: Patch symbol conflicted\n", uuid);
-            patch_symbols_free(patch_syms);
-            meta_put_symbol_collision(collision_syms);
-            return EEXIST;
+            ret = EEXIST;
+            goto out;
         }
     }
 
     // Alloc memory for patch
-    patch_entity_t *patch_entity = calloc(1, sizeof(patch_entity_t));
+    patch_entity = calloc(1, sizeof(patch_entity_t));
     if (patch_entity == NULL) {
         log_warn("{%s}: Failed to alloc memory\n", uuid);
-        patch_symbols_free(patch_syms);
-        return ENOMEM;
+        ret = ENOMEM;
+        goto out;
     }
 
     strncpy(patch_entity->target_path, target, strnlen(target, PATH_MAX));
@@ -91,17 +141,27 @@ int upatch_load(const char *uuid, const char *target, const char *patch, bool fo
     log_normal("patch: %s, patch_path: %s\n", patch, patch_entity->patch_path);
     patch_entity->symbols = patch_syms;
 
-    int ret = meta_create_patch(uuid, patch_entity);
+    ret = meta_create_patch(uuid, patch_entity);
     if (ret != 0) {
         log_warn("{%s}: Failed to create patch entity\n", uuid);
-        free(patch_entity);
-        patch_symbols_free(patch_syms);
-        return ret;
+        goto out;
     }
 
-    free(patch_entity);
     meta_set_patch_status(uuid, UPATCH_PATCH_STATUS_DEACTIVED);
 
+out:
+    if (collision_syms != NULL) {
+        meta_put_symbol_collision(collision_syms);
+    }
+    if (patch_syms != NULL) {
+        patch_symbols_free(patch_syms);
+    }
+    if (patch_entity != NULL) {
+        free(patch_entity);
+    }
+	binary_close(&relf);
+	upatch_close(&uelf);
+
     return ret;
 }
 
-- 
2.33.0

