checkpt

firewall/rules_builder.go

@@ -29,12 +29,8 @@ type RulesBuilder struct {
 	l *slog.Logger
 }
 
-// tableRules are one direction's rules, waiting to be built into a Table
-type tableRules struct {
-	anyRules []rule
-	// protoRules holds the rules for each protocol besides proto `any`, ICMPv6 rules are kept with ICMP
-	protoRules [256][]rule
-}
+// tableRules are one direction's rules in the order they were added, waiting to be built into a Table
+type tableRules []rule
 
 func NewRulesBuilder(l *slog.Logger) *RulesBuilder {
 	return &RulesBuilder{
@@ -145,42 +141,31 @@ func parseRule(startPort, endPort int32, groups []string, host, cidr, localCidr,
 	return r, nil
 }
 
-// add keeps r for proto, or for every protocol when proto is ProtoAny
+// add keeps r for proto, or for every protocol when proto is ProtoAny. ICMPv6 rules are kept as ICMP rules.
 func (tr *tableRules) add(proto uint8, r rule) {
-	switch proto {
-	case ProtoAny:
-		tr.anyRules = append(tr.anyRules, r)
-	case iputil.IPProtocolICMPv6:
-		tr.protoRules[iputil.IPProtocolICMP] = append(tr.protoRules[iputil.IPProtocolICMP], r)
-	default:
-		tr.protoRules[proto] = append(tr.protoRules[proto], r)
+	if proto == iputil.IPProtocolICMPv6 {
+		proto = iputil.IPProtocolICMP
 	}
+	r.proto = proto
+	*tr = append(*tr, r)
 }
 
 // build turns the rules into a Table. Rules without a local cidr allow defaultLocal,
 // or any local address when defaultLocal is nil.
-func (tr *tableRules) build(defaultLocal []netip.Prefix) *Table {
-	t := &Table{}
-
-	// A rule's id is its index in t.rules.
-	add := func(rules []rule) []int {
-		ids := make([]int, 0, len(rules))
-		for _, r := range rules {
-			ids = append(ids, len(t.rules))
-			t.rules = append(t.rules, r.withDefaultLocal(defaultLocal))
+func (tr tableRules) build(defaultLocal []netip.Prefix) *Table {
+	t := &Table{rules: make([]rule, len(tr))}
+	var hasOwn [256]bool
+	for id, r := range tr {
+		t.rules[id] = r.withDefaultLocal(defaultLocal)
+		if r.proto != ProtoAny {
+			hasOwn[r.proto] = true
 		}
-		return ids
-	}
-	anyIDs := add(tr.anyRules)
-	var protoIDs [256][]int
-	for proto, rules := range tr.protoRules {
-		protoIDs[proto] = add(rules)
 	}
 
 	// Every protocol starts out sharing an index of the proto `any` rules, including protos[ProtoAny]. There is
 	// one for protocols with ports and one for protocols without.
-	anyPorts := newProtoIndex(true, t.rules, anyIDs)
-	anyNoPorts := newProtoIndex(false, t.rules, anyIDs)
+	anyPorts := newProtoIndex(true, t.rules, ProtoAny)
+	anyNoPorts := newProtoIndex(false, t.rules, ProtoAny)
 	for proto := range t.protos {
 		if hasPorts(uint8(proto)) {
 			t.protos[proto] = anyPorts
@@ -191,9 +176,9 @@ func (tr *tableRules) build(defaultLocal []netip.Prefix) *Table {
 
 	// A protocol with rules of its own gets an index of those rules plus the proto `any` rules, so a packet
 	// checks a single index.
-	for proto, ids := range protoIDs {
-		if len(ids) > 0 {
-			t.protos[proto] = newProtoIndex(hasPorts(uint8(proto)), t.rules, anyIDs, ids)
+	for proto, own := range hasOwn {
+		if own {
+			t.protos[proto] = newProtoIndex(hasPorts(uint8(proto)), t.rules, uint8(proto))
 		}
 	}
 
@@ -203,23 +188,25 @@ func (tr *tableRules) build(defaultLocal []netip.Prefix) *Table {
 	return t
 }
 
-// newProtoIndex builds an index of the rules with the given ids for a protocol with or without ports, or
-// returns nil when there are none.
-func newProtoIndex(hasPorts bool, rules []rule, idSets ...[]int) *protoIndex {
-	ids := slices.Concat(idSets...)
-	if len(ids) == 0 {
-		return nil
-	}
-
-	pi := &protoIndex{hasPorts: hasPorts, fragment: newRuleSet(len(rules))}
-	if hasPorts {
-		pi.byPort = newRuleSets(math.MaxUint16+1, len(rules))
-	} else {
-		pi.packet = newRuleSet(len(rules))
-	}
+// newProtoIndex builds an index of the proto `any` rules and proto's own rules, for a protocol with or without
+// ports, or returns nil when there are none. With ProtoAny, it indexes only the proto `any` rules.
+func newProtoIndex(hasPorts bool, rules []rule, proto uint8) *protoIndex {
+	var pi *protoIndex
+	for id := range rules {
+		r := &rules[id]
+		if r.proto != ProtoAny && r.proto != proto {
+			continue
+		}
 
-	for _, id := range ids {
-		pi.add(&rules[id], id)
+		if pi == nil {
+			pi = &protoIndex{hasPorts: hasPorts, fragment: newRuleSet(len(rules))}
+			if hasPorts {
+				pi.byPort = newRuleSets(math.MaxUint16+1, len(rules))
+			} else {
+				pi.packet = newRuleSet(len(rules))
+			}
+		}
+		pi.add(r, id)
 	}
 	return pi
 }

firewall/rules_builder_test.go

@@ -279,11 +279,12 @@ func TestRulesBuilder_manyRules(t *testing.T) {
 	rules := rb.Build(nil, nil)
 
 	assert.Equal(t, ruleSetLen(n+1), rules.In.protos[iputil.IPProtocolTCP].byPort.setLen)
+	// Ids are the order the rules were added in, so the proto `any` rule comes last
 	for i := range n {
 		rs := rules.In.rulesAt(iputil.IPProtocolTCP, int32(i+1))
 		require.Len(t, rs, 2, "port %d", i+1)
-		assert.Equal(t, "any", rs[0].host)
-		assert.Equal(t, fmt.Sprintf("h%d", i), rs[1].host)
+		assert.Equal(t, fmt.Sprintf("h%d", i), rs[0].host)
+		assert.Equal(t, "any", rs[1].host)
 	}
 	assert.Len(t, rules.In.rulesAt(iputil.IPProtocolTCP, n+1), 1)
 }

firewall/table.go

@@ -21,7 +21,8 @@ type Rules struct {
 // proto AND port AND (CA SHA or CA name) AND local CIDR AND (group OR groups OR name OR remote CIDR).
 // A Table does not change once built; see RulesBuilder.
 type Table struct {
-	// rules holds every rule for the direction. A rule's index in this slice is its id in the ruleSets.
+	// rules holds every rule for the direction, in the order they were added. A rule's index in this slice is
+	// its id in the ruleSets.
 	rules []rule
 
 	// protos holds the index for each IP protocol number, so a packet is checked against a single index.
@@ -73,6 +74,10 @@ func (pi *protoIndex) whichMatch(p *Packet, incoming bool) ruleSet {
 
 // rule is a parsed rule. A packet is allowed by the rule when its port, certificate, and address clauses all pass.
 type rule struct {
+	// proto is the IP protocol number the rule applies to, or ProtoAny for every protocol. ICMPv6 rules are kept
+	// as ICMP. The protoIndex checks it, so match does not.
+	proto uint8
+
 	// startPort and endPort are the port clause. A startPort of PortAny covers every port, every packet of a
 	// protocol without ports, and every fragment. PortFragment covers only fragments. A range covers the ports
 	// in it, so it never applies to a protocol without ports.