From 91fb7cbfd3b2d02c68a54fb7c35efc9d274a1bc6 Mon Sep 17 00:00:00 2001
From: Huaxin Lu <luhuaxin1@huawei.com>
Date: Wed, 14 Feb 2024 14:52:29 +0800
Subject: [PATCH 18/26] Add safe wapper for some memory and string functions

1. Warp some memory and strings functions to make them more safe.
2. Add detection for memory leakage.

Signed-off-by: Huaxin Lu <luhuaxin1@huawei.com>
---
 src/Makefile                                  |   4 +
 src/common/dim_baseline.c                     |   5 +-
 src/common/dim_baseline.h                     |   2 +-
 src/common/dim_measure_log.c                  |  17 +--
 src/common/dim_measure_log.h                  |   4 +-
 src/common/dim_safe_func.c                    |  18 +++
 src/common/dim_safe_func.h                    | 135 ++++++++++++++++++
 src/common/dim_symbol.c                       |   3 +-
 src/common/dim_tpm.c                          |   9 +-
 src/common/dim_utils.c                        |  28 +---
 src/common/dim_utils.h                        |   2 -
 src/core/dim_core_main.c                      |   6 +
 src/core/dim_core_mem_pool.c                  |   2 +-
 src/core/dim_core_mem_pool.h                  |   2 +-
 src/core/dim_core_sig.c                       |  30 ++--
 src/core/policy/dim_core_policy.c             |  25 ++--
 src/core/policy/dim_core_policy_complex.c     |  11 +-
 .../dim_core_static_baseline.c                |   6 +-
 .../dim_core_static_baseline_complex.c        |   2 +-
 src/core/tasks/dim_core_measure_kernel.c      |   2 +-
 .../dim_core_measure_process.c                |  18 ++-
 .../dim_core_measure_process_elf.c            |  22 +--
 .../dim_core_measure_process/dim_vm_hash.c    |   7 +-
 src/measure/dim_measure_baseline.c            |  10 +-
 src/monitor/dim_monitor_main.c                |   4 +
 25 files changed, 268 insertions(+), 106 deletions(-)
 create mode 100644 src/common/dim_safe_func.c
 create mode 100644 src/common/dim_safe_func.h

diff --git a/src/Makefile b/src/Makefile
index 8f4dce8..6782fd1 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -34,6 +34,7 @@ dim_core-objs += common/dim_hash.o
 dim_core-objs += common/dim_measure_log.o
 dim_core-objs += common/dim_tpm.o
 dim_core-objs += common/dim_symbol.o
+dim_core-objs += common/dim_safe_func.o
 dim_core-objs += measure/dim_measure.o
 dim_core-objs += measure/dim_measure_baseline.o
 dim_core-objs += measure/dim_measure_task.o
@@ -55,6 +56,7 @@ dim_monitor-objs += common/dim_measure_log.o
 dim_monitor-objs += common/dim_baseline.o
 dim_monitor-objs += common/dim_tpm.o
 dim_monitor-objs += common/dim_symbol.o
+dim_monitor-objs += common/dim_safe_func.o
 dim_monitor-objs += measure/dim_measure.o
 dim_monitor-objs += measure/dim_measure_baseline.o
 dim_monitor-objs += measure/dim_measure_task.o
@@ -72,6 +74,8 @@ ccflags-y += -I$(src)/common
 ccflags-y += -I$(src)/measure
 
 ccflags-y += -Wall -Werror -D_FORTIFY_SOURCE=2 -O2 -fstack-protector-strong
+# For check memory leakage
+# ccflags-y += -DDIM_DEBUG_MEMORY_LEAK
 
 KERNEL_SRC ?= /lib/modules/$(shell uname -r)/build
 PWD := $(shell pwd)
diff --git a/src/common/dim_baseline.c b/src/common/dim_baseline.c
index ec53b1c..e79458d 100644
--- a/src/common/dim_baseline.c
+++ b/src/common/dim_baseline.c
@@ -5,6 +5,7 @@
 #include "dim_rb.h"
 #include "dim_baseline.h"
 #include "dim_utils.h"
+#include "dim_safe_func.h"
 
 static int dim_baseline_compare(struct dim_baseline *x,
 				struct dim_baseline *y)
@@ -14,7 +15,7 @@ static int dim_baseline_compare(struct dim_baseline *x,
 	if (x->type != y->type)
 		return x->type > y->type ? 1 : -1;
 
-	ret = strcmp(x->name, y->name);
+	ret = dim_strcmp(x->name, y->name);
 	if (ret != 0)
 		return ret;
 
@@ -150,7 +151,7 @@ int dim_baseline_init_tree(malloc_func malloc, free_func free,
 	rwlock_init(&root->lock);
 	root->rb_root = RB_ROOT;
 	/* use kmalloc by default */
-	root->malloc = malloc == NULL ? dim_kmalloc_gfp : malloc;
+	root->malloc = malloc == NULL ? dim_kzalloc_gfp : malloc;
 	root->free = free == NULL ? dim_kfree : free;
 	return 0;
 }
diff --git a/src/common/dim_baseline.h b/src/common/dim_baseline.h
index 6e9d943..345b348 100644
--- a/src/common/dim_baseline.h
+++ b/src/common/dim_baseline.h
@@ -9,7 +9,7 @@
 #include "dim_hash.h"
 
 typedef void *(*malloc_func)(size_t);
-typedef void (*free_func)(void*);
+typedef void (*free_func)(const void*);
 
 enum dim_baseline_type {
 	DIM_BASELINE_USER, /* baseline of user process */
diff --git a/src/common/dim_measure_log.c b/src/common/dim_measure_log.c
index b4185d2..b84e635 100644
--- a/src/common/dim_measure_log.c
+++ b/src/common/dim_measure_log.c
@@ -6,6 +6,7 @@
 
 #include "dim_rb.h"
 #include "dim_tpm.h"
+#include "dim_safe_func.h"
 #include "dim_measure_log.h"
 
 /*
@@ -102,7 +103,7 @@ static int measure_info_insert(struct dim_measure_name *name,
 
 static void measure_log_destroy_info(struct dim_measure_log *info)
 {
-	kfree(info);
+	dim_kfree(info);
 }
 
 static void measure_log_destroy_name(struct dim_measure_name *name)
@@ -114,8 +115,8 @@ static void measure_log_destroy_name(struct dim_measure_name *name)
 	list_for_each_entry_safe(pos, n, &name->log_root, node)
 		measure_log_destroy_info(pos);
 	/* free self */
-	kfree(name->name);
-	kfree(name);
+	dim_kfree(name->name);
+	dim_kfree(name);
 }
 
 static int measure_log_create_name(const char *name_str,
@@ -123,13 +124,13 @@ static int measure_log_create_name(const char *name_str,
 {
 	struct dim_measure_name *new = NULL;
 
-	new = kzalloc(sizeof(struct dim_measure_name), GFP_KERNEL);
+	new = dim_kzalloc_gfp(sizeof(struct dim_measure_name));
 	if (new == NULL)
 		return -ENOMEM;
 
-	new->name = kstrdup(name_str, GFP_KERNEL);
+	new->name = dim_kstrdup_gfp(name_str);
 	if (new->name == NULL) {
-		kfree(new);
+		dim_kfree(new);
 		return -ENOMEM;
 	}
 
@@ -145,7 +146,7 @@ static int measure_log_create_info(char pcr, struct dim_digest *digest,
 	int ret = 0;
 	struct dim_measure_log *new = NULL;
 
-	new = kzalloc(sizeof(struct dim_measure_log), GFP_KERNEL);
+	new = dim_kzalloc_gfp(sizeof(struct dim_measure_log));
 	if (new == NULL)
 		return -ENOMEM;
 
@@ -153,7 +154,7 @@ static int measure_log_create_info(char pcr, struct dim_digest *digest,
 	new->type = flag;
 	ret = dim_digest_copy(&new->digest, digest);
 	if (ret < 0) {
-		kfree(new);
+		dim_kfree(new);
 		return ret;
 	}
 
diff --git a/src/common/dim_measure_log.h b/src/common/dim_measure_log.h
index 28f8a00..6ea2361 100644
--- a/src/common/dim_measure_log.h
+++ b/src/common/dim_measure_log.h
@@ -7,7 +7,9 @@
 
 #include <linux/list.h>
 #include <linux/rbtree.h>
+
 #include "dim_hash.h"
+#include "dim_safe_func.h"
 
 #define DIM_NG "dim-ng"
 #define LOG_MAX_LENGTH_PCR 3
@@ -61,7 +63,7 @@ struct dim_measure_log {
 static inline int dim_measure_name_compare(struct dim_measure_name *x,
 					   struct dim_measure_name *y)
 {
-	return strcmp(x->name, y->name);
+	return dim_strcmp(x->name, y->name);
 }
 
 static inline const char *dim_measure_log_type_to_name(int type)
diff --git a/src/common/dim_safe_func.c b/src/common/dim_safe_func.c
new file mode 100644
index 0000000..f13168c
--- /dev/null
+++ b/src/common/dim_safe_func.c
@@ -0,0 +1,18 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
+ */
+
+#include "dim_safe_func.h"
+
+#ifdef DIM_DEBUG_MEMORY_LEAK
+atomic_t dim_alloc_num = ATOMIC_INIT(0);;
+
+void dim_check_memory_leak(void)
+{
+	unsigned int n = atomic_read(&dim_alloc_num);
+	if (n != 0)
+		dim_warn("warning: detect %u memory leakage\n", n);
+	else
+		dim_info("not detect memory leakage\n");
+}
+#endif
\ No newline at end of file
diff --git a/src/common/dim_safe_func.h b/src/common/dim_safe_func.h
new file mode 100644
index 0000000..3e97f4e
--- /dev/null
+++ b/src/common/dim_safe_func.h
@@ -0,0 +1,135 @@
+/*
+ * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
+ */
+
+#ifndef __DIM_SAFE_FUNC_H
+#define __DIM_SAFE_FUNC_H
+
+#include <linux/atomic.h>
+#include <linux/slab.h>
+#include <linux/vmalloc.h>
+#include <linux/string.h>
+
+#include "dim_utils.h"
+
+#ifdef DIM_DEBUG_MEMORY_LEAK
+extern atomic_t dim_alloc_num;
+
+static inline void dim_alloc_debug_inc(void)
+{
+	atomic_inc(&dim_alloc_num);
+}
+
+static inline void dim_alloc_debug_dec(void)
+{
+	atomic_dec(&dim_alloc_num);
+}
+
+static inline void dim_print_alloc_num(const char *s)
+{
+	dim_info("%s: dim_alloc_num=%d\n", s, atomic_read(&dim_alloc_num));
+}
+
+void dim_check_memory_leak(void);
+#endif
+
+static inline void *dim_kzalloc_gfp(size_t size)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	void *data = kzalloc(size, GFP_KERNEL);
+	if (data != NULL)
+		dim_alloc_debug_inc();
+	return data;
+#else
+	return kzalloc(size, GFP_KERNEL);
+#endif
+}
+
+static inline void *dim_kcalloc_gfp(size_t n, size_t size)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	void *data = kcalloc(n, size, GFP_KERNEL);
+	if (data != NULL)
+		dim_alloc_debug_inc();
+	return data;
+#else
+	return kcalloc(n, size, GFP_KERNEL);
+#endif
+}
+
+static inline void *dim_krealloc_atom(const void *p, size_t new_size)
+{
+	return krealloc(p, new_size, GFP_ATOMIC);
+}
+
+static inline void *dim_kmemdup_gfp(const void *src, size_t len)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	void *data = kmemdup(src, len, GFP_KERNEL);
+	if (data != NULL)
+		dim_alloc_debug_inc();
+	return data;
+#else
+	return kmemdup(src, len, GFP_KERNEL);
+#endif
+}
+
+static inline void dim_kfree(const void *objp)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	if (objp != NULL)
+		dim_alloc_debug_dec();
+#endif
+	kfree(objp);
+}
+
+static inline void *dim_vzalloc(size_t size)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	void *data = vzalloc(size);
+	if (data != NULL)
+		dim_alloc_debug_inc();
+	return data;
+#else
+	return vzalloc(size);
+#endif
+}
+
+static inline void dim_vfree(void *data)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	if (data != NULL)
+		dim_alloc_debug_dec();
+#endif
+	vfree(data);
+}
+
+static inline char *dim_kstrdup_gfp(const char *s)
+{
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	void *data = kstrdup(s, GFP_KERNEL);
+	if (data != NULL)
+		dim_alloc_debug_inc();
+	return data;
+#else
+	return kstrdup(s, GFP_KERNEL);
+#endif
+}
+
+static inline int dim_strcmp(const char *cs, const char *ct)
+{
+	if (cs == NULL || ct == NULL)
+		return -1;
+
+	return strcmp(cs, ct);
+}
+
+static inline int dim_strncmp(const char *cs, const char *ct, size_t count)
+{
+	if (cs == NULL || ct == NULL)
+		return -1;
+
+	return strncmp(cs, ct, count);
+}
+
+#endif
\ No newline at end of file
diff --git a/src/common/dim_symbol.c b/src/common/dim_symbol.c
index 48f6491..63824e6 100644
--- a/src/common/dim_symbol.c
+++ b/src/common/dim_symbol.c
@@ -4,6 +4,7 @@
 
 #include <linux/kallsyms.h>
 
+#include "dim_safe_func.h"
 #include "dim_symbol.h"
 
 static int find_kernel_symbol(unsigned long addr,
@@ -34,7 +35,7 @@ DIM_SYMBOL_LOOKUP_FUNC dim_get_symbol_lookup_func(void)
 		if (ret < 0 || offset > size)
 			break;
 
-		if (strcmp(symbol_name, DIM_KALLSYMS_LOOKUP_NAME) == 0)
+		if (dim_strcmp(symbol_name, DIM_KALLSYMS_LOOKUP_NAME) == 0)
 			return (DIM_SYMBOL_LOOKUP_FUNC)(kaddr - offset);
 
 		if (kaddr == next) {
diff --git a/src/common/dim_tpm.c b/src/common/dim_tpm.c
index 5d983e8..6d3c255 100644
--- a/src/common/dim_tpm.c
+++ b/src/common/dim_tpm.c
@@ -4,6 +4,7 @@
 
 #include <linux/crypto.h>
 
+#include "dim_safe_func.h"
 #include "dim_tpm.h"
 
 int dim_tpm_init(struct dim_tpm *tpm, int algo)
@@ -15,8 +16,8 @@ int dim_tpm_init(struct dim_tpm *tpm, int algo)
 	if (tpm->chip == NULL)
 		return -ENODEV;
 
-	tpm->digests = kcalloc(tpm->chip->nr_allocated_banks,
-			       sizeof(struct tpm_digest), GFP_KERNEL);
+	tpm->digests = dim_kcalloc_gfp(tpm->chip->nr_allocated_banks,
+				       sizeof(struct tpm_digest));
 	if (tpm->digests == NULL) {
 		ret = -ENOMEM;
 		goto err;
@@ -40,7 +41,7 @@ int dim_tpm_init(struct dim_tpm *tpm, int algo)
 err:
 	put_device(&tpm->chip->dev);
 	if (tpm->digests != NULL) {
-		kfree(tpm->digests);
+		dim_kfree(tpm->digests);
 		tpm->digests = NULL;
 	}
 
@@ -72,5 +73,5 @@ void dim_tpm_destroy(struct dim_tpm *tpm)
 		return;
 
 	put_device(&tpm->chip->dev);
-	kfree(tpm->digests);
+	dim_kfree(tpm->digests);
 }
\ No newline at end of file
diff --git a/src/common/dim_utils.c b/src/common/dim_utils.c
index 4c99879..598e824 100644
--- a/src/common/dim_utils.c
+++ b/src/common/dim_utils.c
@@ -5,22 +5,10 @@
 #include <linux/fs.h>
 #include <linux/err.h>
 #include <linux/namei.h>
-#include <linux/slab.h>
-#include <linux/vmalloc.h>
 
+#include "dim_safe_func.h"
 #include "dim_utils.h"
 
-void *dim_kmalloc_gfp(size_t size)
-{
-	return kmalloc(size, GFP_KERNEL);
-}
-
-void dim_kfree(void *data)
-{
-	if (data != NULL)
-		kfree(data);
-}
-
 int dim_get_absolute_path(const char *path, const char **result)
 {
 	int ret = 0;
@@ -35,7 +23,7 @@ int dim_get_absolute_path(const char *path, const char **result)
 	if (ret < 0)
 		return ret;
 
-	buf = dim_kmalloc_gfp(PATH_MAX);
+	buf = dim_kzalloc_gfp(PATH_MAX);
 	if (buf == NULL) {
 		ret = -ENOMEM;
 		goto out;
@@ -47,16 +35,14 @@ int dim_get_absolute_path(const char *path, const char **result)
 		goto out;
 	}
 
-	*result = kstrdup(apath, GFP_KERNEL);
+	*result = dim_kstrdup_gfp(apath);
 	if (*result == NULL) {
 		ret = -ENOMEM;
 		goto out;
 	}
 out:
 	path_put(&p);
-	if (buf != NULL)
-		dim_kfree(buf);
-
+	dim_kfree(buf);
 	return ret;	
 }
 
@@ -72,7 +58,7 @@ bool dim_string_end_with(const char *str, const char *ext)
 	if (name_len < ext_len)
 		return false;
 
-	return strcmp(str + name_len - ext_len, ext) == 0;
+	return dim_strcmp(str + name_len - ext_len, ext) == 0;
 }
 
 int dim_parse_line_buf(char *buf, loff_t len, int (*line_parser)(char *, int, void *), void *data)
@@ -97,7 +83,7 @@ int dim_parse_line_buf(char *buf, loff_t len, int (*line_parser)(char *, int, vo
 			line = &buf[i + 1];
 		} else {
 			line_len = buf + i - line + 1;
-			line_buf = kzalloc(line_len + 1, GFP_KERNEL);
+			line_buf = dim_kzalloc_gfp(line_len + 1);
 			if (line_buf == NULL)
 				return -ENOMEM;
 
@@ -118,7 +104,7 @@ int dim_parse_line_buf(char *buf, loff_t len, int (*line_parser)(char *, int, vo
 	}
 out:
 	if (line_buf != NULL)
-		kfree(line_buf);
+		dim_kfree(line_buf);
 
 	return ret;
 }
diff --git a/src/common/dim_utils.h b/src/common/dim_utils.h
index 5a9f132..8c7d855 100644
--- a/src/common/dim_utils.h
+++ b/src/common/dim_utils.h
@@ -17,8 +17,6 @@
 #define dim_info(fmt, ...) pr_info(dim_fmt(fmt), THIS_MODULE->name, ##__VA_ARGS__)
 #define dim_devel(fmt, ...)
 
-void *dim_kmalloc_gfp(size_t size);
-void dim_kfree(void *data);
 int dim_get_absolute_path(const char *path, const char **result);
 bool dim_string_end_with(const char *str, const char *ext);
 int dim_parse_line_buf(char *buf, loff_t len, int (*line_parser)(char *, int, void *), void *data);
diff --git a/src/core/dim_core_main.c b/src/core/dim_core_main.c
index ae34e81..c62fa09 100644
--- a/src/core/dim_core_main.c
+++ b/src/core/dim_core_main.c
@@ -4,6 +4,8 @@
 
 #include <linux/module.h>
 
+#include "dim_safe_func.h"
+
 #include "dim_core_policy.h"
 #include "dim_core_symbol.h"
 #include "dim_core_fs.h"
@@ -99,6 +101,10 @@ static void __exit dim_core_exit(void)
 
 	if (signature)
 		dim_core_sig_destroy();
+
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	dim_check_memory_leak();
+#endif
 }
 
 module_init(dim_core_init);
diff --git a/src/core/dim_core_mem_pool.c b/src/core/dim_core_mem_pool.c
index 160f819..974033f 100644
--- a/src/core/dim_core_mem_pool.c
+++ b/src/core/dim_core_mem_pool.c
@@ -112,7 +112,7 @@ out:
 	return data->data;
 }
 
-void dim_mem_pool_free(void *data)
+void dim_mem_pool_free(const void *data)
 {
 	struct dim_pool_mem *mem = NULL;
 
diff --git a/src/core/dim_core_mem_pool.h b/src/core/dim_core_mem_pool.h
index c566dc8..5c4cdea 100644
--- a/src/core/dim_core_mem_pool.h
+++ b/src/core/dim_core_mem_pool.h
@@ -25,7 +25,7 @@ typedef void (*pool_chunk_visitor)(struct gen_pool *,
 int dim_mem_pool_init(void);
 void dim_mem_pool_destroy(void);
 void *dim_mem_pool_alloc(size_t size);
-void dim_mem_pool_free(void *data);
+void dim_mem_pool_free(const void *data);
 void dim_mem_pool_walk_chunk(pool_chunk_visitor f, void *data);
 
 #endif
diff --git a/src/core/dim_core_sig.c b/src/core/dim_core_sig.c
index 70a3469..f142050 100644
--- a/src/core/dim_core_sig.c
+++ b/src/core/dim_core_sig.c
@@ -13,6 +13,7 @@
 
 #include "dim_hash.h"
 #include "dim_utils.h"
+#include "dim_safe_func.h"
 
 #include "dim_core_sig.h"
 
@@ -26,7 +27,7 @@ static char *add_suffix(const char *str, const char *suffix)
 	char *buf = NULL;
 
 	len = strlen(str) + strlen(suffix) + 1;
-	buf = dim_kmalloc_gfp(len);
+	buf = dim_kzalloc_gfp(len);
 	if (buf == NULL)
 		return NULL;
 
@@ -39,10 +40,16 @@ static int read_file_root(struct path *root, const char *name, void **buf)
 	int ret = 0;
 	struct file *file = NULL;
 
-	if (root == NULL)
-		return kernel_read_file_from_path(name, 0, buf,
+	if (root == NULL) {
+		ret = kernel_read_file_from_path(name, 0, buf,
 						  DIM_CORE_MAX_FILE_SIZE,
 						  NULL, READING_UNKNOWN);
+#ifdef DIM_DEBUG_MEMORY_LEAK
+		if (*buf != NULL)
+			dim_alloc_debug_inc();
+#endif
+		return ret;
+	}
 
 	file = file_open_root(root, name, O_RDONLY, 0);
 	if (IS_ERR(file))
@@ -50,6 +57,10 @@ static int read_file_root(struct path *root, const char *name, void **buf)
 
 	ret = kernel_read_file(file, 0, buf, DIM_CORE_MAX_FILE_SIZE,
 			       NULL, READING_UNKNOWN);
+#ifdef DIM_DEBUG_MEMORY_LEAK
+	if (*buf != NULL)
+		dim_alloc_debug_inc();
+#endif
 	(void)filp_close(file, NULL);
 	return ret;
 }
@@ -114,12 +125,10 @@ int dim_read_verify_file(struct path *root, const char *name, void **buf)
 	sig_size = ret;
 	ret = dim_core_sig_verify(file_buf, file_size, sig_buf, sig_size);
 out:
-	if (sig_name != NULL)
-		kfree(sig_name);
-	if (sig_buf != NULL)
-		vfree(sig_buf);
-	if (file_buf != NULL && ret < 0)
-		vfree(file_buf);
+	dim_kfree(sig_name);
+	dim_vfree(sig_buf);
+	if (ret < 0)
+		dim_vfree(file_buf);
 	if (ret == 0) {
 		*buf = file_buf;
 		ret = file_size;
@@ -172,8 +181,7 @@ int dim_core_sig_init(void)
 	dim_info("load DIM cert: %s\n", dim_core_key->description);
 	ret = 0;
 err:
-	if (data != NULL)
-		vfree(data);
+	dim_vfree(data);
 	if (ret < 0)
 		key_put(dim_core_keyring);
 	return ret;
diff --git a/src/core/policy/dim_core_policy.c b/src/core/policy/dim_core_policy.c
index 4d7bcc1..f10a256 100644
--- a/src/core/policy/dim_core_policy.c
+++ b/src/core/policy/dim_core_policy.c
@@ -11,8 +11,9 @@
 #include <linux/namei.h>
 #include <linux/utsname.h>
 
-#include "dim_utils.h"
 #include "dim_rb.h"
+#include "dim_utils.h"
+#include "dim_safe_func.h"
 
 #include "dim_core_sig.h"
 #include "dim_core_policy.h"
@@ -26,13 +27,9 @@ static int dim_policy_compare(struct dim_policy *x, struct dim_policy *y)
 
 	switch (x->obj) {
 	case DIM_POLICY_OBJ_BPRM_TEXT:
-		if (x->path == NULL || y->path == NULL)
-			return -1;
-		return strcmp(x->path, y->path);
+		return dim_strcmp(x->path, y->path);
 	case DIM_POLICY_OBJ_MODULE_TEXT:
-		if (x->name == NULL || y->name == NULL)
-			return -1;
-		return strcmp(x->name, y->name);
+		return dim_strcmp(x->name, y->name);
 	case DIM_POLICY_OBJ_KERNEL_TEXT:
 		return 0;
 	default:
@@ -60,8 +57,8 @@ void policy_destroy(struct dim_policy *policy)
 	if (policy == NULL)
 		return;
 
-	dim_kfree((char *)policy->name);
-	dim_kfree((char *)policy->path);
+	dim_kfree(policy->name);
+	dim_kfree(policy->path);
 	dim_kfree(policy);
 }
 
@@ -100,15 +97,15 @@ static int policy_check_add_bprm_text(struct dim_policy *policy)
 		return 0;
 	}
 
-	if (strcmp(apath, policy->path) == 0) {
+	if (dim_strcmp(apath, policy->path) == 0) {
 		/* the two paths are same, no need to add another policy */
-		dim_kfree((char *)apath);
+		dim_kfree(apath);
 		return 0;
 	}
 
-	p = kmemdup(policy, sizeof(struct dim_policy), GFP_KERNEL);
+	p = dim_kmemdup_gfp(policy, sizeof(struct dim_policy));
 	if (p == NULL) {
-		dim_kfree((char *)apath);
+		dim_kfree(apath);
 		return -ENOMEM;
 	}
 
@@ -191,7 +188,7 @@ int dim_core_policy_load(void)
 		dim_core_policy_destroy();
 	}
 
-	vfree(buf);
+	dim_vfree(buf);
 	return ret;
 }
 
diff --git a/src/core/policy/dim_core_policy_complex.c b/src/core/policy/dim_core_policy_complex.c
index b29483d..18a9e58 100644
--- a/src/core/policy/dim_core_policy_complex.c
+++ b/src/core/policy/dim_core_policy_complex.c
@@ -4,7 +4,9 @@
 
 #include <linux/slab.h>
 
+#include "dim_rb.h"
 #include "dim_utils.h"
+#include "dim_safe_func.h"
 
 #include "dim_core_policy.h"
 
@@ -41,7 +43,7 @@ static const char *dim_policy_action_str[DIM_POLICY_KEY_LAST] = {
 
 static const char *policy_get_string_value(const char *s)
 {
-	return kstrdup(s, GFP_KERNEL);
+	return dim_kstrdup_gfp(s);
 }
 
 static int policy_get_action(const char *s)
@@ -121,7 +123,7 @@ static int parse_line(char *line_str, struct dim_policy *policy)
 	char *p = NULL;
 
 	if ((p = strsep(&line_str, " ")) == NULL ||
-	    strcmp(p, DIM_POLICY_MEASURE) != 0) {
+	    dim_strcmp(p, DIM_POLICY_MEASURE) != 0) {
 		dim_err("invalid policy prefix, must start with %s\n",
 			DIM_POLICY_MEASURE);
 		return -EINVAL;
@@ -156,14 +158,13 @@ static int policy_parse_line(char* line, int line_no, void *data)
 		return -EINVAL;
 	}
 
-	policy = dim_kmalloc_gfp(sizeof(struct dim_policy));
+	policy = dim_kzalloc_gfp(sizeof(struct dim_policy));
 	if (policy == NULL)
 		return -ENOMEM;
 
-	memset(policy, 0, sizeof(struct dim_policy));
-
 	ret = parse_line(line, policy);
 	if (ret < 0) {
+		policy_destroy(policy);
 		dim_err("fail to parse policy at line %d: %d\n", line_no, ret);
 		return ret;
 	}
diff --git a/src/core/static_baseline/dim_core_static_baseline.c b/src/core/static_baseline/dim_core_static_baseline.c
index 49810f3..e33c67c 100644
--- a/src/core/static_baseline/dim_core_static_baseline.c
+++ b/src/core/static_baseline/dim_core_static_baseline.c
@@ -28,12 +28,12 @@ static bool baseline_match_policy(const char *name, int type)
 		return dim_core_policy_match(DIM_POLICY_OBJ_BPRM_TEXT,
 					     DIM_POLICY_KEY_PATH, name);
 
-	if (strcmp(name, kr) == 0)
+	if (dim_strcmp(name, kr) == 0)
 		return dim_core_policy_match(DIM_POLICY_OBJ_KERNEL_TEXT,
 					     DIM_POLICY_KEY_NAME, kr);
 
 	if (name_len <= kr_len + 2 || /* <kernel release>/<mod_name> */
-	    strncmp(kr, name, kr_len) != 0 ||
+	    dim_strncmp(kr, name, kr_len) != 0 ||
 	    *(name + kr_len) != '/')
 		return false;
 
@@ -82,7 +82,7 @@ static_baseline_load(struct dir_context *__ctx,
 		dim_err("failed to parse baseline file %s: %d\n", name, ret);
 out:
 	if (buf != NULL)
-		vfree(buf);
+		dim_vfree(buf);
 
 #if LINUX_VERSION_CODE < KERNEL_VERSION(6, 4, 0)
 	return 0; /* ignore fail */
diff --git a/src/core/static_baseline/dim_core_static_baseline_complex.c b/src/core/static_baseline/dim_core_static_baseline_complex.c
index 685118f..8ff7c86 100644
--- a/src/core/static_baseline/dim_core_static_baseline_complex.c
+++ b/src/core/static_baseline/dim_core_static_baseline_complex.c
@@ -37,7 +37,7 @@ static int parse_line(char* line, int line_no, void *data)
 	}
 
 	if ((p = strsep(&line_str, " ")) == NULL ||
-	    strcmp(p, DIM_STATIC_BASELINE_PREFIX) != 0) {
+	    dim_strcmp(p, DIM_STATIC_BASELINE_PREFIX) != 0) {
 		dim_warn("invalid baseline prefix at line %d\n", line_no);
 		return 0;
 	}
diff --git a/src/core/tasks/dim_core_measure_kernel.c b/src/core/tasks/dim_core_measure_kernel.c
index e13e177..d49095b 100644
--- a/src/core/tasks/dim_core_measure_kernel.c
+++ b/src/core/tasks/dim_core_measure_kernel.c
@@ -131,7 +131,7 @@ static int calc_kernel_digest(struct dim_hash *hash, struct dim_digest *digest)
 	if (ret < 0)
 		dim_err("failed to calculate kernel digest: %d\n", ret);
 
-	vfree(jcode_sort);
+	dim_vfree(jcode_sort);
 	return ret;
 }
 
diff --git a/src/core/tasks/dim_core_measure_process/dim_core_measure_process.c b/src/core/tasks/dim_core_measure_process/dim_core_measure_process.c
index 513f5a0..8522085 100644
--- a/src/core/tasks/dim_core_measure_process/dim_core_measure_process.c
+++ b/src/core/tasks/dim_core_measure_process/dim_core_measure_process.c
@@ -60,9 +60,8 @@ static int store_task_tree(struct task_struct *p, void *data)
 
 		/* realloc to size * 2 */
 		new_size = ctx->size << 1;
-		tmp = krealloc(ctx->buf,
-			       new_size * sizeof(struct task_struct *),
-			       GFP_ATOMIC);
+		tmp = dim_krealloc_atom(ctx->buf,
+				new_size * sizeof(struct task_struct *));
 		if (tmp == NULL)
 			return -ENOMEM;
 
@@ -79,7 +78,7 @@ static int kill_task_tree(struct task_struct *tsk)
 	const int def_size = 32;
 	struct task_kill_ctx ctx = { .size = def_size };
 
-	ctx.buf = dim_kmalloc_gfp(def_size * sizeof(struct task_struct *));
+	ctx.buf = dim_kzalloc_gfp(def_size * sizeof(struct task_struct *));
 	if (ctx.buf == NULL)
 		return -ENOMEM;
 
@@ -91,7 +90,7 @@ static int kill_task_tree(struct task_struct *tsk)
 		}
 	}
 
-	kfree(ctx.buf);
+	dim_kfree(ctx.buf);
 	send_sig(SIGKILL, tsk, 1);
 	return 0;
 }
@@ -166,7 +165,6 @@ static void measure_task_module(struct vm_area_struct *vma,
 	ret = measure_process_text(vma, ctx);
 	if (ret < 0)
 		dim_err("failed to measure module file text: %d", ret);
-
 }
 
 static int measure_task(struct task_struct *task, struct task_measure_ctx *ctx)
@@ -238,7 +236,7 @@ static int store_task_pids(pid_t **pid_buf, unsigned int *pid_cnt)
 	unsigned int max_cnt = (PID_MAX_DEFAULT << 1);
 
 	/* maximum processing of PID_MAX_DEFAULT * 2 pids */
-	buf = vmalloc(max_cnt);
+	buf = dim_vzalloc(max_cnt);
 	if (buf == NULL) {
 		dim_err("failed to allocate memory for pid buffer\n");
 		return -ENOMEM;
@@ -294,7 +292,7 @@ static int walk_measure_tasks(struct task_measure_ctx *ctx)
 		}
 	}
 
-	vfree(pid_buf);
+	dim_vfree(pid_buf);
 	return 0;
 }
 
@@ -306,7 +304,7 @@ static int user_text_measure(int mode, struct dim_measure *m)
 	if (m == NULL)
 		return -EINVAL;
 
-	ctx = vmalloc(sizeof(struct task_measure_ctx));
+	ctx = dim_vzalloc(sizeof(struct task_measure_ctx));
 	if (ctx == NULL)
 		return -ENOMEM;
 
@@ -315,7 +313,7 @@ static int user_text_measure(int mode, struct dim_measure *m)
 	ctx->check = check_process_digest;
 
 	ret = walk_measure_tasks(ctx);
-	vfree(ctx);
+	dim_vfree(ctx);
 	return ret;
 }
 
diff --git a/src/core/tasks/dim_core_measure_process/dim_core_measure_process_elf.c b/src/core/tasks/dim_core_measure_process/dim_core_measure_process_elf.c
index 9210f47..3821c7f 100644
--- a/src/core/tasks/dim_core_measure_process/dim_core_measure_process_elf.c
+++ b/src/core/tasks/dim_core_measure_process/dim_core_measure_process_elf.c
@@ -57,7 +57,7 @@ static int get_elf_phdrs(struct file *elf_file, struct elfhdr *ehdr,
 		return -ENOEXEC;
 
 	phdr_size = sizeof(struct elf_phdr) * ehdr->e_phnum;
-	elf_phdata = dim_kmalloc_gfp(phdr_size);
+	elf_phdata = dim_kzalloc_gfp(phdr_size);
 	if (elf_phdata == NULL)
 		return -ENOMEM;
 
@@ -89,7 +89,7 @@ static int get_elf_section(struct file *elf_file, struct elfhdr *ehdr,
 	if (ehdr->e_shentsize != sizeof(struct elf_shdr))
 		return -EBADF;
 
-	sh_table = dim_kmalloc_gfp(ehdr->e_shentsize);
+	sh_table = dim_kzalloc_gfp(ehdr->e_shentsize);
 	if (sh_table == NULL)
 		return -ENOMEM;
 
@@ -103,21 +103,21 @@ static int get_elf_section(struct file *elf_file, struct elfhdr *ehdr,
 
 	str_size = sh_table->sh_size;
 	if (str_size > i_size_read(file_inode(elf_file))) {
-		kfree(sh_table);
+		dim_kfree(sh_table);
 		return -EBADF;
 	}
 
-	sh_str = vmalloc(str_size);
+	sh_str = dim_vzalloc(str_size);
 	if (sh_str == NULL) {
-		kfree(sh_table);
+		dim_kfree(sh_table);
 		return -ENOMEM;
 	}
 
 	pos = sh_table->sh_offset;
 	size = kernel_read(elf_file, sh_str, sh_table->sh_size, &pos);
 	if (size != sh_table->sh_size) {
-		kfree(sh_table);
-		vfree(sh_str);
+		dim_kfree(sh_table);
+		dim_vfree(sh_str);
 		return size < 0 ? (int)size : -EBADF;
 	}
 
@@ -135,15 +135,15 @@ static int get_elf_section(struct file *elf_file, struct elfhdr *ehdr,
 		    sh_table->sh_name + name_len >= str_size)
 		    	break;
 
-		if (strcmp(name, sh_str + sh_table->sh_name) == 0) {
+		if (dim_strcmp(name, sh_str + sh_table->sh_name) == 0) {
 			memcpy(shdr, sh_table, sizeof(struct elf_shdr));
 			ret = 0;
 			break;
 		}
 	}
 
-	kfree(sh_table);
-	vfree(sh_str);
+	dim_kfree(sh_table);
+	dim_vfree(sh_str);
 	return ret;
 }
 
@@ -179,7 +179,7 @@ static int get_elf_text_phdrs(struct file *elf_file,
 	}
 
 	/* alloc memory buffer for phdrs */
-	phdrs_text = dim_kmalloc_gfp(phdrs_text_num * sizeof(struct elf_phdr));
+	phdrs_text = dim_kzalloc_gfp(phdrs_text_num * sizeof(struct elf_phdr));
 	if (phdrs_text == NULL) {
 		dim_kfree(phdrs_get);
 		return -ENOMEM;
diff --git a/src/core/tasks/dim_core_measure_process/dim_vm_hash.c b/src/core/tasks/dim_core_measure_process/dim_vm_hash.c
index 0c59b9e..c3a8887 100644
--- a/src/core/tasks/dim_core_measure_process/dim_vm_hash.c
+++ b/src/core/tasks/dim_core_measure_process/dim_vm_hash.c
@@ -7,6 +7,7 @@
 #include <linux/highmem.h>
 
 #include "dim_utils.h"
+#include "dim_safe_func.h"
 
 #include "dim_vm_hash.h"
 
@@ -26,7 +27,7 @@ int dim_vm_hash_update_address(struct mm_struct *mm,
 	if (mm == NULL || addr_len == 0 || shash == NULL)
 		return -EINVAL;
 
-	pages = vzalloc(nr_pages * sizeof(struct page *));
+	pages = dim_vzalloc(nr_pages * sizeof(struct page *));
 	if (pages == NULL)
 		return -ENOMEM;
 
@@ -38,7 +39,7 @@ int dim_vm_hash_update_address(struct mm_struct *mm,
 #endif
 	if (ret_pages < 0) {
 		dim_err("failed to get remote pages: %ld\n", ret_pages);
-		vfree(pages);
+		dim_vfree(pages);
 		return ret_pages;
 	} else if (ret_pages != nr_pages) {
 		dim_warn("failed to get all remote pages\n");
@@ -64,7 +65,7 @@ int dim_vm_hash_update_address(struct mm_struct *mm,
 		put_page(pages[i]);
 	}
 
-	vfree(pages);
+	dim_vfree(pages);
 	return 0;
 }
 
diff --git a/src/measure/dim_measure_baseline.c b/src/measure/dim_measure_baseline.c
index dc358a7..c62d6be 100644
--- a/src/measure/dim_measure_baseline.c
+++ b/src/measure/dim_measure_baseline.c
@@ -16,7 +16,7 @@ static const char *process_static_name(const char *name, int type,
 {
 	const char *kr = init_uts_ns.name.release;
 
-	if (type != DIM_BASELINE_KERNEL || strcmp(name, kr) == 0)
+	if (type != DIM_BASELINE_KERNEL || dim_strcmp(name, kr) == 0)
 		return name;
 
 	/* name of kernel module has a kernel prefix in static baseline */
@@ -101,7 +101,7 @@ static int measure_log_add(struct dim_measure *m, const char *name,
 
 /* check dynamic measurement result in baseline stage */
 static int process_dynamic_baseline(struct dim_measure *m, const char *name,
-				    struct dim_digest *digest, int *log_flag) // TODO
+				    struct dim_digest *digest, int *log_flag)
 {
 	int ret = 0;
 	struct dim_digest digest_static = { 0 };
@@ -134,7 +134,7 @@ static int process_dynamic_measure(struct dim_measure *m, const char *name,
 
 	if(!dynamic_baseline_match(m, name, DIM_BASELINE_KERNEL, digest)) {
 		dim_err("mismatch dynamic baseline of kernel %s\n", name);
-		if (log_flag != NULL) // TODO
+		if (log_flag != NULL)
 			*log_flag = LOG_TAMPERED;
 
 		return measure_log_add(m, name, digest, LOG_TAMPERED);
@@ -165,7 +165,7 @@ static int process_static_baseline(struct dim_measure *m, const char *name,
 		return measure_log_add(m, name, digest, LOG_STATIC_BASELINE);
 
 	dim_warn("mismatch static baseline of user process %s\n", name);
-	if (log_flag != NULL) // TODO
+	if (log_flag != NULL)
 		*log_flag = LOG_TAMPERED;
 
 	return measure_log_add(m, name, digest, LOG_TAMPERED);
@@ -177,7 +177,7 @@ static int process_static_measure(struct dim_measure *m, const char *name,
 {
 	if(!dynamic_baseline_match(m, name, DIM_BASELINE_USER, digest)) {
 		dim_err("mismatch dynamic baseline of user %s\n", name);
-		if (log_flag != NULL) // TODO
+		if (log_flag != NULL)
 			*log_flag = LOG_TAMPERED;
 
 		return measure_log_add(m, name, digest, LOG_TAMPERED);
diff --git a/src/monitor/dim_monitor_main.c b/src/monitor/dim_monitor_main.c
index 4b3505d..d0e89f1 100644
--- a/src/monitor/dim_monitor_main.c
+++ b/src/monitor/dim_monitor_main.c
@@ -57,6 +57,10 @@ static void __exit dim_monitor_exit(void)
 {
 	dim_monitor_measure_destroy();
 	dim_monitor_destroy_fs();
+
+	#ifdef DIM_DEBUG_MEMORY_LEAK
+		dim_check_memory_leak();
+	#endif
 }
 
 module_init(dim_monitor_init);
-- 
2.33.0