From fe09dece3a0e25dd3e26e620eb2b533ce30a0a20 Mon Sep 17 00:00:00 2001
From: gamemann <gamemann@gflclan.com>
Date: Fri, 12 Nov 2021 21:12:50 +0000
Subject: [PATCH] Check IP header validity everywhere else.

---
 src/xdpfw_kern.c | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/src/xdpfw_kern.c b/src/xdpfw_kern.c
index 8d98b0a..63722c3 100644
--- a/src/xdpfw_kern.c
+++ b/src/xdpfw_kern.c
@@ -149,11 +149,11 @@ int xdp_prog_main(struct xdp_md *ctx)
     // Check blacklist map.
     __u64 *blocked = NULL;
 
-    if (eth->h_proto == htons(ETH_P_IPV6))
+    if (iph6)
     {
         blocked = bpf_map_lookup_elem(&ip6_blacklist_map, &srcip6);
     }
-    else
+    else if (iph)
     {
         blocked = bpf_map_lookup_elem(&ip_blacklist_map, &iph->saddr);
     }
@@ -167,11 +167,11 @@ int xdp_prog_main(struct xdp_md *ctx)
         if (now > *blocked)
         {
             // Remove element from map.
-            if (eth->h_proto == htons(ETH_P_IPV6))
+            if (iph6)
             {
                 bpf_map_delete_elem(&ip6_blacklist_map, &srcip6);
             }
-            else
+            else if (iph)
             {
                 bpf_map_delete_elem(&ip_blacklist_map, &iph->saddr);
             }
@@ -197,11 +197,11 @@ int xdp_prog_main(struct xdp_md *ctx)
 
     struct ip_stats *ip_stats = NULL;
     
-    if (eth->h_proto == htons(ETH_P_IPV6))
+    if (iph6)
     {
         ip_stats = bpf_map_lookup_elem(&ip6_stats_map, &srcip6);
     }
-    else
+    else if (iph)
     {
         ip_stats = bpf_map_lookup_elem(&ip_stats_map, &iph->saddr);
     }
@@ -235,11 +235,11 @@ int xdp_prog_main(struct xdp_md *ctx)
         pps = new.pps;
         bps = new.bps;
 
-        if (eth->h_proto == htons(ETH_P_IPV6))
+        if (iph6)
         {
             bpf_map_update_elem(&ip6_stats_map, &srcip6, &new, BPF_ANY);
         }
-        else
+        else if (iph)
         {
             bpf_map_update_elem(&ip_stats_map, &iph->saddr, &new, BPF_ANY);
         } 
@@ -251,7 +251,7 @@ int xdp_prog_main(struct xdp_md *ctx)
     struct icmp6hdr *icmp6h = NULL;
     
     // Check protocol.
-    if (eth->h_proto == htons(ETH_P_IPV6))
+    if (iph6)
     {
         switch (iph6->nexthdr)
         {
@@ -292,7 +292,7 @@ int xdp_prog_main(struct xdp_md *ctx)
                 break;
         }
     }
-    else
+    else if (iph)
     {
         switch (iph->protocol)
         {