feat:Prevent segmentation fault in probe_read & probe_write (#368)

* feat:Prevent segmentation fault in probe_read & probe_write * update: update status & probe_test_case * bug fix: return value for syscall & nullptr * update enum class & fix error of compile flag * update CI of runtime-unit-test * fix error of CI in build runtime with mpk enable * update compile settings * fix benchmark docs errors and update more instructions * fix bug of probe size < 0 * update unit-test * add macro definition * supplement the macro definitions
eunomia-bpf · Dec 24, 2024 · 3f6724b · 3f6724b
1 parent 50a1e15
commit 3f6724b
Show file tree

Hide file tree

Showing 8 changed files with 355 additions and 14 deletions.
diff --git a/.github/workflows/test-runtime.yml b/.github/workflows/test-runtime.yml
@@ -40,7 +40,7 @@ jobs:
   
     - name: Build
       run: |
-        cmake -DTEST_LCOV=ON -DBPFTIME_LLVM_JIT=YES -DBPFTIME_ENABLE_UNIT_TESTING=YES -DCMAKE_BUILD_TYPE=Debug -B build
+        cmake -DTEST_LCOV=ON -DBPFTIME_LLVM_JIT=YES -DBPFTIME_ENABLE_UNIT_TESTING=YES -DENABLE_PROBE_WRITE_CHECK=1 -DENABLE_PROBE_READ_CHECK=1 -DCMAKE_BUILD_TYPE=Debug -B build
         cmake --build build --config Debug --target bpftime_runtime_tests -j$(nproc)
     - name: Test Runtime
       run:  ./build/runtime/unit-test/bpftime_runtime_tests
@@ -69,7 +69,7 @@ jobs:
     - name: build runtime with mpk enable
       run: |
         rm -rf build
-        cmake -Bbuild -DTEST_LCOV=ON -DBPFTIME_LLVM_JIT=YES  -DBPFTIME_ENABLE_UNIT_TESTING=YES -DBPFTIME_ENABLE_MPK=YES -DCMAKE_BUILD_TYPE=Debug
+        cmake -Bbuild -DTEST_LCOV=ON -DBPFTIME_LLVM_JIT=YES  -DBPFTIME_ENABLE_UNIT_TESTING=YES -DBPFTIME_ENABLE_MPK=YES -DCMAKE_BUILD_TYPE=Debug -DENABLE_PROBE_WRITE_CHECK=1 -DENABLE_PROBE_READ_CHECK=1
         cmake --build build --config Debug --target bpftime_runtime_tests -j$(nproc)
     - name: test runtime with mpk
       run:  ./build/runtime/unit-test/bpftime_runtime_tests

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -160,6 +160,15 @@ add_subdirectory(attach)
 
 add_subdirectory(runtime)
 
+# option to use probe read/write checks
+if(DEFINED ENABLE_PROBE_READ_CHECK)
+  target_compile_definitions(runtime PRIVATE ENABLE_PROBE_READ_CHECK=1)
+endif()
+
+if(DEFINED ENABLE_PROBE_WRITE_CHECK)
+  target_compile_definitions(runtime PRIVATE ENABLE_PROBE_WRITE_CHECK=1)
+endif()
+
 # add to single archive if option is enabled
 if(${BPFTIME_BUILD_STATIC_LIB})
   set(UBPF_BUILD_DIR ${CMAKE_CURRENT_BINARY_DIR}/vm/ubpf-vm/ubpf)

diff --git a/Makefile b/Makefile
@@ -33,6 +33,10 @@ help:
 	@python3 -c "$$PRINT_HELP_PYSCRIPT" < $(MAKEFILE_LIST)
 
 build-unit-test:
+	cmake -Bbuild  -DBPFTIME_ENABLE_UNIT_TESTING=1 -DCMAKE_BUILD_TYPE:STRING=Debug -DENABLE_PROBE_WRITE_CHECK=1 -DENABLE_PROBE_READ_CHECK=1
+	cmake --build build --config Debug --target bpftime_runtime_tests bpftime_daemon_tests  -j$(JOBS)
+
+build-unit-test-without-probe-check:
 	cmake -Bbuild  -DBPFTIME_ENABLE_UNIT_TESTING=1 -DCMAKE_BUILD_TYPE:STRING=Debug
 	cmake --build build --config Debug --target bpftime_runtime_tests bpftime_daemon_tests  -j$(JOBS)
 

diff --git a/benchmark/README.md b/benchmark/README.md
@@ -30,6 +30,8 @@ cmake -Bbuild -DLLVM_DIR=/usr/lib/llvm-15/cmake -DCMAKE_BUILD_TYPE:STRING=RelWit
 cmake --build build --config RelWithDebInfo --target install -j
 ```
 
+If you fail to build , notice LLVM version.
+
 ## build and run at a click
 
 Build the agent first. In project root:
@@ -195,7 +197,7 @@ Tested on `kernel version 6.2` and `Intel(R) Xeon(R) Gold 5418Y` CPU.
 
 ### Uprobe and read/write with `bpf_probe_write_user` and `bpf_probe_read_user`
 
-Userspace:
+Kernelspace:
 
 ```txt
 Benchmarking __bench_uprobe_uretprobe in thread 1
@@ -233,6 +235,8 @@ Benchmarking __bench_write in thread 1
 Average time usage 389.037170 ns, iter 100000 times
 ```
 
+noted that the performance of `bpf_probe_read_user` and `bpf_probe_write_user` will be influenced by the option of `ENABLE_PROBE_WRITE_USER` and `ENABLE_PROBE_READ_USER`.
+
 ### maps operations
 
 Run the map op 1000 times in one function. Userspace map op is also faster than the kernel in the current version. Current version is 10x faster than stupid old version.

diff --git a/cmake/StandardSettings.cmake b/cmake/StandardSettings.cmake
@@ -87,4 +87,10 @@ option(BPFTIME_BUILD_KERNEL_BPF "Whether to build with bpf share maps" ON)
 option(BPFTIME_BUILD_STATIC_LIB "Whether to build a single static library for different archive files" OFF)
 
 # whether to build bpftime with libbpf and other linux headers
-option(BPFTIME_BUILD_WITH_LIBBPF "Whether to build with libbpf and other linux headers" ON)
+option(BPFTIME_BUILD_WITH_LIBBPF "Whether to build with libbpf and other linux headers" ON)
+
+# whether to enable probe write check
+option(ENABLE_PROBE_WRITE_CHECK "Whether to enable probe write check" OFF)
+
+# whether to enable probe read check
+option(ENABLE_PROBE_READ_CHECK "Whether to enable probe read check" OFF)
diff --git a/runtime/src/bpf_helper.cpp b/runtime/src/bpf_helper.cpp
@@ -3,6 +3,8 @@
  * Copyright (c) 2022, eunomia-bpf org
  * All rights reserved.
  */
+#include <stdexcept>
+#include <system_error>
 #if __APPLE__
 #include <cstdint>
 #include <pthread.h>
@@ -35,6 +37,9 @@
 #include <vector>
 #include <bpftime_shm_internal.hpp>
 #include <chrono>
+#include <thread>
+#include <setjmp.h>
+#include <signal.h>
 
 #define PATH_MAX 4096
 
@@ -64,20 +69,222 @@ long bpftime_strncmp(const char *s1, uint64_t s1_sz, const char *s2)
 	return strncmp(s1, s2, s1_sz);
 }
 
-uint64_t bpftime_probe_read(uint64_t dst, uint64_t size, uint64_t ptr, uint64_t,
-			    uint64_t)
+
+#if defined(ENABLE_PROBE_WRITE_CHECK) || defined(ENABLE_PROBE_READ_CHECK)
+
+/*
+status instruction for probe_read and probe_write
+*/
+enum class PROBE_STATUS {
+	NOT_RUNNING = -1,
+	RUNNING_NO_ERROR = 0,
+	RUNNING_ERROR = 1
+};
+
+/*
+origin handler exist flag for probe_read and probe_write
+*/
+enum class ORIGIN_HANDLER_EXIST_FLAG {
+	NOT_CHECKED = -1,
+	NOT_EXIST = 0,
+	EXIST = 1
+};
+
+#endif
+
+#ifdef ENABLE_PROBE_READ_CHECK
+extern "C" void jump_point_read();
+
+thread_local static PROBE_STATUS status_probe_write = PROBE_STATUS::NOT_RUNNING;
+
+thread_local static ORIGIN_HANDLER_EXIST_FLAG exist_read =
+	ORIGIN_HANDLER_EXIST_FLAG::NOT_CHECKED;
+
+thread_local static void (*origin_segv_read_handler)(int, siginfo_t *,
+						     void *) = nullptr;
+#endif
+#ifdef ENABLE_PROBE_WRITE_CHECK
+
+extern "C" void jump_point_write();
+thread_local static PROBE_STATUS status_probe_read = PROBE_STATUS::NOT_RUNNING;
+
+thread_local static ORIGIN_HANDLER_EXIST_FLAG exist_write =
+	ORIGIN_HANDLER_EXIST_FLAG::NOT_CHECKED;
+
+thread_local static void (*origin_segv_write_handler)(int, siginfo_t *,
+						      void *) = nullptr;
+#endif
+
+#ifdef ENABLE_PROBE_READ_CHECK
+static void segv_read_handler(int sig, siginfo_t *siginfo, void *ctx)
 {
-	memcpy((void *)(uintptr_t)dst, (void *)(uintptr_t)ptr,
-	       (size_t)(uint32_t)(size));
-	return 0;
+	if (status_probe_read == PROBE_STATUS::NOT_RUNNING) {
+		if (origin_segv_read_handler != nullptr) {
+			origin_segv_read_handler(sig, siginfo, ctx);
+		} else {
+			SPDLOG_ERROR("no origin handler for probe_read");
+			throw std::runtime_error(
+				"segv_handler for probe_read called");
+		}
+	} else if (status_probe_read == PROBE_STATUS::RUNNING_NO_ERROR) {
+		// set status to error
+		auto uctx = (ucontext_t *)ctx;
+		auto *rip = (greg_t *)(&uctx->uc_mcontext.gregs[REG_RIP]);
+		status_probe_read = PROBE_STATUS::RUNNING_ERROR;
+		*rip = (greg_t)&jump_point_read;
+	}
 }
+#endif
 
-uint64_t bpftime_probe_write_user(uint64_t dst, uint64_t src, uint64_t len,
-				  uint64_t, uint64_t)
+int64_t bpftime_probe_read(uint64_t dst, int64_t size, uint64_t ptr, uint64_t,
+			   uint64_t)
 {
-	memcpy((void *)(uintptr_t)dst, (void *)(uintptr_t)src,
-	       (size_t)(uint32_t)(len));
-	return 0;
+	if (size < 0) {
+		SPDLOG_ERROR("Invalid size: {}", size);
+		return -EFAULT;
+	}
+	int64_t ret = 0;
+
+#ifdef ENABLE_PROBE_READ_CHECK
+	status_probe_read = PROBE_STATUS::RUNNING_NO_ERROR;
+
+	struct sigaction sa, original_sa;
+	// set up the signal handler
+	if (exist_read == ORIGIN_HANDLER_EXIST_FLAG::NOT_CHECKED) {
+		int err = sigaction(SIGSEGV, nullptr, &original_sa);
+		if (err) {
+			SPDLOG_ERROR("Failed to get signal handler: {}", errno);
+			return -EFAULT;
+		}
+		if (original_sa.sa_sigaction == nullptr) {
+			exist_read = ORIGIN_HANDLER_EXIST_FLAG::NOT_EXIST;
+		} else {
+			exist_read = ORIGIN_HANDLER_EXIST_FLAG::EXIST;
+			origin_segv_read_handler = original_sa.sa_sigaction;
+		}
+	}
+
+	if (original_sa.sa_sigaction != segv_read_handler) {
+		sa.sa_flags = SA_SIGINFO;
+		int err = 0;
+		err = sigemptyset(&sa.sa_mask);
+		if (err) {
+			SPDLOG_ERROR("Failed to set signal handler: {}", errno);
+			return -EFAULT;
+		}
+		sa.sa_sigaction = segv_read_handler;
+		err = sigaction(SIGSEGV, &sa, nullptr);
+		if (err) {
+			SPDLOG_ERROR("Failed to set signal handler: {}", errno);
+			return -EFAULT;
+		}
+	}
+#endif
+	unsigned char *dst_p = (unsigned char *)dst;
+	unsigned char *src_p = (unsigned char *)ptr;
+	while (size--) {
+		*((unsigned char *)dst_p) = *((unsigned char *)src_p);
+		dst_p++;
+		src_p++;
+	}
+	__asm__("jump_point_read:");
+
+#ifdef ENABLE_PROBE_READ_CHECK
+	if (status_probe_read == PROBE_STATUS::RUNNING_ERROR) {
+		ret = -EFAULT;
+	}
+
+	status_probe_read = PROBE_STATUS::NOT_RUNNING;
+#endif
+	return ret;
+}
+
+#ifdef ENABLE_PROBE_WRITE_CHECK
+static void segv_write_handler(int sig, siginfo_t *siginfo, void *ctx)
+{
+	SPDLOG_TRACE("segv_handler for probe_write called");
+	if (status_probe_write == PROBE_STATUS::NOT_RUNNING) {
+		if (origin_segv_write_handler) {
+			origin_segv_write_handler(sig, siginfo, ctx);
+		} else {
+			SPDLOG_ERROR("no origin handler for probe_write");
+			throw std::runtime_error(
+				"segv_handler for probe_write called");
+		}
+	} else if (status_probe_write == PROBE_STATUS::RUNNING_NO_ERROR) {
+		// set status to error
+		auto uctx = (ucontext_t *)ctx;
+		auto *rip = (greg_t *)(&uctx->uc_mcontext.gregs[REG_RIP]);
+		status_probe_write = PROBE_STATUS::RUNNING_ERROR;
+		*rip = (greg_t)&jump_point_write;
+	}
+}
+#endif
+
+int64_t bpftime_probe_write_user(uint64_t dst, uint64_t src, int64_t len,
+				 uint64_t, uint64_t)
+{
+	if (len < 0) {
+		SPDLOG_ERROR("Invalid len: {}", len);
+		return -EFAULT;
+	}
+	int64_t ret = 0;
+
+#ifdef ENABLE_PROBE_WRITE_CHECK
+	status_probe_write = PROBE_STATUS::RUNNING_NO_ERROR;
+
+	struct sigaction sa, original_sa;
+	// set up the signal handler
+	if (exist_write == ORIGIN_HANDLER_EXIST_FLAG::NOT_CHECKED) {
+		int err = sigaction(SIGSEGV, nullptr, &original_sa);
+		if (err) {
+			SPDLOG_ERROR("Failed to get signal handler: {}", errno);
+			return -EFAULT;
+		}
+
+		if (original_sa.sa_sigaction == nullptr) {
+			exist_write = ORIGIN_HANDLER_EXIST_FLAG::NOT_EXIST;
+		} else {
+			exist_write = ORIGIN_HANDLER_EXIST_FLAG::EXIST;
+			origin_segv_write_handler = original_sa.sa_sigaction;
+		}
+	}
+
+	if (original_sa.sa_sigaction != segv_write_handler) {
+		sa.sa_flags = SA_SIGINFO;
+		int err = 0;
+		err = sigemptyset(&sa.sa_mask);
+		if (err) {
+			SPDLOG_ERROR("Failed to set signal handler: {}", errno);
+			return -EFAULT;
+		}
+
+		sa.sa_sigaction = segv_write_handler;
+		err = sigaction(SIGSEGV, &sa, nullptr);
+		if (err) {
+			SPDLOG_ERROR("Failed to set signal handler: {}", errno);
+			return -EFAULT;
+		}
+	}
+#endif
+	unsigned char *dst_p = (unsigned char *)dst;
+	unsigned char *src_p = (unsigned char *)src;
+	while (len--) {
+		*((unsigned char *)dst_p) = *((unsigned char *)src_p);
+		dst_p++;
+		src_p++;
+	}
+
+	__asm__("jump_point_write:");
+#ifdef ENABLE_PROBE_WRITE_CHECK
+	if (status_probe_write == PROBE_STATUS::RUNNING_ERROR) {
+		ret = -EFAULT;
+	}
+
+	status_probe_write = PROBE_STATUS::NOT_RUNNING;
+
+#endif
+	return ret;
 }
 
 uint64_t bpftime_get_prandom_u32()

diff --git a/runtime/unit-test/CMakeLists.txt b/runtime/unit-test/CMakeLists.txt
@@ -22,6 +22,7 @@ set(TEST_SOURCES
     maps/kernel_unit_tests.cpp
 
     test_bpftime_shm_json.cpp
+    test_probe.cpp
     attach_with_ebpf/test_attach_filter_with_ebpf.cpp
     attach_with_ebpf/test_attach_uprobe_with_ebpf.cpp
     attach_with_ebpf/test_helpers.cpp