u6: fix socketpair exit fd capture and socket filtering

4 files changed, 25 insertions, 8 deletions

diff --git a/internal/generate/bpfhandler.go b/internal/generate/bpfhandler.go
index 6c9c314..50917e7 100644
--- a/internal/generate/bpfhandler.go
+++ b/internal/generate/bpfhandler.go
@@ -76,7 +76,7 @@ func generateExtra(tp GeneratedTracepoint, isEnter bool) string {
 	case KindSocket:
 		return generateExtraSocket()
 	case KindSocketpair:
-		return generateExtraSocketpair()
+		return generateExtraSocketpair(isEnter)
 	case KindOpen:
 		return generateExtraOpen(f)
 	case KindPathname:
@@ -157,8 +157,11 @@ func generateExtraSocket() string {
 	return "    ev->family = (__s32)ctx->args[0];\n    ev->type = (__s32)ctx->args[1];\n    ev->protocol = (__s32)ctx->args[2];\n"
 }
 
-func generateExtraSocketpair() string {
-	return "    int sv[2];\n    __builtin_memset(&sv, 0xff, sizeof(sv));\n    if (ctx->args[3] != 0) {\n        bpf_probe_read_user(&sv, sizeof(sv), (void *)ctx->args[3]);\n    }\n    ev->family = (__s32)ctx->args[0];\n    ev->type = (__s32)ctx->args[1];\n    ev->protocol = (__s32)ctx->args[2];\n    ev->sv0 = (__s32)sv[0];\n    ev->sv1 = (__s32)sv[1];\n"
+func generateExtraSocketpair(isEnter bool) string {
+	if isEnter {
+		return "    struct socketpair_ctx pending;\n    pending.usockvec = ctx->args[3];\n    pending.family = (__s32)ctx->args[0];\n    pending.type = (__s32)ctx->args[1];\n    pending.protocol = (__s32)ctx->args[2];\n    bpf_map_update_elem(&socketpair_ctx_map, &tid, &pending, BPF_ANY);\n    ev->family = pending.family;\n    ev->type = pending.type;\n    ev->protocol = pending.protocol;\n    ev->sv0 = -1;\n    ev->sv1 = -1;\n    ev->ret = 0;\n"
+	}
+	return "    __s32 family = -1;\n    __s32 type = -1;\n    __s32 protocol = -1;\n    __s32 sv0 = -1;\n    __s32 sv1 = -1;\n    struct socketpair_ctx *pending = bpf_map_lookup_elem(&socketpair_ctx_map, &tid);\n    if (pending) {\n        family = pending->family;\n        type = pending->type;\n        protocol = pending->protocol;\n        if (ctx->ret == 0 && pending->usockvec != 0) {\n            int sv[2];\n            if (bpf_probe_read_user(&sv, sizeof(sv), (void *)pending->usockvec) == 0) {\n                sv0 = (__s32)sv[0];\n                sv1 = (__s32)sv[1];\n            }\n        }\n        bpf_map_delete_elem(&socketpair_ctx_map, &tid);\n    }\n    ev->family = family;\n    ev->type = type;\n    ev->protocol = protocol;\n    ev->sv0 = sv0;\n    ev->sv1 = sv1;\n    ev->ret = ctx->ret;\n"
 }
 
 // eventStructName returns the C struct name for a TracepointKind. The mapping
diff --git a/internal/generate/classify.go b/internal/generate/classify.go
index 68cd722..56f5cd2 100644
--- a/internal/generate/classify.go
+++ b/internal/generate/classify.go
@@ -85,6 +85,8 @@ func classifyNameOnly(name string) (ClassificationResult, bool) {
 		return ClassificationResult{Kind: KindSocket}, true
 	case "sys_enter_socketpair":
 		return ClassificationResult{Kind: KindSocketpair}, true
+	case "sys_exit_socketpair":
+		return ClassificationResult{Kind: KindSocketpair}, true
 	}
 	if strings.HasPrefix(name, "sys_enter_io_") {
 		return ClassificationResult{Kind: KindNull}, true
diff --git a/internal/generate/classify_test.go b/internal/generate/classify_test.go
index ce6eff1..e6353b5 100644
--- a/internal/generate/classify_test.go
+++ b/internal/generate/classify_test.go
@@ -258,6 +258,13 @@ func TestClassifySocketpair(t *testing.T) {
 	}
 }
 
+func TestClassifyExitSocketpair(t *testing.T) {
+	r := classifyFromData(t, FormatExitSocketpair)
+	if r.Kind != KindSocketpair {
+		t.Errorf("exit_socketpair: got kind %d, want KindSocketpair", r.Kind)
+	}
+}
+
 func TestClassifyKillRequiresGenerationFallback(t *testing.T) {
 	r := classifyFromData(t, FormatKill)
 	if r.Kind != KindNone {
diff --git a/internal/generate/codegen_test.go b/internal/generate/codegen_test.go
index f87a6dc..3e29612 100644
--- a/internal/generate/codegen_test.go
+++ b/internal/generate/codegen_test.go
@@ -200,11 +200,16 @@ func TestGenerateSocketpairHandler(t *testing.T) {
 
 	requireContains(t, output, "struct socketpair_event *ev")
 	requireContains(t, output, "ev->event_type = ENTER_SOCKETPAIR_EVENT;")
-	requireContains(t, output, "int sv[2];")
-	requireContains(t, output, "bpf_probe_read_user(&sv, sizeof(sv), (void *)ctx->args[3]);")
-	requireContains(t, output, "ev->family = (__s32)ctx->args[0];")
-	requireContains(t, output, "ev->sv0 = (__s32)sv[0];")
-	requireContains(t, output, "ev->sv1 = (__s32)sv[1];")
+	requireContains(t, output, "struct socketpair_ctx pending;")
+	requireContains(t, output, "bpf_map_update_elem(&socketpair_ctx_map, &tid, &pending, BPF_ANY);")
+	requireContains(t, output, "ev->sv0 = -1;")
+	requireContains(t, output, "ev->ret = 0;")
+	requireContains(t, output, "SEC(\"tracepoint/syscalls/sys_exit_socketpair\")")
+	requireContains(t, output, "ev->event_type = EXIT_SOCKETPAIR_EVENT;")
+	requireContains(t, output, "struct socketpair_ctx *pending = bpf_map_lookup_elem(&socketpair_ctx_map, &tid);")
+	requireContains(t, output, "if (ctx->ret == 0 && pending->usockvec != 0) {")
+	requireContains(t, output, "ev->ret = ctx->ret;")
+	requireContains(t, output, "ev->family = pending.family;")
 }
 
 func TestGenerateNameToHandleAtHandler(t *testing.T) {