hook exec events to get args of pid

bpf/bpf.go

@@ -9,13 +9,14 @@ import (
 	"github.com/cilium/ebpf"
 	"github.com/cilium/ebpf/link"
 	"github.com/cilium/ebpf/perf"
+	"github.com/cilium/ebpf/ringbuf"
 	"github.com/florianl/go-tc"
 	"github.com/florianl/go-tc/core"
 	"golang.org/x/sys/unix"
 	"golang.org/x/xerrors"
 )
 
-//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -target native -type packet_event_t Bpf ./ptcpdump.c -- -I./headers -I. -Wall
+//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -no-strip -target native -type packet_event_t -type exec_event_t Bpf ./ptcpdump.c -- -I./headers -I. -Wall
 
 const tcFilterName = "ptcpdump"
 
@@ -72,6 +73,16 @@ func (b *BPF) AttachKprobes() error {
 	return nil
 }
 
+func (b *BPF) AttachTracepoints() error {
+	lk, err := link.Tracepoint("sched", "sched_process_exec",
+		b.objs.TracepointSchedSchedProcessExec, nil)
+	if err != nil {
+		return xerrors.Errorf("attach tracepoint/sched/sched_process_exec: %w", err)
+	}
+	b.links = append(b.links, lk)
+	return nil
+}
+
 func (b *BPF) AttachTcHooks(dev *net.Interface) error {
 	closeFunc, err := ensureTcQdisc(dev)
 	if err != nil {
@@ -106,6 +117,14 @@ func (b *BPF) NewPacketEventReader() (*perf.Reader, error) {
 	return reader, nil
 }
 
+func (b *BPF) NewExecEventReader() (*ringbuf.Reader, error) {
+	reader, err := ringbuf.NewReader(b.objs.ExecEvents)
+	if err != nil {
+		return nil, xerrors.Errorf(": %w", err)
+	}
+	return reader, nil
+}
+
 func attachTcHook(dev *net.Interface, prog *ebpf.Program, ingress bool) (func(), error) {
 	tcnl, err := tc.Open(&tc.Config{})
 	if err != nil {

bpf/ptcpdump.c

@@ -25,6 +25,7 @@
 #define MAX_PAYLOAD_SIZE 1500
 #define INGRESS_PACKET 0
 #define EGRESS_PACKET 1
+#define EXEC_ARGS_LEN 4096
 
 char _license[] SEC("license") = "Dual MIT/GPL";
 
@@ -79,6 +80,17 @@ struct packet_event_t {
     u8 payload[MAX_PAYLOAD_SIZE];
 };
 
+struct exec_event_t {
+    u32 pid;
+    u8 truncated;
+    unsigned int args_size;
+    char args[EXEC_ARGS_LEN];
+};
+
+struct {
+    __uint(type, BPF_MAP_TYPE_RINGBUF);
+    __uint(max_entries, 1024 * 512);
+} exec_events SEC(".maps");
 
 struct {
     __uint(type, BPF_MAP_TYPE_LRU_HASH);
@@ -100,6 +112,13 @@ struct {
     __uint(value_size, sizeof(u32));
 } packet_events SEC(".maps");
 
+// force emitting struct into the ELF.
+// the `-type` flag of bpf2go need this
+// avoid "Error: collect C types: type name XXX: not found"
+const struct packet_event_t *unused1 __attribute__((unused));
+const struct exec_event_t *unused2 __attribute__((unused));
+
+
 static __always_inline int parse_skb_l2(struct __sk_buff *skb, struct l2_t *l2, u32 *offset) {
     if (bpf_skb_load_bytes(skb, *offset + offsetof(struct ethhdr, h_proto), &l2->h_protocol, sizeof(l2->h_protocol)) < 0 ) {
         return -1;
@@ -270,7 +289,7 @@ int BPF_KPROBE(kprobe__security_sk_classify_flow, struct sock *sk) {
     return 0;
 };
 
-static __always_inline int handel_tc(struct __sk_buff *skb, bool egress) {
+static __always_inline int handle_tc(struct __sk_buff *skb, bool egress) {
     bpf_skb_pull_data(skb, 0);
 
     struct packet_meta_t packet_meta = {0};
@@ -304,7 +323,7 @@ static __always_inline int handel_tc(struct __sk_buff *skb, bool egress) {
 
     struct flow_pid_value_t *value = bpf_map_lookup_elem(&flow_pid_map, &key);
     if (value) {
-        bpf_printk("[tc] (%s) %pI4 %d", value->comm, &key.saddr[0], key.sport);
+//        bpf_printk("[tc] (%s) %pI4 %d", value->comm, &key.saddr[0], key.sport);
     } else {
         /* bpf_printk("[tc] %pI4 %d bpf_map_lookup_elem is empty", &key.saddr[0], key.sport); */
         return TC_ACT_OK;
@@ -335,12 +354,47 @@ static __always_inline int handel_tc(struct __sk_buff *skb, bool egress) {
     return TC_ACT_OK;
 }
 
+static __always_inline void handle_exec(struct trace_event_raw_sched_process_exec *ctx) {
+    struct exec_event_t *event;
+    event = bpf_ringbuf_reserve(&exec_events, sizeof(*event), 0);
+    if (!event) {
+        bpf_printk("[ptcpdump] bpf_ringbuf_reserve failed");
+        return;
+    }
+
+    struct task_struct *task = (struct task_struct *)bpf_get_current_task();
+    event->pid = bpf_get_current_pid_tgid() >> 32;
+
+    void *arg_start = (void *)BPF_CORE_READ(task, mm, arg_start);
+    void *arg_end = (void *)BPF_CORE_READ(task, mm, arg_end);
+    unsigned long arg_length = arg_end - arg_start;
+    if (arg_length > EXEC_ARGS_LEN) {
+        arg_length = EXEC_ARGS_LEN;
+        event->truncated = 1;
+    }
+    int arg_ret = bpf_probe_read(&event->args, arg_length, arg_start);
+    if (arg_ret < 0) {
+        bpf_printk("[ptcpdump] read exec args failed: %d", arg_ret);
+    } else {
+        event->args_size = arg_length;
+    }
+
+    bpf_ringbuf_submit(event, 0);
+    return;
+}
+
+SEC("tracepoint/sched/sched_process_exec")
+int tracepoint__sched__sched_process_exec(struct trace_event_raw_sched_process_exec *ctx) {
+    handle_exec(ctx);
+    return 0;
+}
+
 SEC("tc")
 int tc_ingress(struct __sk_buff *skb) {
-    return handel_tc(skb, false);
+    return handle_tc(skb, false);
 };
 
 SEC("tc")
 int tc_egress(struct __sk_buff *skb) {
-    return handel_tc(skb, true);
+    return handle_tc(skb, true);
 };

internal/event/process.go

@@ -0,0 +1,36 @@
+package event
+
+import (
+	"bytes"
+	"encoding/binary"
+	"github.com/mozillazg/ptcpdump/bpf"
+	"golang.org/x/xerrors"
+)
+
+type ProcessExec struct {
+	Pid       int
+	Args      []byte
+	Truncated bool
+}
+
+func ParseProcessExecEvent(rawSample []byte) (*ProcessExec, error) {
+	var p ProcessExec
+	event := bpf.BpfExecEventT{}
+	if err := binary.Read(bytes.NewBuffer(rawSample), binary.LittleEndian, &event); err != nil {
+		return nil, xerrors.Errorf("parse event: %w", err)
+	}
+
+	if event.Truncated == 1 {
+		p.Truncated = true
+	}
+	p.Pid = int(event.Pid)
+	for i := 0; i < int(event.ArgsSize); i++ {
+		b := byte(event.Args[i])
+		if b == '\x00' {
+			b = ' '
+		}
+		p.Args = append(p.Args, b)
+	}
+
+	return &p, nil
+}

internal/metadata/process.go

@@ -0,0 +1,30 @@
+package metadata
+
+import (
+	"github.com/mozillazg/ptcpdump/internal/event"
+)
+
+type ProcessCache struct {
+	m map[int]*event.ProcessExec
+}
+
+func NewProcessCache() *ProcessCache {
+	return &ProcessCache{
+		m: make(map[int]*event.ProcessExec),
+	}
+}
+
+func (c *ProcessCache) AddItem(exec event.ProcessExec) {
+	pid := exec.Pid
+	c.m[pid] = &exec
+
+	//log.Printf("add new cache: %d", pid)
+}
+
+func (c *ProcessCache) Get(pid int) event.ProcessExec {
+	p := c.m[pid]
+	if p == nil {
+		return event.ProcessExec{}
+	}
+	return *p
+}

internal/writer/pcapng.go

@@ -5,31 +5,38 @@ import (
 	"github.com/gopacket/gopacket"
 	"github.com/gopacket/gopacket/pcapgo"
 	"github.com/mozillazg/ptcpdump/internal/event"
+	"github.com/mozillazg/ptcpdump/internal/metadata"
 	"golang.org/x/xerrors"
+	"log"
 	"time"
 )
 
 type PcapNGWriter struct {
-	pw *pcapgo.NgWriter
+	pw     *pcapgo.NgWriter
+	pcache *metadata.ProcessCache
 }
 
-func NewPcapNGWriter(pw *pcapgo.NgWriter) *PcapNGWriter {
-	return &PcapNGWriter{pw: pw}
+func NewPcapNGWriter(pw *pcapgo.NgWriter, pcache *metadata.ProcessCache) *PcapNGWriter {
+	return &PcapNGWriter{pw: pw, pcache: pcache}
 }
 
-func (w *PcapNGWriter) Write(p *event.Packet) error {
-	payloadLen := len(p.Data)
+func (w *PcapNGWriter) Write(e *event.Packet) error {
+	payloadLen := len(e.Data)
 	info := gopacket.CaptureInfo{
 		Timestamp:      time.Now(),
 		CaptureLength:  payloadLen,
 		Length:         payloadLen,
 		InterfaceIndex: 0,
 	}
+	p := w.pcache.Get(e.Pid)
+	if p.Pid == 0 {
+		log.Printf("not found pid from cache: %d", e.Pid)
+	}
 	opts := pcapgo.NgPacketOptions{
-		Comment: fmt.Sprintf("PID: %d\nCOMMAND: %s", p.Pid, p.Comm),
+		Comment: fmt.Sprintf("PID: %d\nCOMMAND: %s", e.Pid, string(p.Args)),
 	}
 
-	if err := w.pw.WritePacketWithOptions(info, p.Data, opts); err != nil {
+	if err := w.pw.WritePacketWithOptions(info, e.Data, opts); err != nil {
 		return xerrors.Errorf("writing packet: %w", err)
 	}