Skip to content

Commit

Permalink
bpf: use sched_process_fork instead of clone/clone3/fork/vfork
Browse files Browse the repository at this point in the history
  • Loading branch information
mozillazg committed May 2, 2024
1 parent 5ec5e91 commit 5a80d3a
Show file tree
Hide file tree
Showing 9 changed files with 153 additions and 105 deletions.
35 changes: 13 additions & 22 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -78,18 +78,9 @@ jobs:
uname -a
cat /etc/issue
timeout 20s /host/ptcpdump/ptcpdump -c 1 -i any --print -w ptcpdump.pcapng 'dst host 1.1.1.1 and tcp[tcpflags] = tcp-syn' 2>&1 1>/tmp/log | (read _; curl -m 10 1.1.1.1 &>/dev/null || true)
bash /host/ptcpdump/testdata/test_base.sh /host/ptcpdump/ptcpdump
cat /tmp/log
cat /tmp/log | grep '/usr/bin/curl'
cat /tmp/log | grep -F ' > 1.1.1.1.80: Flags [S],'
apt update || true
apt install -y tcpdump
tcpdump -nr ./ptcpdump.pcapng
tcpdump -nr ./ptcpdump.pcapng | grep -F ' > 1.1.1.1.80: Flags [S],'
- name: Test filter by process
- name: Test filter by process name
uses: cilium/little-vm-helper@908ab1ff8a596a03cd5221a1f8602dc44c3f906d # v0.0.12
with:
provision: 'false'
Expand All @@ -99,19 +90,19 @@ jobs:
uname -a
cat /etc/issue
timeout 20s /host/ptcpdump/ptcpdump -c 3 --pname curl -i any --print -w ptcpdump.pcapng 2>&1 1>/tmp/log | (read _; curl -m 10 1.1.1.1 &>/dev/null || true)
bash /host/ptcpdump/testdata/test_pname_filter.sh /host/ptcpdump/ptcpdump
cat /tmp/log
cat /tmp/log | grep '/usr/bin/curl'
cat /tmp/log | grep -F ' > 1.1.1.1.80: Flags [S],'
cat /tmp/log | grep -P '1.1.1.1.80 > .*: Flags \[S.\],'
- name: Test filter by process id
uses: cilium/[email protected]
with:
provision: 'false'
dns-resolver: '1.1.1.1'
cmd: |
set -ex
uname -a
cat /etc/issue
apt update || true
apt install -y tcpdump
tcpdump -nr ./ptcpdump.pcapng
tcpdump -nr ./ptcpdump.pcapng | grep -F ' > 1.1.1.1.80: Flags [S],'
tcpdump -nr ./ptcpdump.pcapng | grep -P '1.1.1.1.80 > .*: Flags \[S.\],'
bash /host/ptcpdump/testdata/test_pid_filter.sh /host/ptcpdump/ptcpdump
releaser-test:
runs-on: ubuntu-latest
Expand Down
17 changes: 10 additions & 7 deletions bpf/bpf.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ type BPF struct {
type Options struct {
Pid uint32
Comm [16]int8
filterComm uint8
FollowForks uint8
PcapFilter string
}
Expand Down Expand Up @@ -58,6 +59,7 @@ func NewOptions(pid uint, comm string, followForks bool, pcapFilter string) Opti
opts.Comm[i] = int8(s)
}
opts.Comm[15] = '\x00'
opts.filterComm = 1
}
opts.FollowForks = 0
if followForks {
Expand All @@ -72,6 +74,7 @@ func (b *BPF) Load(opts Options) error {
err := b.spec.RewriteConstants(map[string]interface{}{
"filter_pid": opts.Pid,
"filter_comm": opts.Comm,
"filter_comm_enable": opts.filterComm,
"filter_follow_forks": opts.FollowForks,
})
if err != nil {
Expand Down Expand Up @@ -147,14 +150,14 @@ func (b *BPF) AttachTracepoints() error {
b.links = append(b.links, lk)

if b.opts.attachForks() {
for _, name := range []string{"sys_exit_fork", "sys_exit_vfork", "sys_exit_clone", "sys_exit_clone3"} {
lk, err := link.Tracepoint("syscalls", name,
b.objs.TracepointSyscallsSysExitFork, nil)
if err != nil {
return xerrors.Errorf("attach tracepoint/syscalls/%s: %w", name, err)
}
b.links = append(b.links, lk)
lk, err := link.AttachRawTracepoint(link.RawTracepointOptions{
"sched_process_fork",
b.objs.RawTracepointSchedProcessFork,
})
if err != nil {
return xerrors.Errorf("attach raw_tracepoint/sched_process_fork: %w", err)
}
b.links = append(b.links, lk)
}

return nil
Expand Down
15 changes: 3 additions & 12 deletions bpf/bpf_x86_bpfel.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

Binary file modified bpf/bpf_x86_bpfel.o
Binary file not shown.
86 changes: 24 additions & 62 deletions bpf/ptcpdump.c
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
static volatile const u32 filter_pid = 0;
static volatile const u8 filter_follow_forks = 0;
volatile const char filter_comm[TASK_COMM_LEN];
static volatile const u8 filter_comm_enable = 0;
static const u8 u8_zero = 0;
static const u32 u32_zero = 0;

Expand Down Expand Up @@ -306,99 +307,59 @@ static __always_inline bool str_cmp(const char *a, const volatile char *b, int l
return 0;
}

static __always_inline int str_len(const volatile char *s, int max_len)
{
#pragma unroll
for (int i = 0; i < max_len; i++) {
if (s[i] == '\0')
return i;
static __always_inline int process_filter(struct task_struct *task) {
// no filter rules
if (filter_pid < 1 && filter_comm_enable != 1) {
return 0;
}
if (s[max_len - 1] != '\0')
return max_len;
return 0;
}

static __always_inline int process_filter(struct task_struct *task) {
u32 pid = BPF_CORE_READ(task, tgid);
if (bpf_map_lookup_elem(&filter_pid_map, &pid)) {
return 0;
}

bool should_filter = false;
if (filter_pid != 0) {
if (pid != filter_pid) {
return -1;
}
if (filter_pid > 0 && pid == filter_pid) {
should_filter = true;
}

if (!should_filter) {
if (str_len(filter_comm, TASK_COMM_LEN) > 1) {
if (filter_comm_enable == 1) {
char comm[TASK_COMM_LEN];
BPF_CORE_READ_STR_INTO(&comm, task, comm);
if (str_cmp(comm, filter_comm, TASK_COMM_LEN) != 0) {
return -1;
if (str_cmp(comm, filter_comm, TASK_COMM_LEN) == 0) {
should_filter = true;
}
should_filter = true;
}
}

if (should_filter) {
bpf_map_update_elem(&filter_pid_map, &pid, &u8_zero, BPF_NOEXIST);
return 0;
}

return 0;
return -1;
}

static __always_inline void handle_fork(struct trace_event_raw_sys_exit *ctx) {
static __always_inline void handle_fork(struct task_struct *parent, struct task_struct *child) {
if (filter_follow_forks != 1) {
return;
}
u32 child_pid = BPF_CORE_READ(ctx, ret);
if (child_pid <= 0) {
return;
}

bool should_filter = false;
struct task_struct *task = (struct task_struct*)bpf_get_current_task();
if (process_filter(task) == 0) {
should_filter = true;
// bpf_printk("handle fork: parent");
return;
}

if (!should_filter) {
if (filter_pid > 0 && child_pid == filter_pid) {
should_filter = true;
}
}
if (should_filter) {
if (process_filter(parent) == 0) {
bpf_map_update_elem(&filter_pid_map, &child_pid, &u8_zero, BPF_NOEXIST);
// bpf_printk("handle fork: %d", child_pid);
return;
}
if (process_filter(child) == 0) {
return;
}
return;
}

SEC("tracepoint/syscalls/sys_exit_fork")
int tracepoint__syscalls__sys_exit_fork(struct trace_event_raw_sys_exit *ctx) {
handle_fork(ctx);
return 0;
}

SEC("tracepoint/syscalls/sys_exit_vfork")
int tracepoint__syscalls__sys_exit_vfork(struct trace_event_raw_sys_exit *ctx) {
handle_fork(ctx);
return 0;
}

SEC("tracepoint/syscalls/sys_exit_clone")
int tracepoint__syscalls__sys_exit_clone(struct trace_event_raw_sys_exit *ctx) {
handle_fork(ctx);
return 0;
}

SEC("tracepoint/syscalls/sys_exit_clone3")
int tracepoint__syscalls__sys_exit_clone3(struct trace_event_raw_sys_exit *ctx) {
handle_fork(ctx);
SEC("raw_tracepoint/sched_process_fork")
int raw_tracepoint__sched_process_fork(struct bpf_raw_tracepoint_args *ctx) {
struct task_struct *parent = (struct task_struct *)BPF_CORE_READ(ctx, args[0]);
struct task_struct *child = (struct task_struct *)BPF_CORE_READ(ctx, args[1]);
handle_fork(parent, child);
return 0;
}

Expand Down Expand Up @@ -554,6 +515,7 @@ static __always_inline void handle_exec(struct trace_event_raw_sched_process_exe
return;
}

// TODO: change to use raw tracepoint
SEC("tracepoint/sched/sched_process_exec")
int tracepoint__sched__sched_process_exec(struct trace_event_raw_sched_process_exec *ctx) {
handle_exec(ctx);
Expand Down
3 changes: 1 addition & 2 deletions cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -115,8 +115,7 @@ func run(cmd *cobra.Command, args []string) error {
}()

runtime.Gosched()

log.Println("tracing...")
log.Println("capturing...")
<-ctx.Done()
log.Println("bye bye")

Expand Down
30 changes: 30 additions & 0 deletions testdata/test_base.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
#!/usr/bin/env bash

set -ex

CMD="$1"
FNAME="/tmp/base.pcapng"
LNAME="/tmp/base.log"


function test_ptcpdump() {
timeout 20s ${CMD} -c 1 -i any --print -w "${FNAME}" \
'dst host 1.1.1.1 and tcp[tcpflags] = tcp-syn' 2>&1 | tee "${LNAME}" | (read _; curl -m 10 1.1.1.1 &>/dev/null || true)

cat "${LNAME}"
cat "${LNAME}" | grep '/usr/bin/curl'
cat "${LNAME}" | grep -F ' > 1.1.1.1.80: Flags [S],' # SYN
}

function test_tcpdump_read() {
which tcpdump || (apt update || true && apt install -y tcpdump)
tcpdump -nr "${FNAME}"
tcpdump -nr "${FNAME}" | grep -F ' > 1.1.1.1.80: Flags [S],' # SYN
}

function main() {
test_ptcpdump
test_tcpdump_read
}

main
38 changes: 38 additions & 0 deletions testdata/test_pid_filter.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
#!/usr/bin/env bash

set -ex

CMD="$1"
FNAME="/tmp/filter_by_pid.pcapng"
LNAME="/tmp/filter_by_pid.log"


function test_ptcpdump() {
timeout 30s ${CMD} -c 3 --pid $$ -f -i any --print -w "${FNAME}" 2>&1 | tee "${LNAME}" &
sleep 10
curl -m 10 1.1.1.1 &>/dev/null || true
echo $$
wait

cat "${LNAME}"
cat "${LNAME}" | grep '/usr/bin/curl'
cat "${LNAME}" | grep -F ' > 1.1.1.1.80: Flags [S],' # SYN
cat "${LNAME}" | grep -P '1.1.1.1.80 > .*: Flags \[S.\],' # SYN-ACK
cat "${LNAME}" | grep -F ' > 1.1.1.1.80: Flags [.],' # ACK
}

function test_tcpdump_read() {
which tcpdump || (apt update || true && apt install -y tcpdump)
tcpdump -nr "${FNAME}"
tcpdump -nr "${FNAME}" | grep -F ' > 1.1.1.1.80: Flags [S],' # SYN
tcpdump -nr "${FNAME}" | grep -P '1.1.1.1.80 > .*: Flags \[S.\],' # SYN-ACK
tcpdump -nr "${FNAME}" | grep -F ' > 1.1.1.1.80: Flags [.],' # ACK
}


function main() {
test_ptcpdump
test_tcpdump_read
}

main
34 changes: 34 additions & 0 deletions testdata/test_pname_filter.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
#!/usr/bin/env bash

set -ex

CMD="$1"
FNAME="/tmp/filter_by_pname.pcapng"
LNAME="/tmp/filter_by_pname.log"


function test_ptcpdump() {
timeout 20s ${CMD} -c 3 --pname curl -f -i any \
--print -w "${FNAME}" 2>&1 | tee "${LNAME}" | (read _; curl -m 10 1.1.1.1 &>/dev/null || true)

cat "${LNAME}"
cat "${LNAME}" | grep '/usr/bin/curl'
cat "${LNAME}" | grep -F ' > 1.1.1.1.80: Flags [S],' # SYN
cat "${LNAME}" | grep -P '1.1.1.1.80 > .*: Flags \[S.\],' # SYN-ACK
cat "${LNAME}" | grep -F ' > 1.1.1.1.80: Flags [.],' # ACK
}

function test_tcpdump_read() {
which tcpdump || (apt update || true && apt install -y tcpdump)
tcpdump -nr "${FNAME}"
tcpdump -nr "${FNAME}" | grep -F ' > 1.1.1.1.80: Flags [S],' # SYN
tcpdump -nr "${FNAME}" | grep -P '1.1.1.1.80 > .*: Flags \[S.\],' # SYN-ACK
tcpdump -nr "${FNAME}" | grep -F ' > 1.1.1.1.80: Flags [.],' # ACK
}

function main() {
test_ptcpdump
test_tcpdump_read
}

main

0 comments on commit 5a80d3a

Please sign in to comment.