Track drops

[flowspec-xdp] / genrules.py

diff --git a/genrules.py b/genrules.py
index d372806a8e8c99a6b9fed67158cbf6e04e0406ba..40fc5c1f2d84969b959e5d30fdaad8c3c29cdf20 100755 (executable)
--- a/genrules.py
+++ b/genrules.py
@@ -3,17 +3,14 @@
 import sys
 import ipaddress
 from enum import Enum
+import argparse
+
 
 IP_PROTO_ICMP = 1
 IP_PROTO_ICMPV6 = 58
 IP_PROTO_TCP = 6
 IP_PROTO_UDP = 17
 
-if len(sys.argv) > 2 and sys.argv[2].startswith("parse_ihl"):
-    PARSE_IHL = True
-else:
-    PARSE_IHL = False
-
 class ASTAction(Enum):
     OR = 1,
     AND = 2,
@@ -239,12 +236,43 @@ def flow_label_to_rule(rules):
 if (!( {ast.write("((((uint32_t)(ip6->flow_lbl[0] & 0xf)) << 2*8) | (((uint32_t)ip6->flow_lbl[1]) << 1*8) | (uint32_t)ip6->flow_lbl[0])")} )) break;"""
 
 with open("rules.h", "w") as out:
-    if len(sys.argv) > 1 and sys.argv[1] == "parse_8021q":
-        out.write("#define PARSE_8021Q\n")
-    if len(sys.argv) > 1 and sys.argv[1].startswith("req_8021q="):
-        out.write("#define PARSE_8021Q\n")
-        out.write(f"#define REQ_8021Q {sys.argv[1][10:]}\n")
-
+    parse = argparse.ArgumentParser()
+    parse.add_argument("--ihl", dest="ihl", required=True, choices=["drop-options","accept-options","parse-options"])
+    parse.add_argument("--v6frag", dest="v6frag", required=True, choices=["drop-frags","ignore","parse-frags","ignore-parse-if-rule"])
+    parse.add_argument("--8021q", dest="vlan", required=True, choices=["drop-vlan","accept-vlan","parse-vlan"])
+    parse.add_argument("--require-8021q", dest="vlan_tag")
+    args = parse.parse_args(sys.argv[1:])
+
+    if args.ihl == "drop-options":
+        out.write("#define PARSE_IHL XDP_DROP\n")
+    elif args.ihl == "accept-options":
+        out.write("#define PARSE_IHL XDP_PASS\n")
+    elif args.ihl == "parse-options":
+        out.write("#define PARSE_IHL PARSE\n")
+
+    if args.v6frag == "drop-frags":
+        out.write("#define PARSE_V6_FRAG XDP_DROP\n")
+    elif args.v6frag == "ignore":
+        pass
+    elif args.v6frag == "parse-frags":
+        out.write("#define PARSE_V6_FRAG PARSE\n")
+
+    if args.vlan == "drop-vlan":
+        out.write("#define PARSE_8021Q XDP_DROP\n")
+    elif args.vlan == "accept-vlan":
+        out.write("#define PARSE_8021Q XDP_PASS\n")
+    elif args.vlan == "parse-vlan":
+        out.write("#define PARSE_8021Q PARSE\n")
+
+    if args.vlan_tag is not None:
+        if args.vlan != "parse-vlan":
+            assert False
+        out.write("#define REQ_8021Q " + args.vlan_tag + "\n")
+
+    use_v4 = False
+    use_v6 = False
+    use_v6_frags = False
+    rulecnt = 0
 
     out.write("#define RULES \\\n")
 
@@ -257,10 +285,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:
@@ -298,12 +328,24 @@ with open("rules.h", "w") as out:
             elif step.strip().startswith("label"):
                 write_rule(flow_label_to_rule(step.strip()[6:]))
             elif step.strip().startswith("fragment"):
+                if proto == 6:
+                    use_v6_frags = True
                 write_rule(fragment_to_rule(proto, step.strip()[9:]))
             elif step.strip() == "":
                 pass
             else:
                 assert False
-        out.write("\t\treturn XDP_DROP;\\\n")
+        out.write(f"\t\tconst uint32_t ruleidx = STATIC_RULE_CNT + {rulecnt};\\\n")
+        out.write("\t\tDO_RETURN(ruleidx, XDP_DROP);\\\n")
         out.write("\t} while(0);\\\n}\\\n")
+        rulecnt += 1
 
     out.write("\n")
+    out.write(f"#define RULECNT {rulecnt}\n")
+    if use_v4:
+        out.write("#define NEED_V4_PARSE\n")
+    if use_v6:
+        out.write("#define NEED_V6_PARSE\n")
+    if args.v6frag == "ignore-parse-if-rule":
+        if use_v6_frags:
+            out.write("#define PARSE_V6_FRAG PARSE\n")