[CWS] Add tests

DataDog · Jan 30, 2025 · 4e99c4d · 4e99c4d
1 parent 1c58c0f
commit 4e99c4d
Show file tree

Hide file tree

Showing 4 changed files with 858 additions and 6 deletions.
diff --git a/pkg/security/ebpf/c/include/hooks/network/flow.h b/pkg/security/ebpf/c/include/hooks/network/flow.h
@@ -320,18 +320,27 @@ int hook_inet_release(ctx_t *ctx) {
     return handle_sk_release(sk);
 }
 
-HOOK_ENTRY("inet_bind")
-int hook_inet_bind(ctx_t *ctx) {
-    struct socket *sock = (struct socket *)CTX_PARM1(ctx);
+__attribute__((always_inline)) int handle_inet_bind(struct socket *sock) {
     struct inet_bind_args_t args = {};
     args.sock = sock;
     u64 pid = bpf_get_current_pid_tgid();
     bpf_map_update_elem(&inet_bind_args, &pid, &args, BPF_ANY);
     return 0;
 }
 
-HOOK_EXIT("inet_bind")
-int rethook_inet_bind(ctx_t *ctx) {
+HOOK_ENTRY("inet_bind")
+int hook_inet_bind(ctx_t *ctx) {
+    struct socket *sock = (struct socket *)CTX_PARM1(ctx);
+    return handle_inet_bind(sock);
+}
+
+HOOK_ENTRY("inet6_bind")
+int hook_inet6_bind(ctx_t *ctx) {
+    struct socket *sock = (struct socket *)CTX_PARM1(ctx);
+    return handle_inet_bind(sock);
+}
+
+__attribute__((always_inline)) int handle_inet_bind_ret(int ret) {
     // fetch inet_bind arguments
     u64 id = bpf_get_current_pid_tgid();
     u32 tid = (u32)id;
@@ -344,7 +353,6 @@ int rethook_inet_bind(ctx_t *ctx) {
     // delete the entry in inet_bind_args to make sure we always cleanup inet_bind_args and we don't leak entries
     bpf_map_delete_elem(&inet_bind_args, &id);
 
-    int ret = CTX_PARMRET(ctx);
     if (ret < 0) {
         // we only care about successful bind operations
         return 0;
@@ -394,4 +402,16 @@ int rethook_inet_bind(ctx_t *ctx) {
     return 0;
 }
 
+HOOK_EXIT("inet_bind")
+int rethook_inet_bind(ctx_t *ctx) {
+    int ret = CTX_PARMRET(ctx);
+    return handle_inet_bind_ret(ret);
+}
+
+HOOK_EXIT("inet6_bind")
+int rethook_inet6_bind(ctx_t *ctx) {
+    int ret = CTX_PARMRET(ctx);
+    return handle_inet_bind_ret(ret);
+}
+
 #endif
diff --git a/pkg/security/ebpf/probes/event_types.go b/pkg/security/ebpf/probes/event_types.go
@@ -49,6 +49,8 @@ func NetworkSelectors() []manager.ProbesSelector {
 			kprobeOrFentry("inet_shutdown"),
 			kprobeOrFentry("inet_bind"),
 			kretprobeOrFexit("inet_bind"),
+			kprobeOrFentry("inet6_bind"),
+			kretprobeOrFexit("inet6_bind"),
 			kprobeOrFentry("sk_common_release"),
 			kprobeOrFentry("path_get"),
 			kprobeOrFentry("proc_fd_link"),

diff --git a/pkg/security/ebpf/probes/flow.go b/pkg/security/ebpf/probes/flow.go
@@ -48,6 +48,18 @@ func getFlowProbes() []*manager.Probe {
 				EBPFFuncName: "rethook_inet_bind",
 			},
 		},
+		{
+			ProbeIdentificationPair: manager.ProbeIdentificationPair{
+				UID:          SecurityAgentUID,
+				EBPFFuncName: "hook_inet6_bind",
+			},
+		},
+		{
+			ProbeIdentificationPair: manager.ProbeIdentificationPair{
+				UID:          SecurityAgentUID,
+				EBPFFuncName: "rethook_inet6_bind",
+			},
+		},
 		{
 			ProbeIdentificationPair: manager.ProbeIdentificationPair{
 				UID:          SecurityAgentUID,