From 5e9ef66bdc342c6f5852ed8afcec463132e7594a Mon Sep 17 00:00:00 2001
From: Niels Werensteijn <niels.werensteijn@teamspeak.com>
Date: Wed, 12 Oct 2016 15:21:41 +0200
Subject: [PATCH] refactor calculate_cookie

---
 src/ts3init_cookie.c | 39 --------------------------------------
 src/ts3init_cookie.h |  3 ---
 src/ts3init_match.c  | 45 +++++++++++++++++++++++++++++++++++++++++---
 3 files changed, 42 insertions(+), 45 deletions(-)

diff --git a/src/ts3init_cookie.c b/src/ts3init_cookie.c
index 8dbf6ab..4e88367 100644
--- a/src/ts3init_cookie.c
+++ b/src/ts3init_cookie.c
@@ -122,42 +122,3 @@ int ts3init_calculate_cookie_ipv4(const struct iphdr *ip, const struct udphdr *u
     return 0;
 }
 
-int ts3init_calculate_cookie(const struct sk_buff *skb, const struct xt_action_param *par, 
-                             struct udphdr *udp, __u64 k0, __u64 k1, __u64* out)
-{
-    int addr_offset;
-    int addr_len;
-    void* addr_data;
-    __u8 addr_buf[2*16];
-    struct ts3init_siphash_state hash_state;
-
-    switch (par->family)
-    {
-    case NFPROTO_IPV4:
-        addr_offset = 12; /*offset to src and dst address in ipv4 header */
-        addr_len = 2*4; /*size of ipv4 address is 4 bytes */
-        break;
-
-    case NFPROTO_IPV6:
-        addr_offset = 8; /*offset to src and dst address in ipv6 header */
-        addr_len = 2*16; /*size of ipv6 address is 16 bytes */
-        break;
-    default:
-        printk(KERN_ERR KBUILD_MODNAME ": invalid family\n");
-        return -EINVAL;
-    }
-
-    addr_data = skb_header_pointer(skb, skb->network_header+addr_offset, addr_len, addr_buf);
-    if (!addr_data)
-    {
-        printk(KERN_ERR KBUILD_MODNAME ": could not load ip addresses\n");
-        return -EINVAL;
-    }
- 
-    ts3init_siphash_setup(&hash_state, k0, k1);
-    ts3init_siphash_update(&hash_state, (u8 *)addr_data, addr_len);
-    ts3init_siphash_update(&hash_state, (u8 *)&udp->source, 4);
-    *out = ts3init_siphash_finalize(&hash_state);
-    
-    return 0;
-}
diff --git a/src/ts3init_cookie.h b/src/ts3init_cookie.h
index f6cf481..6fec05b 100644
--- a/src/ts3init_cookie.h
+++ b/src/ts3init_cookie.h
@@ -20,9 +20,6 @@ __u64* ts3init_get_cookie_seed(time_t current_time, __u8 packet_index,
                 struct xt_ts3init_cookie_cache* cache,
                 const __u8* cookie_seed);
 
-int ts3init_calculate_cookie(const struct sk_buff *skb,
-                const struct xt_action_param *par, struct udphdr *udp,
-                __u64 k0, __u64 k1, __u64* out);
 int ts3init_calculate_cookie_ipv6(const struct ipv6hdr *ip, const struct udphdr *udp, 
                                   __u64 k0, __u64 k1, __u64* out);
 int ts3init_calculate_cookie_ipv4(const struct iphdr *ip, const struct udphdr *udp, 
diff --git a/src/ts3init_match.c b/src/ts3init_match.c
index 71314bf..04ae031 100644
--- a/src/ts3init_match.c
+++ b/src/ts3init_match.c
@@ -82,6 +82,46 @@ static bool check_header(const struct sk_buff *skb, const struct xt_action_param
     return true;
 }
 
+static int calculate_cookie(const struct sk_buff *skb, const struct xt_action_param *par, 
+                       struct udphdr *udp, __u64 k0, __u64 k1, __u64* out)
+{
+    switch (par->family)
+    {
+    case NFPROTO_IPV4:
+        {
+            const struct iphdr *ip;
+            struct iphdr ip_buf;
+
+            ip  = skb_header_pointer(skb, skb->network_header, sizeof(ip_buf), &ip_buf);
+            if (ip == NULL)
+            {
+                printk(KERN_ERR KBUILD_MODNAME ": could not load ipv4 addresses\n");
+                return -EINVAL;
+            }
+
+            return ts3init_calculate_cookie_ipv4(ip, udp, k0, k1, out);
+        }
+
+    case NFPROTO_IPV6:
+        {
+            const struct ipv6hdr *ip;
+            struct ipv6hdr ip_buf;
+
+            ip  = skb_header_pointer(skb, skb->network_header, sizeof(ip_buf), &ip_buf);
+            if (ip == NULL)
+            {
+                printk(KERN_ERR KBUILD_MODNAME ": could not load ipv6 addresses\n");
+                return -EINVAL;
+            }
+
+            return ts3init_calculate_cookie_ipv6(ip, udp, k0, k1, out);
+        }
+    default:
+        printk(KERN_ERR KBUILD_MODNAME ": invalid family\n");
+        return -EINVAL;
+    }
+}
+
 static bool
 ts3init_get_cookie_mt(const struct sk_buff *skb, struct xt_action_param *par)
 {
@@ -136,8 +176,7 @@ static int ts3init_get_cookie_mt_check(const struct xt_mtchk_param *par)
     return 0;
 }
 
-static bool
-ts3init_get_puzzle_mt(const struct sk_buff *skb, struct xt_action_param *par)
+static bool ts3init_get_puzzle_mt(const struct sk_buff *skb, struct xt_action_param *par)
 {
     const struct xt_ts3init_get_puzzle_mtinfo *info = par->matchinfo;
     struct ts3_init_checked_header_data header_data;
@@ -158,7 +197,7 @@ ts3init_get_puzzle_mt(const struct sk_buff *skb, struct xt_action_param *par)
 
         /* use cookie_seed and ipaddress and port to create a hash
          * (cookie) for this connection */
-        if (ts3init_calculate_cookie(skb, par, header_data.udp, cookie_seed[0], cookie_seed[1], &cookie))
+        if (calculate_cookie(skb, par, header_data.udp, cookie_seed[0], cookie_seed[1], &cookie))
             return false; /*something went wrong*/
 
         /* compare cookie with payload bytes 0-7. if equal, cookie