kmod: refactor, crypto init

Makefile

@@ -2,7 +2,7 @@ mimic_common_headers := $(wildcard common/*.h)
 
 mimic_bpf_src := $(wildcard bpf/*.c)
 mimic_bpf_obj := $(mimic_bpf_src:.c=.o)
-mimic_bpf_headers := bpf/vmlinux.h $(wildcard bpf/*.h) $(mimic_common_headers)
+mimic_bpf_headers := bpf/vmlinux.h $(wildcard bpf/*.h) $(wildcard kmod/*.h) $(mimic_common_headers)
 
 mimic_src := $(wildcard src/*.c)
 mimic_obj := $(mimic_src:.c=.o)

bpf/ingress.c

@@ -128,9 +128,9 @@ int ingress_handler(struct xdp_md* xdp) {
   struct filter_settings* settings = matches_whitelist(QUARTET_TCP);
   if (!settings) return XDP_PASS;
   struct conn_tuple conn_key = gen_conn_key(QUARTET_TCP);
-  __u32 payload_len = ip_payload_len - (tcp->doff << 2);
+  __u32 tcp_payload_len = ip_payload_len - (tcp->doff << 2);
 
-  log_tcp(true, &conn_key, tcp, payload_len);
+  log_tcp(true, &conn_key, tcp, tcp_payload_len);
   struct connection* conn = bpf_map_lookup_elem(&mimic_conns, &conn_key);
 
   struct tcp_options opt = {};
@@ -193,7 +193,7 @@ int ingress_handler(struct xdp_md* xdp) {
         conn->state = CONN_SYN_RECV;
         flags |= TCP_FLAG_SYN | TCP_FLAG_ACK;
         seq = conn->seq = random;
-        ack_seq = conn->ack_seq = next_ack_seq(tcp, payload_len);
+        ack_seq = conn->ack_seq = next_ack_seq(tcp, tcp_payload_len);
         conn->seq += 1;
         conn->peer_mss = opt.mss;
       } else {
@@ -203,7 +203,7 @@ int ingress_handler(struct xdp_md* xdp) {
 
     case CONN_SYN_RECV:
       if (likely(tcp->syn && !tcp->ack)) {
-        __u32 new_ack_seq = next_ack_seq(tcp, payload_len);
+        __u32 new_ack_seq = next_ack_seq(tcp, tcp_payload_len);
         if (unlikely(new_ack_seq != conn->ack_seq)) goto fsm_error;
         flags |= TCP_FLAG_SYN | TCP_FLAG_ACK;
         seq = conn->seq++;
@@ -212,7 +212,7 @@ int ingress_handler(struct xdp_md* xdp) {
       } else if (likely(!tcp->syn && tcp->ack)) {
         will_send_ctrl_packet = false;
         conn->state = CONN_ESTABLISHED;
-        conn->ack_seq = next_ack_seq(tcp, payload_len);
+        conn->ack_seq = next_ack_seq(tcp, tcp_payload_len);
         conn->cooldown_mul = 0;
         newly_estab = true;
         swap(pktbuf, conn->pktbuf);
@@ -235,7 +235,7 @@ int ingress_handler(struct xdp_md* xdp) {
           conn->state = CONN_SYN_RECV;
         }
         seq = conn->seq;
-        ack_seq = conn->ack_seq = next_ack_seq(tcp, payload_len);
+        ack_seq = conn->ack_seq = next_ack_seq(tcp, tcp_payload_len);
         conn->peer_mss = opt.mss;
         conn->cwnd = cwnd = 44 * opt.mss;
       } else {
@@ -253,17 +253,17 @@ int ingress_handler(struct xdp_md* xdp) {
         seq = conn->seq;
         ack_seq = conn->ack_seq;
         cwnd = conn->cwnd;
-      } else if (conn->keepalive_sent && payload_len == 0) {
+      } else if (conn->keepalive_sent && tcp_payload_len == 0) {
         // Received keepalive ACK
         will_send_ctrl_packet = false;
         is_keepalive = true;
         conn->keepalive_sent = false;
-      } else if (!tcp->psh && payload_len == 0) {
+      } else if (!tcp->psh && tcp_payload_len == 0) {
         // Empty segment without PSH will be treated as control packet
         will_send_ctrl_packet = false;
       } else br_likely {
         will_send_ctrl_packet = will_drop = false;
-        conn->ack_seq += payload_len;
+        conn->ack_seq += tcp_payload_len;
         __u32 peer_mss = conn->peer_mss ?: 1460;
         __u32 upper_bound = 20 * peer_mss;
         __u32 lower_bound = 2 * peer_mss;
@@ -274,7 +274,7 @@ int ingress_handler(struct xdp_md* xdp) {
           ack_seq = conn->ack_seq;
           conn->cwnd = cwnd = 44 * peer_mss;
         }
-        conn->cwnd -= payload_len;
+        conn->cwnd -= tcp_payload_len;
       }
       break;
 

kmod/Makefile

@@ -9,9 +9,11 @@ CHECKSUM_HACK ?= kfunc
 ifeq ($(filter kfunc kprobe,$(CHECKSUM_HACK)),)
 $(error unknown checksum hack '$(CHECKSUM_HACK)')
 endif
-mimic-objs += csum-hack-$(CHECKSUM_HACK).o
+mimic-objs += $(addprefix $(CHECKSUM_HACK)/,init.o crypto.o csum-hack.o)
 ccflags-y += -DMIMIC_CHECKSUM_HACK_$(CHECKSUM_HACK)
 
+$(info $(mimic-objs))
+
 ifneq ($(CHECKSUM_HACK),kfunc)
 build_dir := $(SYSTEM_BUILD_DIR)
 else

kmod/crypto.h

@@ -0,0 +1,36 @@
+#ifndef _MIMIC_KMOD_CRYPTO_H
+#define _MIMIC_KMOD_CRYPTO_H
+
+#if defined(_MIMIC_KMOD)
+#include <crypto/skcipher.h>
+#include <linux/refcount.h>
+#elif defined(_MIMIC_BPF)
+// clang-format off
+#include "bpf/vmlinux.h"
+#include <bpf/bpf_helpers.h>
+// clang-format on
+#endif
+
+struct mimic_crypto_state {
+  refcount_t rc;
+  struct crypto_skcipher* tfm;
+};
+
+#ifdef _MIMIC_BPF
+
+#if defined(MIMIC_CHECKSUM_HACK_kfunc)
+struct mimic_crypto_state* mimic_crypto_state_create(void) __ksym;
+int mimic_crypto_set_key(struct mimic_crypto_state* state, void* key, __u32 key__sz) __ksym;
+void mimic_crypto_state_release(struct mimic_crypto_state* state) __ksym;
+int mimic_encrypt_wg_header(struct __sk_buff* skb_bpf, __u32 offset, void* iv, __u32 iv__sz,
+                            struct mimic_crypto_state* state) __ksym;
+int mimic_decrypt_wg_header(struct xdp_md* xdp_bpf, __u32 offset, void* iv, __u32 iv__sz,
+                            struct mimic_crypto_state* state) __ksym;
+
+#elif defined(MIMIC_CHECKSUM_HACK_kprobe)
+#error to be implemented
+#endif
+
+#endif
+
+#endif  // _MIMIC_KMOD_CRYPTO_H

kmod/csum-hack.h

@@ -13,8 +13,8 @@
 #endif
 
 #ifdef _MIMIC_KMOD
-int csum_hack_init(void);
-void csum_hack_exit(void);
+int csum_impl_init(void);
+void csum_impl_exit(void);
 #endif
 
 #define MAGIC_FLAG1 0xfc9e39d5

kmod/impl.h

@@ -0,0 +1,7 @@
+#ifndef _MIMIC_KMOD_IMPL_H
+#define _MIMIC_KMOD_IMPL_H
+
+int impl_init(void);
+void impl_exit(void);
+
+#endif  // _MIMIC_KMOD_IMPL_H

kmod/kfunc/common.h

@@ -0,0 +1,30 @@
+#ifndef _MIMIC_KMOD_KFUNC_KFUNC_H
+#define _MIMIC_KMOD_KFUNC_KFUNC_H
+
+#include <linux/btf.h>  // IWYU pragma: export
+
+#ifndef __bpf_kfunc
+#define __bpf_kfunc __used noinline
+#endif
+
+#ifndef __bpf_kfunc_start_defs
+#define __bpf_kfunc_start_defs()                                          \
+  __diag_push();                                                          \
+  __diag_ignore_all("-Wmissing-declarations",                             \
+                    "Global kfuncs as their definitions will be in BTF"); \
+  __diag_ignore_all("-Wmissing-prototypes", "Global kfuncs as their definitions will be in BTF")
+#endif
+
+#ifndef __bpf_kfunc_end_defs
+#define __bpf_kfunc_end_defs() __diag_pop()
+#endif
+
+#ifndef BTF_KFUNCS_START
+#define BTF_KFUNCS_START BTF_SET8_START
+#endif
+
+#ifndef BTF_KFUNCS_END
+#define BTF_KFUNCS_END BTF_SET8_END
+#endif
+
+#endif  // _MIMIC_KMOD_KFUNC_KFUNC_H

kmod/kfunc/crypto.c

@@ -0,0 +1,93 @@
+#include <crypto/skcipher.h>
+#include <linux/bpf.h>
+#include <linux/btf.h>
+#include <linux/btf_ids.h>
+#include <linux/cfi.h>
+#include <linux/crypto.h>
+#include <linux/err.h>
+#include <linux/gfp_types.h>
+#include <linux/random.h>
+#include <linux/refcount.h>
+#include <linux/scatterlist.h>
+#include <linux/skbuff.h>
+#include <linux/slab.h>
+#include <linux/types.h>
+#include <net/xdp.h>
+
+#include "../crypto.h"
+#include "crypto.h"
+
+static int skcipher(void* data, size_t size, int (*process)(struct skcipher_request*)) {
+  int ret = 0;
+  struct crypto_skcipher* tfm = NULL;
+  struct skcipher_request* req = NULL;
+  struct scatterlist sg;
+
+  if (IS_ERR(tfm = crypto_alloc_skcipher("chacha20", 0, 0))) {
+    pr_err("error allocating chacha20 handle: %ld\n", PTR_ERR(tfm));
+    return PTR_ERR(tfm);
+  }
+
+  __u8 iv[16] = {}, key[32] = {};
+  // get_random_bytes(iv, sizeof(iv));
+  // get_random_bytes(key, sizeof(key));
+
+  if ((ret = crypto_skcipher_setkey(tfm, key, sizeof(key))) < 0) {
+    pr_err("error setting key: %d\n", ret);
+    goto cleanup;
+  }
+
+  req = skcipher_request_alloc(tfm, GFP_ATOMIC);
+  if (!req) {
+    ret = -ENOMEM;
+    goto cleanup;
+  }
+
+  sg_init_one(&sg, data, size);
+  skcipher_request_set_crypt(req, &sg, &sg, size, iv);
+  if ((ret = process(req)) < 0) {
+    pr_err("error encrypting/decrypting data: %d\n", ret);
+    goto cleanup;
+  }
+
+  ret = 0;
+cleanup:
+  crypto_free_skcipher(tfm);
+  skcipher_request_free(req);
+  return ret;
+}
+
+struct mimic_crypto_state* mimic_crypto_state_create(void) {
+  struct mimic_crypto_state* state = kzalloc(sizeof(*state), GFP_KERNEL);
+  state->rc = (typeof(state->rc))REFCOUNT_INIT(1);
+  state->tfm = crypto_alloc_skcipher("chacha20", 0, 0);
+  if (IS_ERR(state->tfm)) {
+    kfree(state);
+    return NULL;
+  }
+  return state;
+}
+
+int mimic_crypto_set_key(struct mimic_crypto_state* state, void* key, __u32 key__sz) {
+  return crypto_skcipher_setkey(state->tfm, key, key__sz);
+}
+
+void mimic_crypto_state_release(struct mimic_crypto_state* state) {
+  if (refcount_dec_and_test(&state->rc)) {
+    crypto_free_skcipher(state->tfm);
+    kfree(state);
+  }
+}
+
+void mimic_crypto_state_dtor(void* p) { mimic_crypto_state_release(p); }
+CFI_NOSEAL(mimic_crypto_state_dtor);
+
+int mimic_encrypt_wg_header(struct __sk_buff* skb_bpf, __u32 offset, void* iv, __u32 iv__sz, struct mimic_crypto_state* state) {
+  struct sk_buff* skb = (typeof(skb))skb_bpf;
+  return skcipher(skb->data + offset, 16, crypto_skcipher_encrypt);
+}
+
+int mimic_decrypt_wg_header(struct xdp_md* xdp_bpf, __u32 offset, void* iv, __u32 iv__sz, struct mimic_crypto_state* state) {
+  struct xdp_buff* xdp = (typeof(xdp))xdp_bpf;
+  return skcipher(xdp->data + offset, 16, crypto_skcipher_decrypt);
+}

kmod/kfunc/crypto.h

@@ -0,0 +1,21 @@
+#ifndef _MIMIC_KMOD_KFUNC_CRYPTO_H
+#define _MIMIC_KMOD_KFUNC_CRYPTO_H
+
+#include <linux/bpf.h>
+#include <linux/btf.h>
+
+#include "../crypto.h"
+#include "common.h"
+
+__bpf_kfunc_start_defs();
+
+__bpf_kfunc struct mimic_crypto_state* mimic_crypto_state_create(void);
+__bpf_kfunc int mimic_crypto_set_key(struct mimic_crypto_state* state, void* key, __u32 key__sz);
+__bpf_kfunc void mimic_crypto_state_release(struct mimic_crypto_state* state);
+__bpf_kfunc void mimic_crypto_state_dtor(void* p);
+__bpf_kfunc int mimic_encrypt_wg_header(struct __sk_buff* skb_bpf, __u32 offset, void* iv, __u32 iv__sz, struct mimic_crypto_state* state);
+__bpf_kfunc int mimic_decrypt_wg_header(struct xdp_md* xdp_bpf, __u32 offset, void* iv, __u32 iv__sz, struct mimic_crypto_state* state);
+
+__bpf_kfunc_end_defs();
+
+#endif  // _MIMIC_KMOD_KFUNC_CRYPTO_H