ecen5273

// eBPF IPsec
// Jack Sangdahl
// 28 November 2024
#include <vmlinux.h>

#include <bpf/bpf_endian.h>
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

#define TC_ACT_OK 0
#define ETH_P_IP 0x0800 /* Internet Protocol packet */
#define UDP_HDR_LEN 8   /* UDP header size */
#define UDP_PORT 12345  /* Port number to filter */

struct {
  __uint(type, BPF_MAP_TYPE_ARRAY);
  __uint(max_entries, 1);
  __type(key, __u32);
  __type(value, __u64);
} ingress_counter SEC(".maps");

struct {
  __uint(type, BPF_MAP_TYPE_ARRAY);
  __uint(max_entries, 1);
  __type(key, __u32);
  __type(value, __u64);
} egress_counter SEC(".maps");

struct {
  __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
  __uint(max_entries, 1);
  __type(key, __u32);
  __type(value, __u32);
} shared_key SEC(".maps");

struct {
  __uint(type, BPF_MAP_TYPE_PERCPU_ARRAY);
  __uint(max_entries, 1);
  __type(key, __u32);
  __type(value, __u32);
} key_ready SEC(".maps");

static __always_inline int handle_packet(struct __sk_buff *ctx, bool ingress);
static __always_inline void count_packets(struct __sk_buff *ctx, bool ingress);
static __always_inline __u32 generate_private_key(void);
static __always_inline __u32 mod_exp(__u32 base, __u32 exp, __u32 mod);
static __always_inline __u32 extract_public_key(void *data, void *data_end);
static __always_inline void xor_encrypt_decrypt(void *payload, __u32 len,
                                                __u32 shared_key);
static __always_inline __u32 derive_shared_key(__u32 public_key,
                                               __u32 private_key);
static __always_inline void replace_with_public_key(void *data, void *data_end,
                                                    __u32 public_key);
static __always_inline void print_udp_content(void *data, void *data_end,
                                              struct udphdr *l4);

/// @tchook {"ifindex":2, "attach_point":"BPF_TC_INGRESS"}
/// @tcopts {"handle":1, "priority":1}
SEC("tc")
int tc_ingress(struct __sk_buff *ctx) { return handle_packet(ctx, true); }

/// @tchook {"ifindex":2, "attach_point":"BPF_TC_EGRESS"}
/// @tcopts {"handle":1, "priority":1}
SEC("tc")
int tc_egress(struct __sk_buff *ctx) { return handle_packet(ctx, false); }

static __always_inline int handle_packet(struct __sk_buff *ctx, bool ingress) {
  void *data_end = (void *)(__u64)ctx->data_end;
  void *data = (void *)(__u64)ctx->data;
  struct ethhdr *l2;
  struct iphdr *l3;
  struct udphdr *l4;
  if (ctx->protocol != bpf_htons(ETH_P_IP))
    return TC_ACT_OK;
  l2 = data;
  if ((void *)(l2 + 1) > data_end)
    return TC_ACT_OK;
  l3 = (struct iphdr *)(l2 + 1);
  if ((void *)(l3 + 1) > data_end)
    return TC_ACT_OK;
  if (l3->protocol != IPPROTO_UDP)
    return TC_ACT_OK;
  l4 = (struct udphdr *)(l3 + 1);
  if ((void *)(l4 + 1) > data_end)
    return TC_ACT_OK;
  if (bpf_ntohs(l4->source) != UDP_PORT && bpf_ntohs(l4->dest) != UDP_PORT)
    return TC_ACT_OK;
  bpf_printk("Received UDP packet on port 12345: src_port: %d, dst_port: %d",
             bpf_ntohs(l4->source), bpf_ntohs(l4->dest));
  __u32 key = 0;
  __u32 *shared_key_ptr = bpf_map_lookup_elem(&shared_key, &key);
  __u32 *key_ready_ptr = bpf_map_lookup_elem(&key_ready, &key);
  if (shared_key_ptr && key_ready_ptr && *key_ready_ptr == 1) {
    void *payload = (void *)(l4 + 1);
    __u16 payload_len = bpf_ntohs(l4->len) - UDP_HDR_LEN;
    if (payload + payload_len <= data_end)
      if (ingress) {
        xor_encrypt_decrypt(payload, payload_len, *shared_key_ptr);
        bpf_printk("Decrypted UDP payload: %s", (char *)payload);
      } else {
        xor_encrypt_decrypt(payload, payload_len, *shared_key_ptr);
        bpf_printk("Encrypted UDP payload: %s", (char *)payload);
      }
  }
  count_packets(ctx, ingress);
  return TC_ACT_OK;
}

static __always_inline void count_packets(struct __sk_buff *ctx, bool ingress) {
  __u32 key = 0;
  __u64 *count;
  if (ingress) {
    count = bpf_map_lookup_elem(&ingress_counter, &key);
    if (count) {
      __sync_fetch_and_add(count, 1);
      bpf_printk("Ingress packet count: %llu", *count);
    }
  } else {
    count = bpf_map_lookup_elem(&egress_counter, &key);
    if (count) {
      __sync_fetch_and_add(count, 1);
      bpf_printk("Egress packet count: %llu", *count);
    }
  }
}

static __always_inline __u32 generate_private_key(void) {
  return bpf_get_prandom_u32() % 23;
}

static __always_inline void xor_encrypt_decrypt(void *payload, __u32 len,
                                                __u32 shared_key) {
  __u8 *data = (__u8 *)payload;
  for (__u32 i = 0; i < len; i++)
    data[i] ^= ((__u8 *)&shared_key)[i % sizeof(shared_key)];
}

static __always_inline __u32 mod_exp(__u32 base, __u32 exp, __u32 mod) {
  __u32 result = 1;
  while (exp) {
    if (exp % 2 == 1)
      result = (result * base) % mod;
    base = (base * base) % mod;
    exp /= 2;
  }
  return result;
}

static __always_inline __u32 extract_public_key(void *data, void *data_end) {
  char buf[16] = {};
  if ((data + 16) > data_end)
    return 0;
  bpf_probe_read_kernel_str(buf, sizeof(buf), data);
  unsigned long result;
  int ret = bpf_strtoul(buf, result, 10, 0);
  if (ret < 0)
    return ret;
  return result;
}

static __always_inline __u32 derive_shared_key(__u32 public_key,
                                               __u32 private_key) {
  return mod_exp(public_key, private_key, 23);
}

static __always_inline void replace_with_public_key(void *data, void *data_end,
                                                    __u32 public_key) {
  char buf[16] = {};
  int len = bpf_snprintf(buf, sizeof(buf), "%u", (__u64 *)&public_key,
                         sizeof(public_key));
  if ((data + len) <= data_end) {
    bpf_probe_read_kernel_str(data, len, buf);
    bpf_printk("Replaced payload with public key: %s", buf);
  }
}

static __always_inline void print_udp_content(void *data, void *data_end,
                                              struct udphdr *l4) {
  void *payload = (void *)(l4 + 1);
  if (payload >= data_end)
    return;
  __u16 payload_len = bpf_ntohs(l4->len) - UDP_HDR_LEN;
  if (payload + payload_len > data_end)
    return;
  char buf[33] = {};
  __u32 len = payload_len < 32 ? payload_len : 32;
  if (payload + len <= data_end) {
    bpf_probe_read_kernel(buf, len, payload);
    buf[len] = '\0';
    bpf_printk("UDP payload: %s", buf);
  }
}

char __license[] SEC("license") = "GPL";