fix(cws): register connected sockets in the flow_pid map

pkg/security/ebpf/c/include/hooks/network/connect.h

@@ -4,6 +4,7 @@
 #include "constants/offsets/netns.h"
 #include "constants/syscall_macro.h"
 #include "helpers/discarders.h"
+#include "hooks/network/flow.h"
 
 int __attribute__((always_inline)) sys_connect(void *ctx, u64 pid_tgid) {
     struct policy_t policy = fetch_policy(EVENT_CONNECT);
@@ -40,6 +41,13 @@ int __attribute__((always_inline)) sys_connect_ret(void *ctx, int retval) {
         return 0;
     }
 
+    register_connecting_flow(syscall->connect.sk, syscall->connect.pid_tgid ? syscall->connect.pid_tgid : bpf_get_current_pid_tgid());
+
+    // these probes are also loaded with the network probes, only send the event when a rule asks for it
+    if (!is_event_enabled(EVENT_CONNECT)) {
+        return 0;
+    }
+
     // emit a sample refresh if the dedup map flagged one
     if (syscall->state == DISCARDED && (syscall->resolver.flags & SAMPLE_REFRESH_NEEDED)) {
         struct sample_refresh_event_t ev = {};
@@ -115,6 +123,7 @@ int hook_security_socket_connect(ctx_t *ctx) {
 
     struct sock *sk = get_sock_from_socket(sock);
     syscall->connect.protocol = get_protocol_from_sock(sk);
+    syscall->connect.sk = sk;
     return 0;
 }
 

pkg/security/ebpf/c/include/hooks/network/flow.h

@@ -695,26 +695,37 @@ __attribute__((always_inline)) int register_connected_flow(struct sock *sk, u64
     return 0;
 }
 
-// The socket returned by accept() holds the concrete local address the connection landed on, which
-// the BIND_ENTRY of a wildcard listener doesn't cover. Before Linux 7.0 that IPv6 socket was classified
-// on its first transmit, causing security_sk_classify_flow to be called and classify the flow.
-// Starting with Linux 7.0 that security_sk_classify_flow call is only done on a route miss in inet6_csk_xmit,
-// this means that we miss the classification in the hit case.
-__attribute__((always_inline)) int register_accepted_flow(struct sock *sk) {
-    // Only native IPv6 sockets lost that classification with Linux 7.0, IPv4 always reaches
-    // security_sk_classify_flow, so the flow registration is already handled by the security_sk_classify_flow hook
+// Before Linux 7.0 an IPv6 socket was classified on its first transmit by security_sk_classify_flow.
+// Starting with Linux 7.0, security_sk_classify_flow is only called by inet6_csk_xmit on a route miss,
+// so IPv6 sockets need to call this helper to register the corresponding flow.
+__attribute__((always_inline)) int register_native_ipv6_flow(struct sock *sk, u64 pid_tgid) {
+    // IPv4 sockets still reach security_sk_classify_flow with a usable flow
     if (get_family_from_sock_common((void *)sk) != AF_INET6) {
         return 0;
     }
 
     u64 addr[2] = {};
     bpf_probe_read(&addr, sizeof(addr), &sk->__sk_common.skc_v6_rcv_saddr);
-    // ipv4 mapped addresses already go through the security_sk_classify_flow path
+    // IPv6 sockets using IPv4 mapped addresses still reach security_sk_classify_flow with a usable flow
     if (is_ipv4_mapped_ipv6_addr(addr)) {
         return 0;
     }
 
-    return register_connected_flow(sk, bpf_get_current_pid_tgid());
+    return register_connected_flow(sk, pid_tgid);
+}
+
+// the BIND_ENTRY of a wildcard listener doesn't cover the local address the connection landed on
+__attribute__((always_inline)) int register_accepted_flow(struct sock *sk) {
+    return register_native_ipv6_flow(sk, bpf_get_current_pid_tgid());
+}
+
+// tcp_v6_connect classifies the flow before inet_hash_connect assigns the ephemeral port
+__attribute__((always_inline)) int register_connecting_flow(struct sock *sk, u64 pid_tgid) {
+    if (sk == NULL) {
+        return 0;
+    }
+
+    return register_native_ipv6_flow(sk, pid_tgid);
 }
 
 #endif

pkg/security/ebpf/c/include/structs/syscalls.h

@@ -242,6 +242,7 @@ struct syscall_cache_t {
             u16 port;
             u16 protocol;
             u64 pid_tgid;
+            struct sock *sk;
         } connect;
 
          struct {

pkg/security/ebpf/probes/event_types.go

@@ -37,7 +37,7 @@ func NetworkVethSelectors() []manager.ProbesSelector {
 }
 
 // NetworkSelectors is the list of probes that should be activated when the network is enabled
-func NetworkSelectors(hasCgroupSocket bool) []manager.ProbesSelector {
+func NetworkSelectors(hasFentry, hasCgroupSocket, haveIOURing bool) []manager.ProbesSelector {
 	ps := []manager.ProbesSelector{
 		// flow classification probes
 		&manager.AllOf{Selectors: []manager.ProbesSelector{
@@ -62,6 +62,9 @@ func NetworkSelectors(hasCgroupSocket bool) []manager.ProbesSelector {
 			hookFunc("rethook_inet6_bind"),
 		}},
 
+		// the connect exit is the only hook that sees the final source address and port of a connecting socket
+		&manager.BestEffort{Selectors: ExpandSyscallProbesSelector(SecurityAgentUID, "connect", hasFentry, EntryAndExit)},
+
 		// network device probes
 		&manager.AllOf{Selectors: []manager.ProbesSelector{
 			hookFunc("hook_register_netdevice"),
@@ -90,6 +93,13 @@ func NetworkSelectors(hasCgroupSocket bool) []manager.ProbesSelector {
 		}})
 	}
 
+	if haveIOURing {
+		ps = append(ps, &manager.BestEffort{Selectors: []manager.ProbesSelector{
+			hookFunc("hook_io_connect"),
+			hookFunc("rethook_io_connect"),
+		}})
+	}
+
 	return ps
 }
 
@@ -147,10 +157,10 @@ func GetCapabilitiesMonitoringSelectors() []manager.ProbesSelector {
 // These probes must be loaded independently of the current ruleset or network filter actions as
 // these are used to track resources that are needed if we later dynamically load network rules
 // or network filter actions.
-func GetNetworkSelectors(hasCgroupSocket bool) []manager.ProbesSelector {
+func GetNetworkSelectors(hasFentry, hasCgroupSocket, haveIOURing bool) []manager.ProbesSelector {
 	selectors := []manager.ProbesSelector{
 		&manager.AllOf{Selectors: []manager.ProbesSelector{
-			&manager.AllOf{Selectors: NetworkSelectors(hasCgroupSocket)},
+			&manager.AllOf{Selectors: NetworkSelectors(hasFentry, hasCgroupSocket, haveIOURing)},
 			&manager.AllOf{Selectors: NetworkVethSelectors()},
 		}},
 	}
@@ -683,6 +693,7 @@ func GetSelectorsPerEventType(hasFentry, haveIOURing bool) map[eval.EventType][]
 			&manager.BestEffort{Selectors: ExpandSyscallProbesSelector(SecurityAgentUID, "bind", hasFentry, EntryAndExit)},
 		},
 		// List of probes required to capture connect events
+		// also part of NetworkSelectors, kept here so that connect events are captured when network tracking is off
 		"connect": {
 			&manager.AllOf{Selectors: []manager.ProbesSelector{
 				hookFunc("hook_security_socket_connect"),

pkg/security/probe/probe_ebpf.go

@@ -2352,7 +2352,7 @@ func (p *EBPFProbe) updateProbes(ruleSetEventTypes []eval.EventType, needRawSysc
 	// network filter actions as these are used to track resources that are needed if we later
 	// dynamically load network rules or network filter actions.
 	if p.config.Probe.NetworkEnabled {
-		activatedProbes = append(activatedProbes, probes.GetNetworkSelectors(p.kernelVersion.HasBpfGetSocketCookieForCgroupSocket())...)
+		activatedProbes = append(activatedProbes, probes.GetNetworkSelectors(p.useFentry, p.kernelVersion.HasBpfGetSocketCookieForCgroupSocket(), p.kernelVersion.HaveIOURing())...)
 	}
 
 	if p.config.Probe.CapabilitiesMonitoringEnabled {