tools: add option --cgroupmap to tcp tools

List of tcp tools updated: tcpaccept, tcpconnect, tcptracer
diff --git a/tools/tcpaccept.py b/tools/tcpaccept.py
index 7c10420..03b05e0 100755
--- a/tools/tcpaccept.py
+++ b/tools/tcpaccept.py
@@ -29,6 +29,7 @@
     ./tcpaccept -t        # include timestamps
     ./tcpaccept -P 80,81  # only trace port 80 and 81
     ./tcpaccept -p 181    # only trace PID 181
+    ./tcpaccept --cgroupmap ./mappath  # only trace cgroups in this BPF map
 """
 parser = argparse.ArgumentParser(
     description="Trace TCP accepts",
@@ -42,6 +43,8 @@
     help="trace this PID only")
 parser.add_argument("-P", "--port",
     help="comma-separated list of local ports to trace")
+parser.add_argument("--cgroupmap",
+    help="trace cgroups in this BPF map only")
 parser.add_argument("--ebpf", action="store_true",
     help=argparse.SUPPRESS)
 args = parser.parse_args()
@@ -77,6 +80,11 @@
     char task[TASK_COMM_LEN];
 };
 BPF_PERF_OUTPUT(ipv6_events);
+
+#if CGROUPSET
+BPF_TABLE_PINNED("hash", u64, u64, cgroupset, 1024, "CGROUPPATH");
+#endif
+
 """
 
 #
@@ -89,6 +97,13 @@
 bpf_text_kprobe = """
 int kretprobe__inet_csk_accept(struct pt_regs *ctx)
 {
+#if CGROUPSET
+    u64 cgroupid = bpf_get_current_cgroup_id();
+    if (cgroupset.lookup(&cgroupid) == NULL) {
+        return 0;
+    }
+#endif
+
     struct sock *newsk = (struct sock *)PT_REGS_RC(ctx);
     u32 pid = bpf_get_current_pid_tgid() >> 32;
 
@@ -184,6 +199,11 @@
     lports_if = ' && '.join(['lport != %d' % lport for lport in lports])
     bpf_text = bpf_text.replace('##FILTER_PORT##',
         'if (%s) { return 0; }' % lports_if)
+if args.cgroupmap:
+    bpf_text = bpf_text.replace('CGROUPSET', '1')
+    bpf_text = bpf_text.replace('CGROUPPATH', args.cgroupmap)
+else:
+    bpf_text = bpf_text.replace('CGROUPSET', '0')
 if debug or args.ebpf:
     print(bpf_text)
     if args.ebpf:
diff --git a/tools/tcpaccept_example.txt b/tools/tcpaccept_example.txt
index 2adee45..5b6b1a6 100644
--- a/tools/tcpaccept_example.txt
+++ b/tools/tcpaccept_example.txt
@@ -33,22 +33,33 @@
 1.984    907    sshd         4  127.0.0.1        51250 127.0.0.1        22
 
 
+The --cgroupmap option filters based on a cgroup set. It is meant to be used
+with an externally created map.
+
+# ./tcpaccept --cgroupmap /sys/fs/bpf/test01
+
+For more details, see docs/filtering_by_cgroups.md
+
+
 USAGE message:
 
 # ./tcpaccept -h
-usage: tcpaccept [-h] [-T] [-t] [-p PID] [-P PORTS]
+usage: tcpaccept.py [-h] [-T] [-t] [-p PID] [-P PORT] [--cgroupmap CGROUPMAP]
 
 Trace TCP accepts
 
 optional arguments:
-  -h, --help              show this help message and exit
-  -T, --time              include time column on output (HH:MM:SS)
-  -t, --timestamp         include timestamp on output
-  -p PID, --pid PID       trace this PID only
-  -P PORTS, --port PORTS  comma-separated list of local ports to trace
+  -h, --help            show this help message and exit
+  -T, --time            include time column on output (HH:MM:SS)
+  -t, --timestamp       include timestamp on output
+  -p PID, --pid PID     trace this PID only
+  -P PORT, --port PORT  comma-separated list of local ports to trace
+  --cgroupmap CGROUPMAP
+                        trace cgroups in this BPF map only
 
 examples:
     ./tcpaccept           # trace all TCP accept()s
     ./tcpaccept -t        # include timestamps
     ./tcpaccept -P 80,81  # only trace port 80 and 81
     ./tcpaccept -p 181    # only trace PID 181
+    ./tcpaccept --cgroupmap ./mappath  # only trace cgroups in this BPF map
diff --git a/tools/tcpconnect.py b/tools/tcpconnect.py
index eb12667..67f2cef 100755
--- a/tools/tcpconnect.py
+++ b/tools/tcpconnect.py
@@ -37,6 +37,7 @@
     ./tcpconnect -U        # include UID
     ./tcpconnect -u 1000   # only trace UID 1000
     ./tcpconnect -c        # count connects per src ip and dest ip/port
+    ./tcpconnect --cgroupmap ./mappath  # only trace cgroups in this BPF map
 """
 parser = argparse.ArgumentParser(
     description="Trace TCP connects",
@@ -54,6 +55,8 @@
     help="trace this UID only")
 parser.add_argument("-c", "--count", action="store_true",
     help="count connects per src ip and dest ip/port")
+parser.add_argument("--cgroupmap",
+    help="trace cgroups in this BPF map only")
 parser.add_argument("--ebpf", action="store_true",
     help=argparse.SUPPRESS)
 args = parser.parse_args()
@@ -67,6 +70,10 @@
 
 BPF_HASH(currsock, u32, struct sock *);
 
+#if CGROUPSET
+BPF_TABLE_PINNED("hash", u64, u64, cgroupset, 1024, "CGROUPPATH");
+#endif
+
 // separate data structs for ipv4 and ipv6
 struct ipv4_data_t {
     u64 ts_us;
@@ -109,6 +116,13 @@
 
 int trace_connect_entry(struct pt_regs *ctx, struct sock *sk)
 {
+#if CGROUPSET
+    u64 cgroupid = bpf_get_current_cgroup_id();
+    if (cgroupset.lookup(&cgroupid) == NULL) {
+      return 0;
+    }
+#endif
+
     u64 pid_tgid = bpf_get_current_pid_tgid();
     u32 pid = pid_tgid >> 32;
     u32 tid = pid_tgid;
@@ -234,6 +248,11 @@
 if args.uid:
     bpf_text = bpf_text.replace('FILTER_UID',
         'if (uid != %s) { return 0; }' % args.uid)
+if args.cgroupmap:
+    bpf_text = bpf_text.replace('CGROUPSET', '1')
+    bpf_text = bpf_text.replace('CGROUPPATH', args.cgroupmap)
+else:
+    bpf_text = bpf_text.replace('CGROUPSET', '0')
 
 bpf_text = bpf_text.replace('FILTER_PID', '')
 bpf_text = bpf_text.replace('FILTER_PORT', '')
diff --git a/tools/tcpconnect_example.txt b/tools/tcpconnect_example.txt
index e88701e..cf97562 100644
--- a/tools/tcpconnect_example.txt
+++ b/tools/tcpconnect_example.txt
@@ -68,10 +68,19 @@
 [...]
 
 
+The --cgroupmap option filters based on a cgroup set. It is meant to be used
+with an externally created map.
+
+# ./tcpconnect --cgroupmap /sys/fs/bpf/test01
+
+For more details, see docs/filtering_by_cgroups.md
+
+
 USAGE message:
 
 # ./tcpconnect -h
-usage: tcpconnect [-h] [-c] [-t] [-p PID] [-P PORT]
+usage: tcpconnect.py [-h] [-t] [-p PID] [-P PORT] [-U] [-u UID] [-c]
+                     [--cgroupmap CGROUPMAP]
 
 Trace TCP connects
 
@@ -79,11 +88,12 @@
   -h, --help            show this help message and exit
   -t, --timestamp       include timestamp on output
   -p PID, --pid PID     trace this PID only
-  -P PORT, --port PORT
-                        comma-separated list of destination ports to trace.
+  -P PORT, --port PORT  comma-separated list of destination ports to trace.
   -U, --print-uid       include UID on output
   -u UID, --uid UID     trace this UID only
   -c, --count           count connects per src ip and dest ip/port
+  --cgroupmap CGROUPMAP
+                        trace cgroups in this BPF map only
 
 examples:
     ./tcpconnect           # trace all TCP connect()s
@@ -94,3 +104,4 @@
     ./tcpconnect -U        # include UID
     ./tcpconnect -u 1000   # only trace UID 1000
     ./tcpconnect -c        # count connects per src ip and dest ip/port
+    ./tcpconnect --cgroupmap ./mappath  # only trace cgroups in this BPF map
diff --git a/tools/tcptracer.py b/tools/tcptracer.py
index e61fe9b..8e6e1ec 100755
--- a/tools/tcptracer.py
+++ b/tools/tcptracer.py
@@ -11,7 +11,7 @@
 # The following code should be replaced, and simplified, when static TCP probes
 # exist.
 #
-# Copyright 2017 Kinvolk GmbH
+# Copyright 2017-2020 Kinvolk GmbH
 #
 # Licensed under the Apache License, Version 2.0 (the "License")
 from __future__ import print_function
@@ -29,6 +29,8 @@
                     help="trace this PID only")
 parser.add_argument("-N", "--netns", default=0, type=int,
                     help="trace this Network Namespace only")
+parser.add_argument("--cgroupmap",
+                    help="trace cgroups in this BPF map only")
 parser.add_argument("-v", "--verbose", action="store_true",
                     help="include Network Namespace in the output")
 parser.add_argument("--ebpf", action="store_true",
@@ -77,6 +79,10 @@
 };
 BPF_PERF_OUTPUT(tcp_ipv6_event);
 
+#if CGROUPSET
+BPF_TABLE_PINNED("hash", u64, u64, cgroupset, 1024, "CGROUPPATH");
+#endif
+
 // tcp_set_state doesn't run in the context of the process that initiated the
 // connection so we need to store a map TUPLE -> PID to send the right PID on
 // the event
@@ -173,6 +179,13 @@
 
 int trace_connect_v4_entry(struct pt_regs *ctx, struct sock *sk)
 {
+#if CGROUPSET
+  u64 cgroupid = bpf_get_current_cgroup_id();
+  if (cgroupset.lookup(&cgroupid) == NULL) {
+      return 0;
+  }
+#endif
+
   u64 pid = bpf_get_current_pid_tgid();
 
   ##FILTER_PID##
@@ -220,6 +233,12 @@
 
 int trace_connect_v6_entry(struct pt_regs *ctx, struct sock *sk)
 {
+#if CGROUPSET
+  u64 cgroupid = bpf_get_current_cgroup_id();
+  if (cgroupset.lookup(&cgroupid) == NULL) {
+      return 0;
+  }
+#endif
   u64 pid = bpf_get_current_pid_tgid();
 
   ##FILTER_PID##
@@ -352,6 +371,13 @@
 
 int trace_close_entry(struct pt_regs *ctx, struct sock *skp)
 {
+#if CGROUPSET
+  u64 cgroupid = bpf_get_current_cgroup_id();
+  if (cgroupset.lookup(&cgroupid) == NULL) {
+      return 0;
+  }
+#endif
+
   u64 pid = bpf_get_current_pid_tgid();
 
   ##FILTER_PID##
@@ -413,6 +439,13 @@
 
 int trace_accept_return(struct pt_regs *ctx)
 {
+#if CGROUPSET
+  u64 cgroupid = bpf_get_current_cgroup_id();
+  if (cgroupset.lookup(&cgroupid) == NULL) {
+      return 0;
+  }
+#endif
+
   struct sock *newsk = (struct sock *)PT_REGS_RC(ctx);
   u64 pid = bpf_get_current_pid_tgid();
 
@@ -581,6 +614,11 @@
 
 bpf_text = bpf_text.replace('##FILTER_PID##', pid_filter)
 bpf_text = bpf_text.replace('##FILTER_NETNS##', netns_filter)
+if args.cgroupmap:
+    bpf_text = bpf_text.replace('CGROUPSET', '1')
+    bpf_text = bpf_text.replace('CGROUPPATH', args.cgroupmap)
+else:
+    bpf_text = bpf_text.replace('CGROUPSET', '0')
 
 if args.ebpf:
     print(bpf_text)
diff --git a/tools/tcptracer_example.txt b/tools/tcptracer_example.txt
index 2873d03..b6e5258 100644
--- a/tools/tcptracer_example.txt
+++ b/tools/tcptracer_example.txt
@@ -35,3 +35,11 @@
 3.546    C    748  curl             6  [::1]            [::1]            42592  80
 4.294    X  31002  telnet           4  192.168.1.2      192.168.1.1      42590  23
 ```
+
+
+The --cgroupmap option filters based on a cgroup set. It is meant to be used
+with an externally created map.
+
+# ./tcptracer --cgroupmap /sys/fs/bpf/test01
+
+For more details, see docs/filtering_by_cgroups.md