From: Matt Corallo <git@bluematt.me>
Date: Sat, 3 Apr 2021 22:09:57 +0000 (-0400)
Subject: Only parse v4/v6 if we have relevant rules for them
X-Git-Url: http://git.bitcoin.ninja/?a=commitdiff_plain;h=b6de860db504248c8eb737887e10a334a2b5c298;p=flowspec-xdp

Only parse v4/v6 if we have relevant rules for them
---

diff --git a/genrules.py b/genrules.py
index c0a5311..82b8d82 100755
--- a/genrules.py
+++ b/genrules.py
@@ -261,6 +261,9 @@ with open("rules.h", "w") as out:
             assert False
         out.write("#define REQ_8021Q " + args.vlan_tag + "\n")
 
+    use_v4 = False
+    use_v6 = False
+
     out.write("#define RULES \\\n")
 
     def write_rule(r):
@@ -272,10 +275,12 @@ with open("rules.h", "w") as out:
             continue
         if t[0].strip() == "flow4":
             proto = 4
+            use_v4 = True
             out.write("if (eth_proto == htons(ETH_P_IP)) { \\\n")
             out.write("\tdo {\\\n")
         elif t[0].strip() == "flow6":
             proto = 6
+            use_v6 = True
             out.write("if (eth_proto == htons(ETH_P_IPV6)) { \\\n")
             out.write("\tdo {\\\n")
         else:
@@ -322,3 +327,7 @@ with open("rules.h", "w") as out:
         out.write("\t} while(0);\\\n}\\\n")
 
     out.write("\n")
+    if use_v4:
+        out.write("#define NEED_V4_PARSE\n")
+    if use_v6:
+        out.write("#define NEED_V6_PARSE\n")
diff --git a/xdp.c b/xdp.c
index 669d3eb..b93b169 100644
--- a/xdp.c
+++ b/xdp.c
@@ -170,14 +170,21 @@ int xdp_drop_prog(struct xdp_md *ctx)
 		}
 	}
 
-	const struct tcphdr *tcp = NULL;
-	const struct udphdr *udp = NULL;
+#ifdef NEED_V4_PARSE
+	const struct iphdr *ip = NULL;
 	const struct icmphdr *icmp = NULL;
+#endif
+#ifdef NEED_V6_PARSE
+	const struct ip6hdr *ip6 = NULL;
 	const struct icmp6hdr *icmpv6 = NULL;
 	const struct ip6_fraghdr *frag6 = NULL;
-	const struct iphdr *ip = NULL;
-	const struct ip6hdr *ip6 = NULL;
+#endif
+
 	const void *l4hdr = NULL;
+	const struct tcphdr *tcp = NULL;
+	const struct udphdr *udp = NULL;
+
+#ifdef NEED_V4_PARSE
 	if (eth_proto == BE16(ETH_P_IP)) {
 		if (unlikely(pktdata + sizeof(struct iphdr) > data_end))
 			return XDP_DROP;
@@ -204,7 +211,10 @@ int xdp_drop_prog(struct xdp_md *ctx)
 				return XDP_DROP;
 			icmp = (struct icmphdr*) l4hdr;
 		}
-	} else if (eth_proto == BE16(ETH_P_IPV6)) {
+	}
+#endif
+#ifdef NEED_V6_PARSE
+	if (eth_proto == BE16(ETH_P_IPV6)) {
 		if (unlikely(pktdata + sizeof(struct ip6hdr) > data_end))
 			return XDP_DROP;
 		ip6 = (struct ip6hdr*) pktdata;
@@ -236,9 +246,8 @@ int xdp_drop_prog(struct xdp_md *ctx)
 			icmpv6 = (struct icmp6hdr*) l4hdr;
 		}
 		// TODO: Handle some options?
-	} else {
-		return XDP_PASS;
 	}
+#endif
 
 	uint16_t sport, dport; // Host Endian! Only valid with tcp || udp
 	if (tcp != NULL) {