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

pkg/security/ebpf/c/include/helpers/network/utils.h

@@ -16,4 +16,10 @@ __attribute__((always_inline)) void flip(struct flow_t *flow) {
     flow->daddr[1] = tmp;
 }
 
+// addr holds an in6_addr read as two host-order u64: ::ffff:0:0/96 puts the 0xffff marker in the
+// low half of the second one
+__attribute__((always_inline)) u8 is_ipv4_mapped_ipv6_addr(u64 *addr) {
+    return addr[0] == 0 && (addr[1] & 0xffffffff) == 0xffff0000;
+}
+
 #endif

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

@@ -2,12 +2,10 @@
 #define _HOOKS_ACCEPT_H_
 
 #include "constants/offsets/network.h"
+#include "helpers/events.h"
+#include "hooks/network/flow.h"
 
 int __attribute__((always_inline)) read_sock_and_send_event(ctx_t * ctx, struct sock * sock) {
-    if(sock == NULL) {
-        return 0;
-    }
-
     struct accept_event_t event = {0};
 
     // Extract family from the socket
@@ -39,7 +37,17 @@ int __attribute__((always_inline)) read_sock_and_send_event(ctx_t * ctx, struct
 HOOK_EXIT("inet_csk_accept")
 int hook_accept(ctx_t *ctx) {
     struct sock *sock = (struct sock*)CTX_PARMRET(ctx);
-    return read_sock_and_send_event(ctx, sock);
+    if (sock == NULL || IS_ERR(sock)) {
+        return 0;
+    }
+
+    // this hook is loaded along with the network probes so that flow_pid stays up to date, hence the
+    // event itself is only sent when a rule asks for it
+    if (is_event_enabled(EVENT_ACCEPT)) {
+        read_sock_and_send_event(ctx, sock);
+    }
+
+    return register_accepted_flow(sock);
 }
 
 #endif /* _HOOKS_ACCEPT_H_ */

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

@@ -2,6 +2,7 @@
 #define _HOOKS_NETWORK_FLOW_H_
 #include "constants/offsets/network.h"
 #include "constants/offsets/netns.h"
+#include "helpers/network/parser.h"
 #include "helpers/network/pid_resolver.h"
 #include "helpers/network/utils.h"
 #include "helpers/network/flow.h"
@@ -630,4 +631,90 @@ int rethook_inet6_bind(ctx_t *ctx) {
     return handle_inet_bind_ret(ret);
 }
 
+__attribute__((always_inline)) int register_connected_flow(struct sock *sk, u64 pid_tgid) {
+    struct pid_route_t route = {};
+
+    route.netns = get_netns_from_sock(sk);
+    route.l4_protocol = get_protocol_from_sock(sk);
+    route.port = get_skc_num_from_sock_common((void *)sk);
+    if (route.port == 0) {
+        return 0;
+    }
+
+    u16 family = get_family_from_sock_common((void *)sk);
+    if (family == AF_INET) {
+        bpf_probe_read(&route.addr, sizeof(sk->__sk_common.skc_rcv_saddr), &sk->__sk_common.skc_rcv_saddr);
+    } else if (family == AF_INET6) {
+        bpf_probe_read(&route.addr, sizeof(u64) * 2, &sk->__sk_common.skc_v6_rcv_saddr);
+    } else {
+        return 0;
+    }
+
+    struct sock_meta_t *meta = get_sock_meta(sk);
+    if (meta != NULL) {
+        struct pid_route_t previous = meta->existing_route;
+        if (previous.port != 0 || previous.addr[0] != 0 || previous.addr[1] != 0) {
+            if (can_delete_route(&previous, sk)) {
+
+                #if defined(DEBUG_NETWORK_FLOW)
+                bpf_printk("|    flushing route registered before the source address was known:");
+                print_route(&previous);
+                #endif
+
+                bpf_map_delete_elem(&flow_pid, &previous);
+            }
+        }
+    }
+
+    if (!can_delete_route(&route, sk)) {
+        // we don't want to override the existing entry
+        return 0;
+    }
+
+    struct pid_route_entry_t value = {};
+    value.pid = pid_tgid >> 32;
+    value.type = FLOW_CLASSIFICATION_ENTRY;
+    value.owner_sk = sk;
+    bpf_map_update_elem(&flow_pid, &route, &value, BPF_ANY);
+
+    if (meta != NULL) {
+        meta->existing_route = route;
+    }
+
+    if (route.netns != 0) {
+        u32 tid = (u32)pid_tgid;
+        bpf_map_update_elem(&netns_cache, &tid, &route.netns, BPF_ANY);
+    }
+
+    #if defined(DEBUG_NETWORK_FLOW)
+    bpf_printk("register_connected_flow: @:0x%p", sk);
+    print_route(&route);
+    print_route_entry(&value);
+    #endif
+
+    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
+    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
+    if (is_ipv4_mapped_ipv6_addr(addr)) {
+        return 0;
+    }
+
+    return register_connected_flow(sk, bpf_get_current_pid_tgid());
+}
+
 #endif

pkg/security/ebpf/probes/event_types.go

@@ -51,6 +51,7 @@ func NetworkSelectors(hasCgroupSocket bool) []manager.ProbesSelector {
 			hookFunc("hook_inet_shutdown"),
 			hookFunc("hook_inet_bind"),
 			hookFunc("rethook_inet_bind"),
+			hookFunc("hook_accept"),
 			hookFunc("hook_sk_common_release"),
 			hookFunc("hook_path_get"),
 			hookFunc("hook_proc_fd_link"),
@@ -663,6 +664,8 @@ func GetSelectorsPerEventType(hasFentry, haveIOURing bool) map[eval.EventType][]
 			}}},
 
 		// List of probes required to capture accept events
+		// hook_accept is also part of NetworkSelectors because it keeps flow_pid up to date, it is
+		// kept here so that accept events are still captured when network tracking is off
 		"accept": {
 			&manager.AllOf{Selectors: []manager.ProbesSelector{
 				hookFunc("hook_accept"),