feat(filter): Add support for capturing traffic based on user ID (#233)

Author: mozillazg

README.md

@@ -76,11 +76,12 @@ Multiple interfaces:
 sudo ptcpdump -i eth0 -i lo
 ```
 
-Filter by process:
+Filter by process or user:
 
 ```
 sudo ptcpdump -i any --pid 1234 --pid 233 -f
 sudo ptcpdump -i any --pname curl
+sudo ptcpdump -i any --uid 1000
 ```
 
 Capture by process via run target program:
@@ -146,6 +147,7 @@ With `-v`:
 13:44:41.529003 eth0 In IP (tos 0x4, ttl 45, id 45428, offset 0, flags [DF], proto TCP (6), length 52)
     139.178.84.217.443 > 172.19.0.2.42606: Flags [.], cksum 0x5284, seq 3173118145, ack 1385712707, win 118, options [nop,nop,TS val 134560683 ecr 1627716996], length 0
     Process (pid 553587, cmd /usr/bin/wget, args wget kernel.org)
+    User (uid 1000)
     ParentProc (pid 553296, cmd /bin/sh, args sh)
     Container (name test, id d9028334568bf75a5a084963a8f98f78c56bba7f45f823b3780a135b71b91e95, image docker.io/library/alpine:3.18, labels {"io.cri-containerd.kind":"container","io.kubernetes.container.name":"test","io.kubernetes.pod.name":"test","io.kubernetes.pod.namespace":"default","io.kubernetes.pod.uid":"9e4bc54b-de48-4b1c-8b9e-54709f67ed0c"})
     Pod (name test, namespace default, UID 9e4bc54b-de48-4b1c-8b9e-54709f67ed0c, labels {"run":"test"}, annotations {"kubernetes.io/config.seen":"2024-07-21T12:41:00.460249620Z","kubernetes.io/config.source":"api"})
@@ -277,7 +279,7 @@ Flags:
       --container-id string                          Filter by container id (only TCP and UDP packets are supported)
       --container-name string                        Filter by container name (only TCP and UDP packets are supported)
       --containerd-address string                    Address of containerd service (default "/run/containerd/containerd.sock")
-      --context strings                              Specify which context information to include in the output (default [process,parentproc,container,pod])
+      --context strings                              Specify which context information to include in the output (default [process,thread,parentproc,user,container,pod])
       --count                                        Print only on stdout the packet count when reading capture file instead of parsing/printing the packets
       --cri-runtime-address string                   Address of CRI container runtime service (default: uses in order the first successful one of [/var/run/dockershim.sock, /var/run/cri-dockerd.sock, /run/crio/crio.sock, /run/containerd/containerd.sock])
       --delay-before-handle-packet-events duration   Delay some durations before handle packet events
@@ -311,6 +313,7 @@ Flags:
   -c, --receive-count uint                           Exit after receiving count packets
   -s, --snapshot-length uint32                       Snarf snaplen bytes of data from each packet rather than the default of 262144 bytes (default 262144)
       --time-stamp-precision string                  When capturing, set the time stamp precision for the capture to the format (default "micro")
+      --uid uints                                    Filter by user IDs (only TCP and UDP packets are supported) (default [])
   -v, --verbose count                                When parsing and printing, produce (slightly more) verbose output
       --version                                      Print the ptcpdump and libpcap version strings and exit
   -w, --write-file string                            Write the raw packets to file rather than parsing and printing them out. They can later be printed with the -r option. Standard output is used if file is '-'. e.g. ptcpdump.pcapng
@@ -334,6 +337,7 @@ Flags:
 | -r *-*                                            | ✅       |                          |
 | --pid *process_id*                                |         | ✅                        |
 | --pname *process_name*                            |         | ✅                        |
+| --uid *user_id*                                   |         | ✅                        |
 | --container-id *container_id*                     |         | ✅                        |
 | --container-name *container_name*                 |         | ✅                        |
 | --pod-name *pod_name.namespace*                   |         | ✅                        |

bpf/bpf.go

@@ -42,6 +42,7 @@ type BPF struct {
 type Options struct {
 	haveFilter     uint8
 	pids           []uint32
+	uids           []uint32
 	comm           [16]int8
 	filterComm     uint8
 	followForks    uint8
@@ -517,6 +518,14 @@ func (b *BPF) applyFilters() error {
 		}
 	}
 
+	log.Infof("start to update FilterUidMap with %+v", opts.pids)
+	for _, uid := range opts.uids {
+		uid := uid
+		if err := b.objs.BpfMaps.FilterUidMap.Update(uid, value, ebpf.UpdateAny); err != nil {
+			return fmt.Errorf("update FilterUidMap: %w", err)
+		}
+	}
+
 	log.Infof("start to update FilterPidnsMap with %+v", opts.pidnsIds)
 	for _, id := range opts.pidnsIds {
 		id := id
@@ -565,6 +574,16 @@ func (opts *Options) WithPids(pids []uint) *Options {
 	return opts
 }
 
+func (opts *Options) WithUids(uids []uint) *Options {
+	for _, id := range uids {
+		opts.uids = append(opts.uids, uint32(id))
+	}
+	if len(opts.uids) > 0 {
+		opts.haveFilter = 1
+	}
+	return opts
+}
+
 func (opts *Options) WithComm(comm string) *Options {
 	opts.comm = [16]int8{}
 	if len(comm) > 0 {

bpf/bpf_arm64_bpfel.go

@@ -75,6 +75,13 @@ struct {
     __type(value, u8);
 } filter_pid_map SEC(".maps");
 
+struct {
+    __uint(type, BPF_MAP_TYPE_HASH);
+    __uint(max_entries, 10);
+    __type(key, u32);
+    __type(value, u8);
+} filter_uid_map SEC(".maps");
+
 struct {
     __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
     __uint(max_entries, 1);
@@ -111,6 +118,13 @@ static __always_inline int filter_pid(u32 pid) {
     return -1;
 }
 
+static __always_inline int filter_uid(u32 uid) {
+    if (bpf_map_lookup_elem(&filter_uid_map, &uid)) {
+        return 0;
+    }
+    return -1;
+}
+
 static __always_inline int filter_mntns(u32 ns) {
     if (bpf_map_lookup_elem(&filter_mntns_map, &ns)) {
         return 0;
@@ -154,7 +168,9 @@ static __always_inline int process_filter(struct task_struct *task) {
     u32 mntns_id = BPF_CORE_READ(task, nsproxy, mnt_ns, ns.inum);
     u32 netns_id = BPF_CORE_READ(task, nsproxy, net_ns, ns.inum);
     u32 pidns_id = BPF_CORE_READ(task, nsproxy, pid_ns_for_children, ns.inum);
-    if ((filter_pidns(pidns_id) == 0) || (filter_mntns(mntns_id) == 0) || (filter_netns(netns_id) == 0)) {
+    u32 uid = BPF_CORE_READ(task, cred, uid.val);
+    if ((filter_pidns(pidns_id) == 0) || (filter_mntns(mntns_id) == 0) || (filter_netns(netns_id) == 0) ||
+        (filter_uid(uid) == 0)) {
         // debug_log("%u %u %u\n", mntns_id, netns_id, pidns_id);
         should_filter = true;
     }

bpf/bpf_arm64_bpfel.o

@@ -175,6 +175,13 @@ func getCurrentConnects(ctx context.Context, pcache *metadata.ProcessCache, opts
 		ps := pcache.GetPidsByComm(opts.comm)
 		pids = append(pids, ps...)
 	}
+	if len(opts.uids) > 0 {
+		filterPid = true
+		for _, uid := range opts.uids {
+			ps := pcache.GetPidsByUid(int(uid))
+			pids = append(pids, ps...)
+		}
+	}
 	if len(opts.pidnsIds) > 0 {
 		filterPid = true
 		for _, id := range opts.pidnsIds {

bpf/bpf_legacy_arm64_bpfel.go

@@ -86,6 +86,7 @@ type Options struct {
 	mntnsIds []uint32
 	netnsIds []uint32
 	pidnsIds []uint32
+	uids     []uint
 
 	stdout           io.Writer
 	enhancedContexts []string
@@ -356,6 +357,7 @@ func (o *Options) ToCapturerOptions() *capturer.Options {
 		MntnsIds:                      o.mntnsIds,
 		NetnsIds:                      o.netnsIds,
 		PidnsIds:                      o.pidnsIds,
+		Uids:                          o.uids,
 		BTFPath:                       o.btfPath,
 		AllDev:                        o.allDev,
 		AllNetNs:                      o.allNetNs,

bpf/bpf_legacy_arm64_bpfel.o

@@ -56,6 +56,7 @@ func init() {
 		"Interfaces to capture")
 	rootCmd.Flags().UintSliceVar(&opts.pids, "pid", nil, "Filter by process IDs (only TCP and UDP packets are supported)")
 	rootCmd.Flags().StringVar(&opts.comm, "pname", "", "Filter by process name (only TCP and UDP packets are supported)")
+	rootCmd.Flags().UintSliceVar(&opts.uids, "uid", nil, "Filter by user IDs (only TCP and UDP packets are supported)")
 	rootCmd.Flags().BoolVarP(&opts.followForks, "follow-forks", "f", false,
 		"Trace child processes as they are created by currently traced processes when filter by process")
 	rootCmd.Flags().BoolVarP(&opts.listInterfaces, "list-interfaces", "D", false,

bpf/bpf_legacy_x86_bpfel.go

@@ -43,6 +43,7 @@ type Capturer struct {
 
 type Options struct {
 	Pids           []uint
+	Uids           []uint
 	Comm           string
 	FollowForks    bool
 	WriteFilePath  string
@@ -156,6 +157,7 @@ func (c *Capturer) Prepare() error {
 	bpfopts = bpfopts.WithPids(c.opts.Pids).
 		WithComm(c.opts.Comm).
 		WithFollowFork(c.opts.FollowForks).
+		WithUids(c.opts.Uids).
 		WithPidNsIds(c.opts.PidnsIds).
 		WithMntNsIds(c.opts.MntnsIds).
 		WithNetNsIds(c.opts.NetnsIds).
@@ -518,6 +520,7 @@ func updateFlowPidMapValues(bf *bpf.BPF, conns []metadata.Connection) error {
 		}
 		v := bpf.BpfProcessMetaT{
 			Pid:     uint32(conn.Pid),
+			Uid:     uint32(conn.Uid),
 			MntnsId: uint32(conn.MntNs),
 			NetnsId: uint32(conn.NetNs),
 		}

bpf/bpf_legacy_x86_bpfel.o

@@ -14,6 +14,7 @@ type Connection struct {
 	LocalIP   netip.Addr
 	LocalPort int
 	Pid       int
+	Uid       int
 	MntNs     int64
 	NetNs     int64
 }
@@ -58,6 +59,9 @@ func convertConnectionStat(stat net.ConnectionStat) (Connection, error) {
 	conn.LocalIP = addr
 	conn.LocalPort = port
 	conn.Pid = int(stat.Pid)
+	if len(stat.Uids) > 0 {
+		conn.Uid = int(stat.Uids[0])
+	}
 	conn.MntNs = utils.GetMountNamespaceFromPid(conn.Pid)
 	conn.NetNs = utils.GetNetworkNamespaceFromPid(conn.Pid)
 	return conn, nil

bpf/bpf_no_tracing_arm64_bpfel.go

@@ -325,6 +325,27 @@ func (c *ProcessCache) GetPidsByComm(name string) []int {
 	return pids
 }
 
+func (c *ProcessCache) GetPidsByUid(uid int) []int {
+	var pids []int
+
+	c.pids.Range(func(key, value any) bool {
+		pid, ok := key.(int)
+		if !ok {
+			return true
+		}
+		info, ok := value.(*types.PacketContext)
+		if !ok {
+			return true
+		}
+		if info.UserId == uid {
+			pids = append(pids, pid)
+		}
+		return true
+	})
+
+	return pids
+}
+
 func (c *ProcessCache) GetPidsByPidNsId(nsid int64) []int {
 	var pids []int