refactor(bpf): Restructure BPF code (#207)

mozillazg · Dec 7, 2024 · 1d3d114 · 1d3d114
1 parent 7bd3525
commit 1d3d114
Show file tree

Hide file tree

Showing 10 changed files with 601 additions and 560 deletions.
diff --git a/bpf/bpf_arm64_bpfel.o b/bpf/bpf_arm64_bpfel.o
diff --git a/bpf/bpf_legacy_arm64_bpfel.o b/bpf/bpf_legacy_arm64_bpfel.o
diff --git a/bpf/bpf_legacy_x86_bpfel.o b/bpf/bpf_legacy_x86_bpfel.o
diff --git a/bpf/bpf_no_tracing_arm64_bpfel.o b/bpf/bpf_no_tracing_arm64_bpfel.o
diff --git a/bpf/bpf_no_tracing_x86_bpfel.o b/bpf/bpf_no_tracing_x86_bpfel.o
diff --git a/bpf/bpf_x86_bpfel.o b/bpf/bpf_x86_bpfel.o
diff --git a/bpf/headers/helpers.h b/bpf/headers/helpers.h
@@ -0,0 +1,37 @@
+#ifndef __PTCPDUMP_HELPERS_H__
+#define __PTCPDUMP_HELPERS_H__
+
+#include "vmlinux.h"
+#include <bpf/bpf_core_read.h>
+#include <bpf/bpf_endian.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+#define EEXIST 17 /* File exists */
+
+static __always_inline void *bpf_map_lookup_or_try_init(void *map, const void *key, const void *init) {
+    void *value;
+    value = bpf_map_lookup_elem(map, key);
+    if (value) {
+        return value;
+    }
+
+    int err = bpf_map_update_elem(map, key, init, BPF_NOEXIST);
+    if (err && err != -EEXIST)
+        return 0;
+
+    return bpf_map_lookup_elem(map, key);
+}
+
+static __always_inline int str_cmp(const char *a, const volatile char *b, int len) {
+#pragma unroll
+    for (int i = 0; i < len; i++) {
+        if (a[i] != b[i])
+            return -1;
+        if (a[i] == '\0')
+            break;
+    }
+    return 0;
+}
+
+#endif /* __PTCPDUMP_HELPERS_H__ */
diff --git a/bpf/net.h b/bpf/net.h
@@ -0,0 +1,196 @@
+#ifndef __PTCPDUMP_NET_H__
+#define __PTCPDUMP_NET_H__
+
+#include "compat.h"
+#include "vmlinux.h"
+#include <bpf/bpf_core_read.h>
+#include <bpf/bpf_endian.h>
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+#define ETH_HLEN 14       /* Total octets in header.	 */
+#define ETH_P_IP 0x0800   /* Internet Protocol packet	*/
+#define ETH_P_IPV6 0x86DD /* IPv6 over bluebook		*/
+#define IPPROTO_ICMP 1    /* Internet Control Message Protocol	*/
+#define IPPROTO_ICMPV6 58 /* ICMPv6			*/
+#define IPPROTO_TCP 6     /* Transmission Control Protocol	*/
+#define IPPROTO_UDP 17    /* User Datagram Protocol		*/
+#define IPPROTO_SCTP 132  /* Stream Control Transport Protocol	*/
+#define AF_INET 2
+#define AF_INET6 10
+
+struct l2_t {
+    u16 h_protocol;
+};
+
+struct l3_t {
+    u8 protocol;
+    u64 saddr[2];
+    u64 daddr[2];
+};
+
+struct l4_t {
+    u16 sport;
+    u16 dport;
+};
+
+struct packet_meta_t {
+    u32 ifindex;
+
+    struct l2_t l2;
+    struct l3_t l3;
+    struct l4_t l4;
+
+    u32 offset;
+};
+
+struct flow_pid_key_t {
+    u64 saddr[2];
+    u16 sport;
+};
+
+const struct flow_pid_key_t *unused3 __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) {
+        // debug_log("parse_skb_l2 1 failed:\n");
+        return -1;
+    }
+    l2->h_protocol = bpf_ntohs(l2->h_protocol);
+    *offset += sizeof(struct ethhdr);
+    return 0;
+}
+
+static __always_inline int parse_skb_l3(struct __sk_buff *skb, u16 protocol, struct l3_t *l3, u32 *offset) {
+    switch (protocol) {
+    case ETH_P_IP: {
+        struct iphdr ip_hdr;
+        if (bpf_skb_load_bytes(skb, *offset, &ip_hdr, sizeof(struct iphdr)) < 0) {
+            // debug_log("parse_skb_l3 1 failed:\n");
+            return -1;
+        }
+        l3->protocol = ip_hdr.protocol;
+        l3->saddr[0] = ip_hdr.saddr;
+        l3->daddr[0] = ip_hdr.daddr;
+        *offset += sizeof(struct iphdr);
+        return 0;
+    }
+    case ETH_P_IPV6: {
+        struct ipv6hdr ip_hdr;
+        if (bpf_skb_load_bytes(skb, *offset, &ip_hdr, sizeof(struct ipv6hdr)) < 0) {
+            // debug_log("parse_skb_l3 2 failed:\n");
+            return -1;
+        }
+        l3->protocol = ip_hdr.nexthdr;
+        if (bpf_skb_load_bytes(skb, *offset + offsetof(struct ipv6hdr, saddr), &l3->saddr, sizeof(l3->saddr)) < 0) {
+            // debug_log("parse_skb_l3 3 failed:\n");
+            return -1;
+        }
+        if (bpf_skb_load_bytes(skb, *offset + offsetof(struct ipv6hdr, daddr), &l3->daddr, sizeof(l3->daddr)) < 0) {
+            // debug_log("parse_skb_l3 4 failed:\n");
+            return -1;
+        }
+        *offset += sizeof(struct ipv6hdr);
+        return 0;
+    }
+    default: {
+        return 0;
+    }
+    }
+
+    return 0;
+}
+
+static __always_inline int parse_skb_l4(struct __sk_buff *skb, u8 protocol, struct l4_t *l4, u32 *offset) {
+    switch (protocol) {
+        //    case IPPROTO_ICMP: {
+        //        l4->sport = 0;
+        //        l4->dport = 0;
+        //        *offset += sizeof(struct icmphdr);
+        //        return 0;
+        //     }
+        //    case IPPROTO_ICMPV6: {
+        //        l4->sport = 0;
+        //        l4->dport = 0;
+        //        *offset += sizeof(struct icmp6hdr);
+        //        return 0;
+        //     }
+    case IPPROTO_TCP: {
+        struct tcphdr tcp_hdr;
+        if (bpf_skb_load_bytes(skb, *offset, &tcp_hdr, sizeof(struct tcphdr)) < 0) {
+            // debug_log("parse_skb_l4 1 failed:\n");
+            return -1;
+        }
+        l4->sport = bpf_ntohs(tcp_hdr.source);
+        l4->dport = bpf_ntohs(tcp_hdr.dest);
+        *offset += sizeof(struct tcphdr);
+    }
+        return 0;
+    case IPPROTO_UDP: {
+        struct udphdr udp_hdr;
+        if (bpf_skb_load_bytes(skb, *offset, &udp_hdr, sizeof(struct udphdr)) < 0) {
+            // debug_log("parse_skb_l4 2 failed:\n");
+            return -1;
+        }
+        l4->sport = bpf_ntohs(udp_hdr.source);
+        l4->dport = bpf_ntohs(udp_hdr.dest);
+        *offset += sizeof(struct udphdr);
+        return 0;
+    }
+    case IPPROTO_SCTP: {
+        struct sctphdr sctp_hdr;
+        if (bpf_skb_load_bytes(skb, *offset, &sctp_hdr, sizeof(struct sctphdr)) < 0) {
+            // debug_log("parse_skb_l4 3 failed:\n");
+            return -1;
+        }
+        l4->sport = bpf_ntohs(sctp_hdr.source);
+        l4->dport = bpf_ntohs(sctp_hdr.dest);
+        *offset += sizeof(struct sctphdr);
+        return 0;
+    }
+    default: {
+        return 0;
+    }
+    }
+
+    return 0;
+}
+
+static __always_inline int parse_skb_meta(struct __sk_buff *skb, struct packet_meta_t *meta) {
+    meta->ifindex = skb->ifindex;
+
+    if (parse_skb_l2(skb, &meta->l2, &meta->offset) < 0) {
+        return -1;
+    }
+
+    if (parse_skb_l3(skb, meta->l2.h_protocol, &meta->l3, &meta->offset) < 0) {
+        return -1;
+    }
+
+    if (parse_skb_l4(skb, meta->l3.protocol, &meta->l4, &meta->offset) < 0) {
+        return -1;
+    }
+    return 0;
+}
+
+static __always_inline void fill_sk_meta(struct sock *sk, struct flow_pid_key_t *meta) {
+    BPF_CORE_READ_INTO(&meta->sport, sk, __sk_common.skc_num);
+    u32 family = BPF_CORE_READ(sk, __sk_common.skc_family);
+    switch (family) {
+    case AF_INET: {
+        bpf_probe_read(&meta->saddr, sizeof(sk->__sk_common.skc_rcv_saddr), &sk->__sk_common.skc_rcv_saddr);
+        break;
+    }
+    case AF_INET6: {
+        bpf_probe_read(&meta->saddr, sizeof(sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32),
+                       &sk->__sk_common.skc_v6_rcv_saddr.in6_u.u6_addr32);
+        break;
+    }
+    default: {
+        break;
+    }
+    }
+}
+
+#endif /* __PTCPDUMP_NET_H__ */