Experiment with kprobe to modify BPF helper behaviour

bpf/egress.c

@@ -6,6 +6,7 @@
 #include "common/defs.h"
 #include "common/try.h"
 #include "kmod/main.h"
+#include "kmod/kprobe.h"
 #include "main.h"
 
 // Extend socket buffer and move n bytes from front to back.
@@ -159,7 +160,9 @@ int egress_handler(struct __sk_buff* skb) {
   csum_diff = bpf_csum_diff((__be32*)&old_ph, sizeof(old_ph), (__be32*)&new_ph, sizeof(new_ph), 0);
   bpf_l4_csum_replace(skb, csum_off, 0, csum_diff, BPF_F_PSEUDO_HDR);
 
-  mimic_change_csum_offset(skb, IPPROTO_TCP);
+  // bpf_skb_change_proto(skb, IPPROTO_TCP, MAGIC_FLAG);
+  // log_info("%llu", retval);
+  // mimic_change_csum_offset(skb, IPPROTO_TCP);
 
   return TC_ACT_OK;
 }

kmod/Makefile

@@ -1,5 +1,5 @@
 obj-m += mimic.o
-mimic-objs := main.o
+mimic-objs := main.o kprobe.o
 ccflags-y := -D_MIMIC_KMOD
 
 KERNEL_UNAME ?= $(shell uname -r)

kmod/kprobe.c

@@ -0,0 +1,74 @@
+#include <asm/ptrace.h>
+#include <linux/bpf.h>
+#include <linux/filter.h>
+#include <linux/kprobes.h>
+#include <linux/printk.h>
+#include <linux/stddef.h>
+#include <linux/tcp.h>
+#include <linux/udp.h>
+
+#include "kprobe.h"
+
+struct bpf_skb_change_proto_params {
+  struct sk_buff* skb;
+  __be16 proto;
+  u64 flags;
+};
+
+static int entry_handler(struct kretprobe_instance* ri, struct pt_regs* regs) {
+  struct bpf_skb_change_proto_params* params = (typeof(params))ri->data;
+  params->skb = (void*)regs_get_kernel_argument(regs, 0);
+  params->proto = regs_get_kernel_argument(regs, 1);
+#if defined(CONFIG_32BIT)
+  unsigned long a2 = regs_get_kernel_argument(regs, 2);
+  unsigned long a3 = regs_get_kernel_argument(regs, 3);
+#if defined(CONFIG_CPU_BIG_ENDIAN)
+  params->flags = ((u64)a2 << 32) + a3;
+#else
+  params->flags = ((u64)a3 << 32) + a2;
+#endif
+#else
+  params->flags = regs_get_kernel_argument(regs, 2);
+#endif
+  return 0;
+}
+NOKPROBE_SYMBOL(entry_handler);
+
+static int ret_handler(struct kretprobe_instance* ri, struct pt_regs* regs) {
+  unsigned long retval = regs_return_value(regs);
+  if (retval != -EINVAL) return 0;
+
+  struct bpf_skb_change_proto_params* params = (typeof(params))ri->data;
+  if (!params->skb || params->flags != MAGIC_FLAG) return 0;
+
+  if (params->skb->ip_summed != CHECKSUM_PARTIAL) {
+    regs_set_return_value(regs, -1);
+    return 0;
+  }
+
+  switch (params->proto) {
+    case IPPROTO_TCP:
+      printk("proto TCP, prev offset = %d\n", params->skb->csum_offset);
+      params->skb->csum_offset = offsetof(struct tcphdr, check);
+      // printk("proto TCP, after offset = %d\n", params->skb->csum_offset);
+      break;
+    case IPPROTO_UDP:
+      printk("proto UDP\n");
+      params->skb->csum_offset = offsetof(struct udphdr, check);
+      break;
+    default:
+      regs_set_return_value(regs, -1);
+      return 0;
+  }
+  regs_set_return_value(regs, 0);
+  return 0;
+}
+NOKPROBE_SYMBOL(ret_handler);
+
+struct kretprobe kp = {
+  .kp.symbol_name = "bpf_skb_change_proto",
+  .entry_handler = entry_handler,
+  .handler = ret_handler,
+  .data_size = sizeof(struct bpf_skb_change_proto_params),
+  .maxactive = 32,
+};

kmod/kprobe.h

@@ -0,0 +1,14 @@
+#ifndef _MIMIC_KMOD_KPROBE_H
+#define _MIMIC_KMOD_KPROBE_H
+
+#ifdef _MIMIC_KMOD
+#include <linux/kprobes.h>
+#endif
+
+#define MAGIC_FLAG 0x4eb37b03751ff785
+
+#ifdef _MIMIC_KMOD
+extern struct kretprobe kp;
+#endif
+
+#endif  // _MIMIC_KMOD_KPROBE_H

kmod/main.c

@@ -4,8 +4,10 @@
 #include <linux/skbuff.h>
 #include <linux/tcp.h>
 #include <linux/udp.h>
+#include <linux/kprobes.h>
 
 #include "main.h"
+#include "kprobe.h"
 
 #ifndef __bpf_kfunc
 #define __bpf_kfunc __used noinline
@@ -74,10 +76,13 @@ static const struct btf_kfunc_id_set mimic_tc_kfunc_set = {
 static int __init mimic_init(void) {
   int ret = 0;
   ret = ret ?: register_btf_kfunc_id_set(BPF_PROG_TYPE_SCHED_CLS, &mimic_tc_kfunc_set);
+  ret = ret ?: register_kretprobe(&kp);
   return ret;
 }
 
-static void __exit mimic_exit(void) {}
+static void __exit mimic_exit(void) {
+  unregister_kretprobe(&kp);
+}
 
 module_init(mimic_init);
 module_exit(mimic_exit);