feat: support fallback probes when a probe attach fails

Signed-off-by: Cameron Hall <[email protected]>
pixie-io · Jan 6, 2024 · 3c44477 · 3c44477
1 parent 2e26307
commit 3c44477
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 45 deletions.
diff --git a/src/stirling/bpf_tools/bcc_wrapper.cc b/src/stirling/bpf_tools/bcc_wrapper.cc
@@ -181,6 +181,11 @@ Status BCCWrapperImpl::AttachKProbe(const KProbeSpec& probe) {
       bpf_.attach_kprobe(GetKProbeTargetName(probe), std::string(probe.probe_fn), 0 /* offset */,
                          static_cast<bpf_probe_attach_type>(probe.attach_type), kKprobeMaxActive);
 
+  if (!status.ok() && probe.fallback_probe != nullptr) {
+    VLOG(1) << "kprobe attach failed... attempting fallback.";
+    return BCCWrapperImpl::AttachKProbe(*probe.fallback_probe);
+  }
+
   // Don't return error if the probe is optional.
   if (!probe.is_optional) {
     PX_RETURN_IF_ERROR(status);

diff --git a/src/stirling/bpf_tools/probe_specs/probe_specs.h b/src/stirling/bpf_tools/probe_specs/probe_specs.h
@@ -75,6 +75,9 @@ struct KProbeSpec {
   // kernels.
   bool is_optional = false;
 
+  // If the kernel function is not found, then this fallback function will be used instead.
+  struct KProbeSpec* fallback_probe = nullptr;
+
   std::string ToString() const {
     return absl::Substitute("[kernel_function=$0 type=$1 probe=$2]", kernel_fn,
                             magic_enum::enum_name(attach_type), probe_fn);

diff --git a/src/stirling/source_connectors/socket_tracer/bcc_bpf/socket_trace.c b/src/stirling/source_connectors/socket_tracer/bcc_bpf/socket_trace.c
@@ -1627,8 +1627,9 @@ int probe_ret_sock_alloc(struct pt_regs* ctx) {
 
 // Trace kernel function:
 // int security_socket_sendmsg(struct socket *sock, struct msghdr *msg, int size)
+// int sock_sendmsg(struct socket *sock, struct msghdr *msg)
 // which is called by write/writev/send/sendmsg.
-int probe_entry_security_socket_sendmsg(struct pt_regs* ctx) {
+int probe_entry_socket_sendmsg(struct pt_regs* ctx) {
   uint64_t id = bpf_get_current_pid_tgid();
   struct data_args_t* write_args = active_write_args_map.lookup(&id);
   if (write_args != NULL) {
@@ -1639,7 +1640,8 @@ int probe_entry_security_socket_sendmsg(struct pt_regs* ctx) {
 
 // Trace kernel function:
 // int security_socket_recvmsg(struct socket *sock, struct msghdr *msg, int size)
-int probe_entry_security_socket_recvmsg(struct pt_regs* ctx) {
+// int sock_recvmsg(struct socket *sock, struct msghdr *msg, int flags)
+int probe_entry_socket_recvmsg(struct pt_regs* ctx) {
   uint64_t id = bpf_get_current_pid_tgid();
   struct data_args_t* read_args = active_read_args_map.lookup(&id);
   if (read_args != NULL) {

diff --git a/src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc b/src/stirling/source_connectors/socket_tracer/socket_trace_connector.cc
@@ -293,49 +293,54 @@ void SocketTraceConnector::InitProtocolTransferSpecs() {
 }
 
 using ProbeType = bpf_tools::BPFProbeAttachType;
-const auto kProbeSpecs = MakeArray<bpf_tools::KProbeSpec>(
-    {{"connect", ProbeType::kEntry, "syscall__probe_entry_connect"},
-     {"connect", ProbeType::kReturn, "syscall__probe_ret_connect"},
-     {"accept", ProbeType::kEntry, "syscall__probe_entry_accept"},
-     {"accept", ProbeType::kReturn, "syscall__probe_ret_accept"},
-     {"accept4", ProbeType::kEntry, "syscall__probe_entry_accept4"},
-     {"accept4", ProbeType::kReturn, "syscall__probe_ret_accept4"},
-     {"write", ProbeType::kEntry, "syscall__probe_entry_write"},
-     {"write", ProbeType::kReturn, "syscall__probe_ret_write"},
-     {"writev", ProbeType::kEntry, "syscall__probe_entry_writev"},
-     {"writev", ProbeType::kReturn, "syscall__probe_ret_writev"},
-     {"send", ProbeType::kEntry, "syscall__probe_entry_send"},
-     {"send", ProbeType::kReturn, "syscall__probe_ret_send"},
-     {"sendto", ProbeType::kEntry, "syscall__probe_entry_sendto"},
-     {"sendto", ProbeType::kReturn, "syscall__probe_ret_sendto"},
-     {"sendmsg", ProbeType::kEntry, "syscall__probe_entry_sendmsg"},
-     {"sendmsg", ProbeType::kReturn, "syscall__probe_ret_sendmsg"},
-     {"sendmmsg", ProbeType::kEntry, "syscall__probe_entry_sendmmsg"},
-     {"sendmmsg", ProbeType::kReturn, "syscall__probe_ret_sendmmsg"},
-     {"sendfile", ProbeType::kEntry, "syscall__probe_entry_sendfile"},
-     {"sendfile", ProbeType::kReturn, "syscall__probe_ret_sendfile"},
-     {"sendfile64", ProbeType::kEntry, "syscall__probe_entry_sendfile"},
-     {"sendfile64", ProbeType::kReturn, "syscall__probe_ret_sendfile"},
-     {"read", ProbeType::kEntry, "syscall__probe_entry_read"},
-     {"read", ProbeType::kReturn, "syscall__probe_ret_read"},
-     {"readv", ProbeType::kEntry, "syscall__probe_entry_readv"},
-     {"readv", ProbeType::kReturn, "syscall__probe_ret_readv"},
-     {"recv", ProbeType::kEntry, "syscall__probe_entry_recv"},
-     {"recv", ProbeType::kReturn, "syscall__probe_ret_recv"},
-     {"recvfrom", ProbeType::kEntry, "syscall__probe_entry_recvfrom"},
-     {"recvfrom", ProbeType::kReturn, "syscall__probe_ret_recvfrom"},
-     {"recvmsg", ProbeType::kEntry, "syscall__probe_entry_recvmsg"},
-     {"recvmsg", ProbeType::kReturn, "syscall__probe_ret_recvmsg"},
-     {"recvmmsg", ProbeType::kEntry, "syscall__probe_entry_recvmmsg"},
-     {"recvmmsg", ProbeType::kReturn, "syscall__probe_ret_recvmmsg"},
-     {"close", ProbeType::kEntry, "syscall__probe_entry_close"},
-     {"close", ProbeType::kReturn, "syscall__probe_ret_close"},
-     {"mmap", ProbeType::kEntry, "syscall__probe_entry_mmap"},
-     {"sock_alloc", ProbeType::kReturn, "probe_ret_sock_alloc", /*is_syscall*/ false},
-     {"security_socket_sendmsg", ProbeType::kEntry, "probe_entry_security_socket_sendmsg",
-      /*is_syscall*/ false},
-     {"security_socket_recvmsg", ProbeType::kEntry, "probe_entry_security_socket_recvmsg",
-      /*is_syscall*/ false}});
+const auto kProbeSpecs = MakeArray<bpf_tools::KProbeSpec>({
+    {"connect", ProbeType::kEntry, "syscall__probe_entry_connect"},
+    {"connect", ProbeType::kReturn, "syscall__probe_ret_connect"},
+    {"accept", ProbeType::kEntry, "syscall__probe_entry_accept"},
+    {"accept", ProbeType::kReturn, "syscall__probe_ret_accept"},
+    {"accept4", ProbeType::kEntry, "syscall__probe_entry_accept4"},
+    {"accept4", ProbeType::kReturn, "syscall__probe_ret_accept4"},
+    {"write", ProbeType::kEntry, "syscall__probe_entry_write"},
+    {"write", ProbeType::kReturn, "syscall__probe_ret_write"},
+    {"writev", ProbeType::kEntry, "syscall__probe_entry_writev"},
+    {"writev", ProbeType::kReturn, "syscall__probe_ret_writev"},
+    {"send", ProbeType::kEntry, "syscall__probe_entry_send"},
+    {"send", ProbeType::kReturn, "syscall__probe_ret_send"},
+    {"sendto", ProbeType::kEntry, "syscall__probe_entry_sendto"},
+    {"sendto", ProbeType::kReturn, "syscall__probe_ret_sendto"},
+    {"sendmsg", ProbeType::kEntry, "syscall__probe_entry_sendmsg"},
+    {"sendmsg", ProbeType::kReturn, "syscall__probe_ret_sendmsg"},
+    {"sendmmsg", ProbeType::kEntry, "syscall__probe_entry_sendmmsg"},
+    {"sendmmsg", ProbeType::kReturn, "syscall__probe_ret_sendmmsg"},
+    {"sendfile", ProbeType::kEntry, "syscall__probe_entry_sendfile"},
+    {"sendfile", ProbeType::kReturn, "syscall__probe_ret_sendfile"},
+    {"sendfile64", ProbeType::kEntry, "syscall__probe_entry_sendfile"},
+    {"sendfile64", ProbeType::kReturn, "syscall__probe_ret_sendfile"},
+    {"read", ProbeType::kEntry, "syscall__probe_entry_read"},
+    {"read", ProbeType::kReturn, "syscall__probe_ret_read"},
+    {"readv", ProbeType::kEntry, "syscall__probe_entry_readv"},
+    {"readv", ProbeType::kReturn, "syscall__probe_ret_readv"},
+    {"recv", ProbeType::kEntry, "syscall__probe_entry_recv"},
+    {"recv", ProbeType::kReturn, "syscall__probe_ret_recv"},
+    {"recvfrom", ProbeType::kEntry, "syscall__probe_entry_recvfrom"},
+    {"recvfrom", ProbeType::kReturn, "syscall__probe_ret_recvfrom"},
+    {"recvmsg", ProbeType::kEntry, "syscall__probe_entry_recvmsg"},
+    {"recvmsg", ProbeType::kReturn, "syscall__probe_ret_recvmsg"},
+    {"recvmmsg", ProbeType::kEntry, "syscall__probe_entry_recvmmsg"},
+    {"recvmmsg", ProbeType::kReturn, "syscall__probe_ret_recvmmsg"},
+    {"close", ProbeType::kEntry, "syscall__probe_entry_close"},
+    {"close", ProbeType::kReturn, "syscall__probe_ret_close"},
+    {"mmap", ProbeType::kEntry, "syscall__probe_entry_mmap"},
+    {"sock_alloc", ProbeType::kReturn, "probe_ret_sock_alloc", /*is_syscall*/ false},
+    {"security_socket_sendmsg", ProbeType::kEntry, "probe_entry_socket_sendmsg",
+     /*is_syscall*/ false, /* is_optional */ false,
+     new bpf_tools::KProbeSpec{"sock_sendmesg", ProbeType::kEntry, "probe_entry_socket_sendmsg",
+                               false, true}},
+    {"security_socket_recvmsg", ProbeType::kEntry, "probe_entry_socket_recvmsg",
+     /*is_syscall*/ false, /* is_optional */ false,
+     new bpf_tools::KProbeSpec{"sock_recvmsg", ProbeType::kEntry, "probe_entry_socket_recvmsg",
+                               false, true}},
+});
 
 using bpf_tools::PerfBufferSizeCategory;