From 4b9bd48053e27068025ddeccc1ea4d7460678521 Mon Sep 17 00:00:00 2001
From: Quentin Dufour <quentin@deuxfleurs.fr>
Date: Sun, 25 Apr 2021 16:56:35 +0200
Subject: [PATCH] Working checksum computation

---
 xdp/test.c    | 57 +++++++++++++++++++++------------------------------
 xdp/xdp_udp.c | 30 +++++++++++----------------
 2 files changed, 35 insertions(+), 52 deletions(-)

diff --git a/xdp/test.c b/xdp/test.c
index 80a909a..fe02c8f 100644
--- a/xdp/test.c
+++ b/xdp/test.c
@@ -20,40 +20,30 @@ char pkt3[] = { 0x50,0x9a,0x4c,0x24,0x1a,0x72,0x04,0xf0,0x21,0x65,0x18,0x12,0x08
 
 char csDesired[] = { 0x38, 0xf1 };
 
-// Checksum utilities
-__attribute__((__always_inline__))
-static inline __u16 csum_fold_helper(__u64 csum) {
-  int i;
-  #pragma unroll
-  for (i = 0; i < 4; i ++) {
-    if (csum >> 16)
-      csum = (csum & 0xffff) + (csum >> 16);
-  }
-  return ~csum;
+static __always_inline __u16 csum_fold_helper(__u32 csum)
+{
+	return ~((csum & 0xffff) + (csum >> 16));
+}
+static __always_inline void ipv4_csum(void *data_start, int data_size, __u32 *csum) {
+	*csum = bpf_csum_diff(0, 0, data_start, data_size, *csum);
+	*csum = csum_fold_helper(*csum);
 }
 
-// Update checksum following RFC 1624 (Eqn. 3): https://tools.ietf.org/html/rfc1624
-//     HC' = ~(~HC + ~m + m')
-// Where :
-//   HC  - old checksum in header
-//   HC' - new checksum in header
-//   m   - old value
-//   m'  - new value
-__attribute__((__always_inline__))
-static inline void update_csum(__u64 *csum, __be32 old_flag,__be32 new_flag ) {
-    // ~HC 
-    *csum = ~*csum;
-    *csum = *csum & 0xffff;
-    // + ~m
-    __u32 tmp;
-    tmp = ~old_addr;
-    *csum += tmp;
-    // + m
-    *csum += new_addr;
-    // then fold and complement result ! 
-    *csum = csum_fold_helper(*csum);
+int checksum(__u16* ipbuf, __u32 len) {
+	// https://tools.ietf.org/html/rfc1071
+	__u32 sum = 0;
+	for (int i = 0; i < sizeof(struct iphdr)/2; i++) {
+		sum += ipbuf[i];
+	}
+
+	sum = (sum & 0xffff) + (sum >> 16);
+	sum = (sum & 0xffff) + (sum >> 16);
+
+	return ~sum;
 }
 
+// @TODO: Implement https://tools.ietf.org/html/rfc1624 instead
+
 
 int main(void) {
 	void *data = pkt3;
@@ -62,10 +52,9 @@ int main(void) {
 	struct ethhdr *eth = data; //structure ethernet from if_ether.h
 	struct iphdr *ip = (struct iphdr*)(eth + 1);
 	printf("%#06x\n", bpf_ntohs(ip->check));
-	__u64 cs = iph->check;
-	update_csum(&cs, old_addr, new_addr);
-	iph->check = cs;
-	printf("new cs: %x\n", ip->check);
+	ip->check = 0;
+	ip->check = checksum((__u16*)ip, sizeof(*ip));
+	printf("%#06x\n", bpf_ntohs(ip->check));
 }
 
 int main2(void) {
diff --git a/xdp/xdp_udp.c b/xdp/xdp_udp.c
index 2cdca60..c766886 100644
--- a/xdp/xdp_udp.c
+++ b/xdp/xdp_udp.c
@@ -32,24 +32,18 @@
                          ##__VA_ARGS__);                        \
 })
 
-static inline unsigned short checksum(unsigned short *buf, int bufsz) {
-    unsigned long sum = 0;
+static __always_inline int checksum_ipv4(__u16* ipbuf) {
+	// https://tools.ietf.org/html/rfc1071
+	__u32 sum = 0;
+	for (int i = 0; i < sizeof(struct iphdr)/2; i++) {
+		sum += ipbuf[i];
+	}
+	sum = (sum & 0xffff) + (sum >> 16);
+	sum = (sum & 0xffff) + (sum >> 16);
 
-    while (bufsz > 1) {
-        sum += *buf;
-        buf++;
-        bufsz -= 2;
-    }
-
-    if (bufsz == 1) {
-        sum += *(unsigned char *)buf;
-    }
-
-    sum = (sum & 0xffff) + (sum >> 16);
-    sum = (sum & 0xffff) + (sum >> 16);
-
-    return ~sum;
+	return ~sum;
 }
+// @TODO: Implement https://tools.ietf.org/html/rfc1624 instead
 
 SEC("xdp_udp")
 int xdp_udp_func(struct xdp_md *ctx)
@@ -74,8 +68,8 @@ int xdp_udp_func(struct xdp_md *ctx)
 
 		if(ip->protocol == IPPROTO_UDP) {
 			ip->frag_off |= bpf_htons(0x4000); //Set the DF flag to 1 -> 0b 0100 0000 0000 0000
-			ip->check = 0x0000;
-			ip->check = checksum((unsigned short*)ip, sizeof(ip));
+			ip->check = 0;
+			ip->check = checksum_ipv4((__u16*)ip);
 		}
 	}
 	return XDP_PASS;