+class RuleAction(Enum):
+ CONDITIONS = 1
+ ACTION = 2
+ LIST = 3
+class RuleNode:
+ def __init__(self, ty, action, inner):
+ self.ty = ty
+ self.action = action
+ self.inner = inner
+ if ty == RuleAction.ACTION:
+ assert inner is None
+ assert type(action) == str
+ elif ty == RuleAction.LIST:
+ assert type(inner) == list
+ assert action is None
+ for item in inner:
+ assert type(item) == RuleNode
+ else:
+ assert ty == RuleAction.CONDITIONS
+ assert type(action) == list
+ assert type(inner) == RuleNode
+
+ def __lt__(self, other):
+ assert self.ty == RuleAction.CONDITIONS
+ assert other.ty == RuleAction.CONDITIONS
+
+ o = ORD_EQUAL
+
+ # RFC first has us sort by dest, then source, then other conditions. We don't implement the
+ # other conditions because doing so requires re-implementing the Flowspec wire format,
+ # which isn't trivial. However, we do implement the source/dest sorting in the hopes it
+ # allows us to group rules according to source/dest IP and hopefully LLVM optimizes out
+ # later rules.
+
+ selfdest = next(filter(lambda a : type(a) == IpRule and a.ty == "daddr", self.action), None)
+ otherdest = next(filter(lambda a : type(a) == IpRule and a.ty == "daddr", self.action), None)
+ if o == ORD_EQUAL and selfdest is not None and otherdest is not None:
+ o = selfdest.ord(otherdest)
+
+ if o == ORD_EQUAL:
+ selfsrc = next(filter(lambda a : type(a) == IpRule and a.ty == "saddr", self.action), None)
+ othersrc = next(filter(lambda a : type(a) == IpRule and a.ty == "saddr", self.action), None)
+ if selfsrc is not None and othersrc is None:
+ return True
+ elif selfsrc is None and othersrc is not None:
+ return False
+ elif selfsrc is not None and othersrc is not None:
+ o = selfsrc.ord(othersrc)
+
+ if o == ORD_LESS:
+ return True
+ return self.action < other.action
+
+ def maybe_join(self, neighbor):
+ if self.ty == RuleAction.CONDITIONS and neighbor.ty == RuleAction.CONDITIONS:
+ overlapping_conditions = [x for x in self.action if x in neighbor.action]
+ if len(overlapping_conditions) != 0:
+ us = RuleNode(RuleAction.CONDITIONS, [x for x in self.action if x not in overlapping_conditions], self.inner)
+ them = RuleNode(RuleAction.CONDITIONS, [x for x in neighbor.action if x not in overlapping_conditions], neighbor.inner)
+ self.action = overlapping_conditions
+ if self.inner.ty == RuleAction.LIST and us.action == []:
+ self.inner.inner.append(them)
+ else:
+ self.inner = RuleNode(RuleAction.LIST, None, [us, them])
+ self.inner.flatten()
+ return True
+ return False
+
+ def flatten(self):
+ # LLVM can be pretty bad at optimizing out common subexpressions. Thus, we have to do a
+ # pass here to toptimize out common subexpressions in back-to-back rules.
+ # See https://bugs.llvm.org/show_bug.cgi?id=52455
+ assert self.ty == RuleAction.LIST
+ did_update = True
+ while did_update:
+ did_update = False
+ for i in range(0, len(self.inner) - 1):
+ if self.inner[i].maybe_join(self.inner[i + 1]):
+ del self.inner[i + 1]
+ did_update = True
+ break
+
+ def write(self, out, pfx="\t"):
+ if self.ty == RuleAction.CONDITIONS:
+ out.write(pfx + "do {\\\n")
+ for cond in self.action:
+ out.write("\t" + pfx + str(cond).replace("\n", " \\\n\t" + pfx) + " \\\n")
+ self.inner.write(out, pfx)
+ out.write(pfx + "} while(0);\\\n")
+ elif self.ty == RuleAction.LIST:
+ for item in self.inner:
+ item.write(out, pfx + "\t")
+ else:
+ out.write("\t" + pfx + self.action.strip().replace("\n", " \\\n\t" + pfx) + " \\\n")