Skip to content

Commit 1d998ef

Browse files
kylewanginchinakylewanginchina
committed
fix(ebpf): avoid double output on 5.2+ syscall path
1 parent 507fdcd commit 1d998ef

3 files changed

Lines changed: 117 additions & 9 deletions

File tree

‎agent/src/ebpf/kernel/files_rw.bpf.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ static __inline int trace_io_event_common(void *ctx,
438438
v->thread_trace_id = trace_id;
439439
v->msg_type = MSG_COMMON;
440440
bpf_get_current_comm(v->comm, sizeof(v->comm));
441-
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
441+
#ifdef USE_SOCKET_TRACE_TAIL_CALLS
442442
struct tail_calls_context *context =
443443
(struct tail_calls_context *)v->data;
444444
context->max_size_limit = data_max_sz;

‎agent/src/ebpf/kernel/socket_trace.bpf.c‎

Lines changed: 79 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,24 @@
3939

4040
#define __user
4141

42+
/*
43+
* Linux 5.2+ kernels allow programs larger than the old 4096-insn cap, but
44+
* Linux 5.15 still has a 1,000,000 processed-insn verifier complexity limit.
45+
* AI Agent governance adds state to syscall socket hot paths, so split those
46+
* paths in the 5.2_plus object while leaving uprobe paths on the old layout.
47+
*/
48+
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
49+
#define USE_SOCKET_TRACE_TAIL_CALLS 1
50+
#endif
51+
52+
#if defined(LINUX_VER_5_2_PLUS) && defined(EXTENDED_AI_AGENT_FILE_IO)
53+
#define USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS 1
54+
#endif
55+
56+
#if defined(LINUX_VER_KFUNC) || defined(LINUX_VER_5_2_PLUS)
57+
#define USE_SOCKET_TRACE_INLINE_OUTPUT 1
58+
#endif
59+
4260
#ifdef EXTENDED_AI_AGENT_FILE_IO
4361
#ifndef AI_AGENT_PROC_FORK
4462
#define AI_AGENT_PROC_FORK 1
@@ -1329,7 +1347,7 @@ static __inline void trace_process(struct socket_info_s *socket_info_ptr,
13291347
}
13301348
}
13311349

1332-
#if defined(LINUX_VER_KFUNC) || defined(LINUX_VER_5_2_PLUS)
1350+
#ifdef USE_SOCKET_TRACE_INLINE_OUTPUT
13331351
static __inline int
13341352
__output_data_common(void *ctx, struct tracer_ctx_s *tracer_ctx,
13351353
struct __socket_data_buffer *v_buff,
@@ -1751,7 +1769,23 @@ __data_submit(struct pt_regs *ctx, struct conn_info_s *conn_info,
17511769
p->protocols[idx] = (__u8) v->data_type;
17521770
}
17531771
}
1772+
#endif
1773+
1774+
#ifdef USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS
1775+
if (extra->source == DATA_SOURCE_SYSCALL) {
1776+
struct tail_calls_context *context =
1777+
(struct tail_calls_context *)v->data;
1778+
context->max_size_limit = data_max_sz;
1779+
context->push_reassembly_bytes = send_reasm_bytes;
1780+
context->vecs = (bool) vecs;
1781+
context->is_close = false;
1782+
context->dir = conn_info->direction;
1783+
1784+
return SUBMIT_OK;
1785+
}
1786+
#endif
17541787

1788+
#ifdef USE_SOCKET_TRACE_INLINE_OUTPUT
17551789
return __output_data_common(ctx, tracer_ctx, v_buff, args,
17561790
conn_info->direction, (bool) vecs,
17571791
tracer_ctx->data_limit_max, false,
@@ -1843,7 +1877,11 @@ static __inline int process_data(struct pt_regs *ctx, __u64 id,
18431877
if (!(sk != NULL &&
18441878
((sock_state = is_tcp_udp_data(sk, offset, conn_info))
18451879
!= SOCK_CHECK_TYPE_ERROR))) {
1846-
#if defined(LINUX_VER_KFUNC) || defined(LINUX_VER_5_2_PLUS)
1880+
#ifdef USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS
1881+
if (extra->source == DATA_SOURCE_SYSCALL)
1882+
return -2;
1883+
#endif
1884+
#ifdef USE_SOCKET_TRACE_INLINE_OUTPUT
18471885
return trace_io_event_common(ctx, offset, args, direction, id);
18481886
#else
18491887
return -2; // This means attempting to handle I/O events.
@@ -1914,7 +1952,7 @@ static __inline int process_data(struct pt_regs *ctx, __u64 id,
19141952

19151953
if (act == INFER_TERMINATE)
19161954
return -1;
1917-
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
1955+
#ifdef USE_SOCKET_TRACE_TAIL_CALLS
19181956
if (disable_kprobe && extra->source == DATA_SOURCE_SYSCALL)
19191957
return -1;
19201958

@@ -1938,6 +1976,26 @@ static __inline int process_data(struct pt_regs *ctx, __u64 id,
19381976
PROG_PROTO_INFER_KP_2_IDX);
19391977
}
19401978
}
1979+
#elif defined(USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS)
1980+
if (extra->source == DATA_SOURCE_SYSCALL) {
1981+
if (disable_kprobe)
1982+
return -1;
1983+
1984+
if (act == INFER_CONTINUE) {
1985+
ctx_map->tail_call.conn_info = __conn_info;
1986+
ctx_map->tail_call.extra = *extra;
1987+
ctx_map->tail_call.bytes_count = bytes_count;
1988+
ctx_map->tail_call.offset = offset;
1989+
ctx_map->tail_call.dir = direction;
1990+
#ifdef SUPPORTS_KPROBE_ONLY
1991+
bpf_tail_call(ctx, &NAME(progs_jmp_kp_map),
1992+
PROG_PROTO_INFER_KP_2_IDX);
1993+
#else
1994+
bpf_tail_call(ctx, &NAME(progs_jmp_tp_map),
1995+
PROG_PROTO_INFER_TP_2_IDX);
1996+
#endif
1997+
}
1998+
}
19411999
#endif
19422000

19432001
if (conn_info->protocol == PROTO_CUSTOM) {
@@ -1950,7 +2008,7 @@ static __inline int process_data(struct pt_regs *ctx, __u64 id,
19502008
// data_submit can be performed, otherwise MySQL data may be lost
19512009
if (conn_info->protocol != PROTO_UNKNOWN ||
19522010
conn_info->message_type != MSG_UNKNOWN) {
1953-
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
2011+
#ifdef USE_SOCKET_TRACE_TAIL_CALLS
19542012
/*
19552013
* Fill in tail call context information.
19562014
*/
@@ -1959,6 +2017,17 @@ static __inline int process_data(struct pt_regs *ctx, __u64 id,
19592017
ctx_map->tail_call.bytes_count = bytes_count;
19602018
ctx_map->tail_call.offset = offset;
19612019
return 0;
2020+
#elif defined(USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS)
2021+
if (extra->source == DATA_SOURCE_SYSCALL) {
2022+
ctx_map->tail_call.conn_info = __conn_info;
2023+
ctx_map->tail_call.extra = *extra;
2024+
ctx_map->tail_call.bytes_count = bytes_count;
2025+
ctx_map->tail_call.offset = offset;
2026+
return 0;
2027+
}
2028+
return __data_submit(ctx, conn_info, args, extra->vecs,
2029+
bytes_count, offset, args->enter_ts,
2030+
extra);
19622031
#else
19632032
return __data_submit(ctx, conn_info, args, extra->vecs,
19642033
bytes_count, offset, args->enter_ts,
@@ -1983,7 +2052,8 @@ static __inline void process_syscall_data(struct pt_regs *ctx, __u64 id,
19832052

19842053
int result = process_data(ctx, id, direction, args, bytes_count, &extra);
19852054
if (result == 0) {
1986-
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
2055+
#if defined(USE_SOCKET_TRACE_TAIL_CALLS) || \
2056+
defined(USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS)
19872057
#ifdef SUPPORTS_KPROBE_ONLY
19882058
bpf_tail_call(ctx, &NAME(progs_jmp_kp_map),
19892059
PROG_DATA_SUBMIT_KP_IDX);
@@ -2017,7 +2087,8 @@ static __inline void process_syscall_data_vecs(struct pt_regs *ctx, __u64 id,
20172087

20182088
int result = process_data(ctx, id, direction, args, bytes_count, &extra);
20192089
if (result == 0) {
2020-
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
2090+
#if defined(USE_SOCKET_TRACE_TAIL_CALLS) || \
2091+
defined(USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS)
20212092
#ifdef SUPPORTS_KPROBE_ONLY
20222093
bpf_tail_call(ctx, &NAME(progs_jmp_kp_map),
20232094
PROG_DATA_SUBMIT_KP_IDX);
@@ -2960,7 +3031,7 @@ static __inline void __push_close_event(__u64 pid_tgid, __u64 uid, __u64 seq,
29603031
v->fd = fd;
29613032
bpf_get_current_comm(v->comm, sizeof(v->comm));
29623033

2963-
#if !defined(LINUX_VER_KFUNC) && !defined(LINUX_VER_5_2_PLUS)
3034+
#ifdef USE_SOCKET_TRACE_TAIL_CALLS
29643035
struct tail_calls_context *context =
29653036
(struct tail_calls_context *)v->data;
29663037
context->max_size_limit = data_max_sz;
@@ -3343,7 +3414,7 @@ static __inline int output_extra_data_common(struct data_args_t *args, struct __
33433414
return 0;
33443415
}
33453416

3346-
#if defined(LINUX_VER_KFUNC) || defined(LINUX_VER_5_2_PLUS)
3417+
#ifdef USE_SOCKET_TRACE_INLINE_OUTPUT
33473418
static __inline int __output_data_common(void *ctx,
33483419
struct tracer_ctx_s *tracer_ctx,
33493420
struct __socket_data_buffer *v_buff,

‎agent/src/ebpf/test/test_ai_agent_source_contracts.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -73,6 +73,18 @@ def read_source(path: Path) -> str:
7373
"__u64 pid_tgid = bpf_get_current_pid_tgid();" in data_submit_text,
7474
"__data_submit must define pid_tgid before EXTENDED_AI_AGENT_FILE_IO branch uses it",
7575
)
76+
require(
77+
re.search(
78+
r"#ifdef USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS\s+"
79+
r"if \(extra->source == DATA_SOURCE_SYSCALL\) \{\s+"
80+
r"struct tail_calls_context \*context =\s+"
81+
r"\(struct tail_calls_context \*\)v->data;.*?"
82+
r"return SUBMIT_OK;\s+\}\s+#endif\s+#ifdef USE_SOCKET_TRACE_INLINE_OUTPUT",
83+
data_submit_text,
84+
re.S,
85+
),
86+
"__data_submit must keep 5.2_plus syscall traffic on the tail-call output path before inline output fallback",
87+
)
7688

7789
push_close_end = socket_trace_text.find("\n}\n\n#ifdef SUPPORTS_KPROBE_ONLY", push_close_start)
7890
require(push_close_end != -1, "missing __push_close_event end")
@@ -141,6 +153,31 @@ def read_source(path: Path) -> str:
141153
"AI Agent access_permission extraction must be guarded by EXTENDED_AI_AGENT_FILE_IO_FULL",
142154
)
143155

156+
require(
157+
"#define USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS 1" in socket_trace_text
158+
and "defined(LINUX_VER_5_2_PLUS)" in socket_trace_text
159+
and "defined(EXTENDED_AI_AGENT_FILE_IO)" in socket_trace_text,
160+
"5.2_plus socket-trace must re-enable syscall tail-call splitting when AI Agent governance is compiled in",
161+
)
162+
163+
for helper_name in (
164+
"process_syscall_data",
165+
"process_syscall_data_vecs",
166+
):
167+
helper_idx = socket_trace_text.find(helper_name)
168+
require(helper_idx != -1, f"missing {helper_name}")
169+
helper_text = socket_trace_text[helper_idx : helper_idx + 2200]
170+
require(
171+
"USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS" in helper_text,
172+
f"{helper_name} must use the syscall tail-call split guard",
173+
)
174+
175+
require(
176+
"USE_SOCKET_TRACE_SYSCALL_TAIL_CALLS" in socket_trace_text
177+
and "extra->source == DATA_SOURCE_SYSCALL" in socket_trace_text,
178+
"5.2_plus AI Agent tail-call split must be limited to syscall source paths",
179+
)
180+
144181
require('"lsm/"' in load_text, "load.c must recognize lsm/ section prefix")
145182
require(
146183
"BPF_PROG_TYPE_LSM" in load_text,

0 commit comments

Comments
 (0)