Merge branch 'fix-issue-268' of https://github.com/NobinPegasus/bpftime…

… into fix-issue-268

attach/frida_uprobe_attach_impl/src/frida_internal_attach_entry.cpp

@@ -79,6 +79,7 @@ frida_internal_attach_entry::~frida_internal_attach_entry()
 		SPDLOG_DEBUG("Reverted function replace");
 	}
 	gum_object_unref(interceptor);
+	SPDLOG_DEBUG("Destructor of frida_internal_attach_entry exiting..");
 }
 
 bool frida_internal_attach_entry::has_override() const

attach/text_segment_transformer/agent-transformer.cpp

@@ -13,6 +13,10 @@ using shm_destroy_func_t = void (*)(void);
 
 static main_func_t orig_main_func = nullptr;
 static shm_destroy_func_t shm_destroy_func = nullptr;
+
+// Whether syscall server was injected using frida. Defaults to true. If
+// __libc_start_main was called, it will be set to false
+static bool injected_with_frida = true;
 extern "C" void bpftime_agent_main(const gchar *data, gboolean *stay_resident);
 
 extern "C" int bpftime_hooked_main(int argc, char **argv, char **envp)
@@ -29,6 +33,7 @@ extern "C" int __libc_start_main(int (*main)(int, char **, char **), int argc,
 				 void (*fini)(void), void (*rtld_fini)(void),
 				 void *stack_end)
 {
+	injected_with_frida = false;
 	SPDLOG_INFO("Entering bpftime syscal transformer agent");
 	orig_main_func = main;
 	using this_func_t = decltype(&__libc_start_main);
@@ -70,6 +75,14 @@ extern "C" void bpftime_agent_main(const gchar *data, gboolean *stay_resident)
 		SPDLOG_ERROR("Failed to open agent: {}", dlerror());
 		exit(1);
 	}
+	// Set the flag `injected_with_frida` for agent
+	bool *injected_with_frida__agent =
+		(bool *)dlsym(next_handle, "injected_with_frida");
+	if (!injected_with_frida__agent) {
+		SPDLOG_WARN(
+			"Agent does not expose a symbol named injected_with_frida, so we can't let agent know whether it was loaded using frida");
+	}
+	*injected_with_frida__agent = injected_with_frida;
 	auto entry_func = (void (*)(syscall_hooker_func_t *))dlsym(
 		next_handle, "_bpftime__setup_syscall_trace_callback");
 

runtime/agent/agent.cpp

@@ -1,14 +1,19 @@
 #include "attach_private_data.hpp"
 #include "bpf_attach_ctx.hpp"
+#include "bpftime_shm_internal.hpp"
 #include "frida_attach_private_data.hpp"
 #include "frida_uprobe_attach_impl.hpp"
 #include "spdlog/common.h"
 #include "syscall_trace_attach_impl.hpp"
 #include "syscall_trace_attach_private_data.hpp"
+#include <chrono>
+#include <csignal>
 #include <exception>
 #include <fcntl.h>
 #include <memory>
+#include <pthread.h>
 #include <string_view>
+#include <thread>
 #include <unistd.h>
 #include <frida-gum.h>
 #include <cstdint>
@@ -22,6 +27,12 @@ using main_func_t = int (*)(int, char **, char **);
 
 static main_func_t orig_main_func = nullptr;
 
+// Whether this injected process was operated through frida?
+// Defaults to true. If __libc_start_main was called, it should be set to false;
+// Besides, if agent was loaded by text-transformer, this variable will be set
+// by text-transformer
+bool injected_with_frida = true;
+
 union bpf_attach_ctx_holder {
 	bpf_attach_ctx ctx;
 	bpf_attach_ctx_holder()
@@ -48,7 +59,7 @@ extern "C" void bpftime_agent_main(const gchar *data, gboolean *stay_resident);
 extern "C" int bpftime_hooked_main(int argc, char **argv, char **envp)
 {
 	int stay_resident = 0;
-
+	injected_with_frida = false;
 	bpftime_agent_main("", &stay_resident);
 	int ret = orig_main_func(argc, argv, envp);
 	ctx_holder.destroy();
@@ -70,8 +81,22 @@ extern "C" int __libc_start_main(int (*main)(int, char **, char **), int argc,
 		    stack_end);
 }
 
+static void sig_handler_sigusr1(int sig)
+{
+	SPDLOG_INFO("Detaching..");
+	if (int err = ctx_holder.ctx.destroy_all_attach_links(); err < 0) {
+		SPDLOG_ERROR("Unable to detach: {}", err);
+		return;
+	}
+	shm_holder.global_shared_memory.remove_pid_from_alive_agent_set(getpid());
+	SPDLOG_DEBUG("Detaching done");
+}
+
 extern "C" void bpftime_agent_main(const gchar *data, gboolean *stay_resident)
 {
+	SPDLOG_DEBUG("Registering signal handler");
+	// We use SIGUSR1 to indicate the detaching
+	signal(SIGUSR1, sig_handler_sigusr1);
 	spdlog::cfg::load_env_levels();
 	try {
 		// If we are unable to initialize shared memory..
@@ -81,6 +106,12 @@ extern "C" void bpftime_agent_main(const gchar *data, gboolean *stay_resident)
 			     ex.what());
 		return;
 	}
+	// Only agents injected through frida could be detached
+	if (injected_with_frida) {
+		// Record the pid
+		shm_holder.global_shared_memory.add_pid_into_alive_agent_set(
+			getpid());
+	}
 	ctx_holder.init();
 	// Register syscall trace impl
 	auto syscall_trace_impl = std::make_unique<syscall_trace_attach_impl>();

runtime/include/bpf_attach_ctx.hpp

@@ -52,6 +52,8 @@ class bpf_attach_ctx {
 			private_data_creator);
 	// Destroy a specific attach link
 	int destroy_instantiated_attach_link(int link_id);
+	// Destroy all instantiated attach links
+	int destroy_all_attach_links();
 
     private:
 	constexpr static int CURRENT_ID_OFFSET = 65536;
@@ -64,7 +66,7 @@ class bpf_attach_ctx {
 	std::map<int, std::unique_ptr<bpftime_prog> > instantiated_progs;
 	// handler_id -> (instantiated attaches id, attach_impl*)
 	std::map<int, std::pair<int, attach::base_attach_impl *> >
-		instantiated_attach_ids;
+		instantiated_attach_links;
 	// handler_id -> instantiated attach private data & attach type
 	std::map<int,
 		 std::pair<std::unique_ptr<attach::attach_private_data>, int> >

runtime/src/attach/bpf_attach_ctx.cpp

@@ -236,7 +236,7 @@ int bpf_attach_ctx::instantiate_bpf_link_handler_at(
 			     id, attach_id);
 		return attach_id;
 	}
-	instantiated_attach_ids[id] = std::make_pair(attach_id, attach_impl);
+	instantiated_attach_links[id] = std::make_pair(attach_id, attach_impl);
 	return 0;
 }
 int bpf_attach_ctx::instantiate_perf_event_handler_at(
@@ -315,21 +315,37 @@ int bpf_attach_ctx::instantiate_perf_event_handler_at(
 int bpf_attach_ctx::destroy_instantiated_attach_link(int link_id)
 {
 	SPDLOG_DEBUG("Destroy attach link {}", link_id);
-	if (auto itr = instantiated_attach_ids.find(link_id);
-	    itr != instantiated_attach_ids.end()) {
+	if (auto itr = instantiated_attach_links.find(link_id);
+	    itr != instantiated_attach_links.end()) {
 		auto [attach_id, impl] = itr->second;
 		if (int err = impl->detach_by_id(attach_id); err < 0) {
 			SPDLOG_ERROR(
 				"Failed to detach attach link id {}, attach-specified id {}: {}",
 				link_id, attach_id, err);
 			return err;
 		}
-		instantiated_attach_ids.erase(itr);
+		instantiated_attach_links.erase(itr);
 		return 0;
 	} else {
 		SPDLOG_ERROR("Unable to find instantiated attach link id {}",
 			     link_id);
 		return -ENOENT;
 	}
 }
+int bpf_attach_ctx::destroy_all_attach_links()
+{
+	// Avoid modifying along with iterating..
+	std::vector<int> to_detach;
+	for (const auto &[k, _] : instantiated_attach_links)
+		to_detach.push_back(k);
+	for (auto k : to_detach) {
+		SPDLOG_DEBUG("Destrying attach link {}", k);
+		if (int err = destroy_instantiated_attach_link(k); err < 0) {
+			SPDLOG_ERROR("Unable to destroy attach link {}: {}", k,
+				     err);
+			return err;
+		}
+	}
+	return 0;
+}
 } // namespace bpftime

runtime/src/bpftime_shm_internal.cpp

@@ -5,6 +5,7 @@
  */
 #include "bpftime_shm.hpp"
 #include "handler/epoll_handler.hpp"
+#include "handler/handler_manager.hpp"
 #include "handler/link_handler.hpp"
 #include "handler/perf_event_handler.hpp"
 #include "spdlog/spdlog.h"
@@ -51,6 +52,17 @@ extern "C" void bpftime_remove_global_shm()
 
 static __attribute__((destructor(65535))) void __destruct_shm()
 {
+	// This usually indicates that the living shared memory object is used
+	// by an agent instance
+	if (bpftime::shm_holder.global_shared_memory.get_open_type() ==
+	    bpftime::shm_open_type::SHM_OPEN_ONLY) {
+		// Try our best to remove the current pid from alive agent's set
+		int self_pid = getpid();
+		// It doesn't matter if the current pid is not in the set
+		bpftime::shm_holder.global_shared_memory
+			.remove_pid_from_alive_agent_set(self_pid);
+	}
+
 	bpftime_destroy_global_shm();
 }
 
@@ -527,6 +539,7 @@ bool bpftime_shm::is_exist_fake_fd(int fd) const
 }
 
 bpftime_shm::bpftime_shm(const char *shm_name, shm_open_type type)
+	: open_type(type)
 {
 	if (type == shm_open_type::SHM_OPEN_ONLY) {
 		SPDLOG_DEBUG("start: bpftime_shm for client setup");
@@ -544,6 +557,11 @@ bpftime_shm::bpftime_shm(const char *shm_name, shm_open_type type)
 			segment.find<struct agent_config>(
 				       bpftime::DEFAULT_AGENT_CONFIG_NAME)
 				.first;
+
+		injected_pids =
+			segment.find<alive_agent_pids>(
+				       bpftime::DEFAULT_ALIVE_AGENT_PIDS_NAME)
+				.first;
 		SPDLOG_DEBUG("done: bpftime_shm for client setup");
 	} else if (type == shm_open_type::SHM_CREATE_OR_OPEN) {
 		SPDLOG_DEBUG("start: bpftime_shm for create or open setup");
@@ -567,6 +585,12 @@ bpftime_shm::bpftime_shm(const char *shm_name, shm_open_type type)
 
 		agent_config = segment.find_or_construct<struct agent_config>(
 			bpftime::DEFAULT_AGENT_CONFIG_NAME)();
+
+		injected_pids = segment.find_or_construct<alive_agent_pids>(
+			bpftime::DEFAULT_ALIVE_AGENT_PIDS_NAME)(
+			std::less<int>(),
+			alive_agent_pid_set_allocator(
+				segment.get_segment_manager()));
 		SPDLOG_DEBUG("done: bpftime_shm for open_or_create setup");
 	} else if (type == shm_open_type::SHM_REMOVE_AND_CREATE) {
 		SPDLOG_DEBUG("start: bpftime_shm for server setup");
@@ -596,6 +620,12 @@ bpftime_shm::bpftime_shm(const char *shm_name, shm_open_type type)
 			bpftime::DEFAULT_AGENT_CONFIG_NAME)();
 		SPDLOG_DEBUG(
 			"done: bpftime_shm for server setup: agent_config");
+
+		injected_pids = segment.construct<alive_agent_pids>(
+			bpftime::DEFAULT_ALIVE_AGENT_PIDS_NAME)(
+			std::less<int>(),
+			alive_agent_pid_set_allocator(
+				segment.get_segment_manager()));
 		SPDLOG_DEBUG("done: bpftime_shm for server setup.");
 	} else if (type == shm_open_type::SHM_NO_CREATE) {
 		// not create any shm
@@ -735,4 +765,20 @@ int bpftime_shm::add_custom_perf_event(int type, const char *attach_argument)
 		segment);
 	return fd;
 }
+
+void bpftime_shm::add_pid_into_alive_agent_set(int pid)
+{
+	injected_pids->insert(pid);
+}
+void bpftime_shm::remove_pid_from_alive_agent_set(int pid)
+{
+	injected_pids->erase(pid);
+}
+void bpftime_shm::iterate_all_pids_in_alive_agent_set(
+	std::function<void(int)> &&cb)
+{
+	for (auto x : *injected_pids) {
+		cb(x);
+	}
+}
 } // namespace bpftime

runtime/src/bpftime_shm_internal.hpp

@@ -23,8 +23,16 @@ using syscall_pid_set_allocator = boost::interprocess::allocator<
 using syscall_pid_set =
 	boost::interprocess::set<int, std::less<int>, syscall_pid_set_allocator>;
 
+using alive_agent_pid_set_allocator = boost::interprocess::allocator<
+	int, boost::interprocess::managed_shared_memory::segment_manager>;
+
+using alive_agent_pids =
+	boost::interprocess::set<int, std::less<int>,
+				 alive_agent_pid_set_allocator>;
+
 // global bpftime share memory
 class bpftime_shm {
+	bpftime::shm_open_type open_type;
 	// shared memory segment
 	boost::interprocess::managed_shared_memory segment;
 
@@ -37,6 +45,9 @@ class bpftime_shm {
 	// Configuration for the agent. e.g, which helpers are enabled
 	struct bpftime::agent_config *agent_config = nullptr;
 
+	// Record which pids are injected by agent
+	alive_agent_pids *injected_pids;
+
 #if BPFTIME_ENABLE_MPK
 	// mpk key for protect shm
 	bool is_mpk_init = false;
@@ -55,6 +66,13 @@ class bpftime_shm {
 	// Using a set stored in the shared memory
 	void set_syscall_trace_setup(int pid, bool whether);
 
+	// Add a pid into alive agent set
+	void add_pid_into_alive_agent_set(int pid);
+	// Remove a pid from alive agent set
+	void remove_pid_from_alive_agent_set(int pid);
+	// Iterate over all pids from the alive agent set
+	void iterate_all_pids_in_alive_agent_set(std::function<void(int)> &&cb);
+
 	const handler_variant &get_handler(int fd) const;
 	bool is_epoll_fd(int fd) const;
 
@@ -163,6 +181,10 @@ class bpftime_shm {
 	get_software_perf_event_raw_buffer(int fd, size_t buffer_sz) const;
 
 	int add_custom_perf_event(int type, const char *attach_argument);
+	bpftime::shm_open_type get_open_type() const
+	{
+		return open_type;
+	}
 };
 
 // memory region for maps and prog info

runtime/src/handler/handler_manager.hpp

@@ -53,6 +53,7 @@ constexpr const char *DEFAULT_GLOBAL_SHM_NAME = "bpftime_maps_shm";
 constexpr const char *DEFAULT_GLOBAL_HANDLER_NAME = "bpftime_handler";
 constexpr const char *DEFAULT_SYSCALL_PID_SET_NAME = "bpftime_syscall_pid_set";
 constexpr const char *DEFAULT_AGENT_CONFIG_NAME = "bpftime_agent_config";
+constexpr const char* DEFAULT_ALIVE_AGENT_PIDS_NAME = "bpftime_alive_agent_pids";
 inline const char *get_global_shm_name()
 {
 	const char *name = getenv("BPFTIME_GLOBAL_SHM_NAME");

tools/cli/CMakeLists.txt

@@ -5,12 +5,12 @@ add_executable(
 
 set_target_properties(bpftime-cli-cpp PROPERTIES OUTPUT_NAME "bpftime")
 
-target_include_directories(bpftime-cli-cpp PRIVATE ${FRIDA_CORE_INSTALL_DIR} ${SPDLOG_INCLUDE} ${argparse_INCLUDE})
-target_link_libraries(bpftime-cli-cpp PRIVATE spdlog::spdlog ${FRIDA_CORE_INSTALL_DIR}/libfrida-core.a argparse)
+target_include_directories(bpftime-cli-cpp PRIVATE ${FRIDA_CORE_INSTALL_DIR} ${SPDLOG_INCLUDE} ${argparse_INCLUDE} ../../runtime/include)
+target_link_libraries(bpftime-cli-cpp PRIVATE spdlog::spdlog ${FRIDA_CORE_INSTALL_DIR}/libfrida-core.a argparse runtime)
 set_property(TARGET bpftime-cli-cpp PROPERTY CXX_STANDARD 20)
 
 target_compile_definitions(bpftime-cli-cpp PRIVATE _GNU_SOURCE)
 
-add_dependencies(bpftime-cli-cpp spdlog::spdlog FridaCore argparse)
+add_dependencies(bpftime-cli-cpp spdlog::spdlog FridaCore argparse runtime)
 
 install(TARGETS bpftime-cli-cpp CONFIGURATIONS Release Debug RelWithDebInfo DESTINATION ~/.bpftime)