netd.c - nettrace - add support for icmp type/code & udplite/sctp ports.

Test: TreeHugger
Signed-off-by: Maciej Żenczykowski <maze@google.com>
Change-Id: I37c2aefe8759fdbab3078c60e0f97b57799697fc
diff --git a/bpf_progs/netd.c b/bpf_progs/netd.c
index e2e6d02..41a93f4 100644
--- a/bpf_progs/netd.c
+++ b/bpf_progs/netd.c
@@ -56,6 +56,8 @@
 // see include/uapi/linux/tcp.h
 #define TCP_FLAG32_OFF 12
 
+#define TCP_FLAG8_OFF (TCP_FLAG32_OFF + 1)
+
 // For maps netd does not need to access
 #define DEFINE_BPF_MAP_NO_NETD(the_map, TYPE, TypeOfKey, TypeOfValue, num_entries)      \
     DEFINE_BPF_MAP_EXT(the_map, TYPE, TypeOfKey, TypeOfValue, num_entries,              \
@@ -274,13 +276,27 @@
 
     uint8_t flags = 0;
     __be16 sport = 0, dport = 0;
-    if (proto == IPPROTO_TCP && L4_off >= 20) {
-        (void)bpf_skb_load_bytes_net(skb, L4_off + TCP_FLAG32_OFF + 1, &flags, sizeof(flags), kver);
-        (void)bpf_skb_load_bytes_net(skb, L4_off + TCP_OFFSET(source), &sport, sizeof(sport), kver);
-        (void)bpf_skb_load_bytes_net(skb, L4_off + TCP_OFFSET(dest), &dport, sizeof(dport), kver);
-    } else if (proto == IPPROTO_UDP && L4_off >= 20) {
-        (void)bpf_skb_load_bytes_net(skb, L4_off + UDP_OFFSET(source), &sport, sizeof(sport), kver);
-        (void)bpf_skb_load_bytes_net(skb, L4_off + UDP_OFFSET(dest), &dport, sizeof(dport), kver);
+    if (L4_off >= 20) {
+      switch (proto) {
+        case IPPROTO_TCP:
+          (void)bpf_skb_load_bytes_net(skb, L4_off + TCP_FLAG8_OFF, &flags, sizeof(flags), kver);
+          // fallthrough
+        case IPPROTO_DCCP:
+        case IPPROTO_UDP:
+        case IPPROTO_UDPLITE:
+        case IPPROTO_SCTP:
+          // all of these L4 protocols start with be16 src & dst port
+          (void)bpf_skb_load_bytes_net(skb, L4_off + 0, &sport, sizeof(sport), kver);
+          (void)bpf_skb_load_bytes_net(skb, L4_off + 2, &dport, sizeof(dport), kver);
+          break;
+        case IPPROTO_ICMP:
+        case IPPROTO_ICMPV6:
+          // Both IPv4 and IPv6 icmp start with u8 type & code, which we store in the bottom
+          // (ie. second) byte of sport/dport (which are be16s), the top byte is already zero.
+          (void)bpf_skb_load_bytes_net(skb, L4_off + 0, (char *)&sport + 1, 1, kver); //type
+          (void)bpf_skb_load_bytes_net(skb, L4_off + 1, (char *)&dport + 1, 1, kver); //code
+          break;
+      }
     }
 
     pkt->timestampNs = bpf_ktime_get_boot_ns();