1 files changed, 21 insertions, 17 deletions

diff --git a/src/extra/checksum.c b/src/extra/checksum.c
index f367f75..33480af 100644
--- a/src/extra/checksum.c
+++ b/src/extra/checksum.c
@@ -11,11 +11,13 @@
 
 #include <stdio.h>
 #include <stdbool.h>
+#include <endian.h>
 #include <arpa/inet.h>
 #include <netinet/ip.h>
 #include <netinet/ip6.h>
 #include <netinet/tcp.h>
 
+#include <linux/netfilter/nfnetlink_queue.h>
 #include <libnetfilter_queue/libnetfilter_queue.h>
 
 #include "internal.h"
@@ -26,8 +28,13 @@ uint16_t nfq_checksum(uint32_t sum, uint16_t *buf, int size)
 		sum += *buf++;
 		size -= sizeof(uint16_t);
 	}
-	if (size)
-		sum += *(uint8_t *)buf;
+	if (size) {
+#if __BYTE_ORDER == __BIG_ENDIAN
+		sum += (uint16_t)*(uint8_t *)buf << 8;
+#else
+		sum += (uint16_t)*(uint8_t *)buf;
+#endif
+	}
 
 	sum = (sum >> 16) + (sum & 0xffff);
 	sum += (sum >>16);
@@ -35,7 +42,7 @@ uint16_t nfq_checksum(uint32_t sum, uint16_t *buf, int size)
 	return (uint16_t)(~sum);
 }
 
-uint16_t nfq_checksum_tcpudp_ipv4(struct iphdr *iph)
+uint16_t nfq_checksum_tcpudp_ipv4(struct iphdr *iph, uint16_t protonum)
 {
 	uint32_t sum = 0;
 	uint32_t iph_len = iph->ihl*4;
@@ -46,34 +53,31 @@ uint16_t nfq_checksum_tcpudp_ipv4(struct iphdr *iph)
 	sum += (iph->saddr) & 0xFFFF;
 	sum += (iph->daddr >> 16) & 0xFFFF;
 	sum += (iph->daddr) & 0xFFFF;
-	sum += htons(IPPROTO_TCP);
+	sum += htons(protonum);
 	sum += htons(len);
 
 	return nfq_checksum(sum, (uint16_t *)payload, len);
 }
 
-uint16_t nfq_checksum_tcpudp_ipv6(struct ip6_hdr *ip6h, void *transport_hdr)
+uint16_t nfq_checksum_tcpudp_ipv6(struct ip6_hdr *ip6h, void *transport_hdr,
+				  uint16_t protonum)
 {
 	uint32_t sum = 0;
-	uint32_t hdr_len = (uint32_t *)transport_hdr - (uint32_t *)ip6h;
-	uint32_t len = ip6h->ip6_plen - hdr_len;
+	uint32_t hdr_len = (uint8_t *)transport_hdr - (uint8_t *)ip6h;
+	/* Allow for extra headers before the UDP header */
+	/* TODO: Deal with routing headers */
+	uint32_t len = ntohs(ip6h->ip6_plen) - (hdr_len - sizeof *ip6h);
 	uint8_t *payload = (uint8_t *)ip6h + hdr_len;
 	int i;
 
 	for (i=0; i<8; i++) {
-		sum += (ip6h->ip6_src.s6_addr16[i] >> 16) & 0xFFFF;
-		sum += (ip6h->ip6_src.s6_addr16[i]) & 0xFFFF;
+		sum += (ip6h->ip6_src.s6_addr16[i]);
 	}
 	for (i=0; i<8; i++) {
-		sum += (ip6h->ip6_dst.s6_addr16[i] >> 16) & 0xFFFF;
-		sum += (ip6h->ip6_dst.s6_addr16[i]) & 0xFFFF;
+		sum += (ip6h->ip6_dst.s6_addr16[i]);
 	}
-	sum += htons(IPPROTO_TCP);
-	sum += htons(ip6h->ip6_plen);
+	sum += htons(protonum);
+	sum += htons(len);
 
 	return nfq_checksum(sum, (uint16_t *)payload, len);
 }
-
-/**
- * @}
- */