From b8bc301301608b630e13121de00829d4424e280d Mon Sep 17 00:00:00 2001
From: sy03 <1370724210@qq.com>
Date: Mon, 16 Dec 2024 22:13:15 +0800
Subject: [PATCH] feature:Prevent immediate Segmentation fault in probe_write
 and probe_read. We will get error-code when Getting fault in
 probe_write&read.

---
 runtime/src/bpf_helper.cpp | 86 +++++++++++++++++++++++++++++++++++---
 1 file changed, 80 insertions(+), 6 deletions(-)

diff --git a/runtime/src/bpf_helper.cpp b/runtime/src/bpf_helper.cpp
index 981cddaf..748103c0 100644
--- a/runtime/src/bpf_helper.cpp
+++ b/runtime/src/bpf_helper.cpp
@@ -35,6 +35,8 @@
 #include <vector>
 #include <bpftime_shm_internal.hpp>
 #include <chrono>
+#include <setjmp.h>
+#include <signal.h>
 
 #define PATH_MAX 4096
 
@@ -64,20 +66,92 @@ long bpftime_strncmp(const char *s1, uint64_t s1_sz, const char *s2)
 	return strncmp(s1, s2, s1_sz);
 }
 
+static sigjmp_buf jump_buffer;
+
+// used for sigsetjmp to avoid faulting in probe_write_user & probe_read_user
+void segv_handler(int signum)
+{
+    siglongjmp(jump_buffer, 1);
+}
+
+
 uint64_t bpftime_probe_read(uint64_t dst, uint64_t size, uint64_t ptr, uint64_t,
 			    uint64_t)
 {
-	memcpy((void *)(uintptr_t)dst, (void *)(uintptr_t)ptr,
+	if(size<=0) {
+		return -EINVAL;
+	}
+	struct sigaction sa, old_sa;
+	uint64_t ret = 0;
+	
+	sa.sa_handler = segv_handler; // set signal handler
+	sigemptyset(&sa.sa_mask); // clear signal set
+	sa.sa_flags = 0;
+
+	if (sigaction(SIGSEGV, &sa, &old_sa) == -1) {
+		perror("sigaction");
+		return -EFAULT;
+	}
+	
+	if (sigsetjmp(jump_buffer, 1) == 0) {
+        // to avoid faulting for src & dst (which is NULL)	
+		memcpy((void *)(uintptr_t)dst, (void *)(uintptr_t)ptr,
 	       (size_t)(uint32_t)(size));
-	return 0;
+		spdlog::debug("probe_read: dst={}, src={}, len={}", dst, ptr, size);
+    } else {
+		spdlog::error("probe_read: failed to read from src={}", ptr);
+        ret = -EFAULT; // TO BE checked later
+    }
+
+    // restore the origin signal handler
+    if (sigaction(SIGSEGV, &old_sa, NULL) == -1) {
+        perror("sigaction restore");
+		// if failed to restore, return -EFAULT
+        return -EFAULT;
+    }
+
+    return ret;
 }
 
+
+
+
 uint64_t bpftime_probe_write_user(uint64_t dst, uint64_t src, uint64_t len,
-				  uint64_t, uint64_t)
+                                  uint64_t arg4, uint64_t arg5)
 {
-	memcpy((void *)(uintptr_t)dst, (void *)(uintptr_t)src,
-	       (size_t)(uint32_t)(len));
-	return 0;
+	if(len<=0) {
+		return -EINVAL;
+	}
+    struct sigaction sa, old_sa;
+    uint64_t ret = 0;
+
+    sa.sa_handler = segv_handler; // set signal handler
+    sigemptyset(&sa.sa_mask); // clear signal set
+    sa.sa_flags = 0; 
+
+    if (sigaction(SIGSEGV, &sa, &old_sa) == -1) {
+        perror("sigaction");
+        return -EFAULT;
+    }
+
+    if (sigsetjmp(jump_buffer, 1) == 0) {
+        // to avoid faulting for src & dst (which is NULL)	
+        memcpy((void *)(uintptr_t)dst, (void *)(uintptr_t)src,
+               (size_t)(uint32_t)(len));
+    	spdlog::debug("probe_write_user: dst={}, src={}, len={}", dst, src, len);
+    } else {
+		spdlog::error("probe_write_user: failed to write to dst={}", dst);
+        ret = -EFAULT; // TO BE checked later
+    }
+
+    // restore the origin signal handler
+    if (sigaction(SIGSEGV, &old_sa, NULL) == -1) {
+        perror("sigaction restore");
+		// if failed to restore, return -EFAULT
+        return -EFAULT;
+    }
+
+    return ret;
 }
 
 uint64_t bpftime_get_prandom_u32()