From d08bab23c4b4a49a7e81528c142ec515b48b396a Mon Sep 17 00:00:00 2001
From: guo-zhicheng666 <1678717630@qq.com>
Date: Wed, 18 Sep 2024 20:43:24 +0800
Subject: [PATCH] support monitor process function by kprobe

Signed-off-by: guo-zhicheng666 <1678717630@qq.com>
---
 module/fdstat.c | 123 ++++++++++++++++++++++++++++++++++++++++--------
 module/fdstat.h |   2 -
 2 files changed, 104 insertions(+), 21 deletions(-)

diff --git a/module/fdstat.c b/module/fdstat.c
index 5b1aa9a..c2f7d1d 100644
--- a/module/fdstat.c
+++ b/module/fdstat.c
@@ -11,40 +11,125 @@
 #include <linux/poll.h>
 #include <linux/proc_fs.h>
 #include <linux/uaccess.h>
+#include <linux/fdtable.h>
+#include <linux/atomic.h>
+#include <linux/kernel.h>
+#include <linux/module.h>
+#include <linux/fs.h>
+#include <linux/seq_file.h>
+#include <linux/uaccess.h>
+#include <linux/resource.h>
+
+#include <asm/ptrace.h>
+#include <asm/current.h>
+#include <linux/kprobes.h>
 
 #include "sysmonitor_main.h"
 
-#ifdef CONFIG_EULEROS_SYSMONITOR_FD
-static int do_fdstat(struct notifier_block *self, unsigned long val, void *data)
+static unsigned int fdthreshold = 80;
+struct proc_dir_entry *fdthreshold_entry = NULL;
+
+static unsigned int count_fd_num(struct fdtable *fdt)
+{
+	unsigned int size = fdt->max_fds;
+	unsigned int fd;
+	unsigned int handle_count = 0;
+
+	/* Find the last open fd */
+	for (fd = 0; fd < size; fd++) {
+		if (fdt->fd[fd])
+			handle_count++;
+	}
+	return handle_count;
+}
+
+static int handler_pre(struct kprobe *p, struct pt_regs *regs)
 {
-	struct fdstat *notifier_call_data = (struct fdstat *)data;
 	struct fdstat msg;
+	unsigned int total_fd_num = 0;
+	struct files_struct *files = current->files;
+	struct fdtable *fdt = files_fdtable(files);
+
+	if (fdt) {
+		total_fd_num = count_fd_num(fdt);
+	}
+
+	if (total_fd_num == (unsigned int)rlimit(RLIMIT_NOFILE) * fdthreshold / 100) {
+		pr_err("total_fd_num : %u\n", total_fd_num);
+		(void)memset(&msg, 0, sizeof(struct fdstat));
+		msg.pid = current->pid;
+		msg.total_fd_num =  total_fd_num;
+		(void)memcpy(msg.comm, current->comm, TASK_COMM_LEN);
+		(void)save_msg(FDSTAT, &msg, sizeof(struct fdstat));
+	}
+
+	return 0;
+}
+
+static int fdthreshold_show(struct seq_file *userfile, void *v)
+{
+	seq_printf(userfile, "%u\n", fdthreshold);
+	return 0;
+}
+
+static int fdthreshold_open(struct inode *inode, struct file *file)
+{
+	single_open(file, fdthreshold_show, NULL);
+
+	return 0;
+}
+
+static ssize_t fdthreshold_write(struct file *file, const char __user *ubuf,
+		size_t cnt, loff_t *ppos)
+{
 	int ret;
+	unsigned int val;
+
+	ret = kstrtouint_from_user(ubuf, cnt, 10, &val);
+	if (ret) {
+		pr_err("[fdstat] parse fdthreshold failed\n");
+		return ret;
+	}
 
-	(void)memset(&msg, 0, sizeof(struct fdstat));
-	msg.pid = notifier_call_data->pid;
-	msg.total_fd_num = notifier_call_data->total_fd_num + 1;
-	(void)memcpy(msg.comm, notifier_call_data->comm, TASK_COMM_LEN);
-	(void)save_msg(FDSTAT, &msg, sizeof(struct fdstat));
-	return NOTIFY_DONE;
+	if (val < 1 || val > 99) {
+		pr_err("[fdstat] fdthreshold is invalid\n");
+		return -EINVAL;
+	}
+
+	fdthreshold = val;
+
+	return cnt;
 }
 
-static struct notifier_block g_fdstat_nb = {
-	.notifier_call = do_fdstat,
-	.priority = NOTIFY_CALL_PRIORITY,
+static const struct proc_ops fdthreshold_operations = {
+	.proc_open		= fdthreshold_open,
+	.proc_read		= seq_read,
+	.proc_lseek		= seq_lseek,
+	.proc_release	= single_release,
+	.proc_write		= fdthreshold_write,
+};
+
+static struct kprobe kp = {
+	.symbol_name    = "do_sys_openat2",
+	.pre_handler    = handler_pre,
 };
-#endif
 
 void fdstat_init(void)
 {
-#ifdef CONFIG_EULEROS_SYSMONITOR_FD
-	(void)register_fdstat_notifier(&g_fdstat_nb);
-#endif
+	if (register_kprobe(&kp) < 0)
+		pr_err("Failed to register handler for %s\n", kp.symbol_name);
+
+	fdthreshold_entry = proc_create("fdthreshold", 0600, NULL, &fdthreshold_operations);
+	if (!fdthreshold_entry) {
+		pr_err("[fdstat]: create /proc/fdthreshold failed\n");
+	}
 }
 
 void fdstat_exit(void)
 {
-#ifdef CONFIG_EULEROS_SYSMONITOR_FD
-	(void)unregister_fdstat_notifier(&g_fdstat_nb);
-#endif
+	unregister_kprobe(&kp);
+
+	if (fdthreshold_entry != NULL) {
+		proc_remove(fdthreshold_entry);
+	}
 }
diff --git a/module/fdstat.h b/module/fdstat.h
index ae9ccc3..1a389ca 100644
--- a/module/fdstat.h
+++ b/module/fdstat.h
@@ -10,13 +10,11 @@
 #include <linux/types.h>
 #include <linux/sched.h>
 
-#ifndef CONFIG_EULEROS_SYSMONITOR_FD
 struct fdstat {
 	pid_t pid;
 	unsigned int total_fd_num;
 	char comm[TASK_COMM_LEN];
 };
-#endif
 
 void fdstat_init(void);
 void fdstat_exit(void);
-- 
2.33.0