Fold proto any firewall rules into each protocol's rules so a packet checks one table

firewall.go

@@ -89,31 +89,88 @@ type FirewallConntrack struct {
 // FirewallTable is the entry point for a rule, the evaluation order is:
 // Proto AND port AND (CA SHA or CA name) AND local CIDR AND (group OR groups OR name OR remote CIDR)
 type FirewallTable struct {
-	// Protos holds the rules for each IP protocol number, nil for protocols without rules.
+	// Protos holds the rules for each IP protocol number, so a packet is only checked against one set of rules.
+	// A protocol with rules of its own also holds a copy of every proto `any` rule, the rest share AnyProto.
 	// ICMP and ICMPv6 share one set of rules.
-	Protos   [256]firewallPort
+	Protos [256]firewallPort
+
+	// AnyProto holds the proto `any` rules, nil until there is one
 	AnyProto firewallPort
+
+	// own marks the protocols in Protos with rules of their own
+	own [256]bool
+
+	// anyRules are the proto `any` rules, kept to copy into protocols that get rules of their own later
+	anyRules []firewallPortRule
+}
+
+// firewallPortRule holds the arguments to firewallPort.addRule
+type firewallPortRule struct {
+	startPort, endPort                   int32
+	groups                               []string
+	host, cidr, localCidr, caName, caSha string
+}
+
+func (r firewallPortRule) addTo(f *Firewall, fp firewallPort) error {
+	return fp.addRule(f, r.startPort, r.endPort, r.groups, r.host, r.cidr, r.localCidr, r.caName, r.caSha)
 }
 
 func newFirewallTable() *FirewallTable {
-	return &FirewallTable{
-		AnyProto: firewallPort{},
-	}
+	return &FirewallTable{}
 }
 
-// protoPorts returns the rules for proto, creating them on first use.
-func (ft *FirewallTable) protoPorts(proto uint8) firewallPort {
-	if ft.Protos[proto] == nil {
+// addRule adds r to the rules for proto, or to the rules for every protocol when proto is firewall.ProtoAny
+func (ft *FirewallTable) addRule(f *Firewall, proto uint8, r firewallPortRule) error {
+	if proto == firewall.ProtoAny {
+		return ft.addAnyProtoRule(f, r)
+	}
+
+	if !ft.own[proto] {
 		fp := firewallPort{}
+		for _, ar := range ft.anyRules {
+			if err := ar.addTo(f, fp); err != nil {
+				return err
+			}
+		}
+
 		ft.Protos[proto] = fp
+		ft.own[proto] = true
 		switch proto {
 		case iputil.IPProtocolICMP:
 			ft.Protos[iputil.IPProtocolICMPv6] = fp
+			ft.own[iputil.IPProtocolICMPv6] = true
 		case iputil.IPProtocolICMPv6:
 			ft.Protos[iputil.IPProtocolICMP] = fp
+			ft.own[iputil.IPProtocolICMP] = true
+		}
+	}
+
+	return r.addTo(f, ft.Protos[proto])
+}
+
+func (ft *FirewallTable) addAnyProtoRule(f *Firewall, r firewallPortRule) error {
+	if ft.AnyProto == nil {
+		ft.AnyProto = firewallPort{}
+	}
+	if err := r.addTo(f, ft.AnyProto); err != nil {
+		return err
+	}
+	ft.anyRules = append(ft.anyRules, r)
+
+	for proto := range ft.Protos {
+		switch {
+		case !ft.own[proto]:
+			ft.Protos[proto] = ft.AnyProto
+		case proto == iputil.IPProtocolICMPv6:
+			// Shares its rules with ICMP, which gets the copy
+		default:
+			if err := r.addTo(f, ft.Protos[proto]); err != nil {
+				return err
+			}
 		}
 	}
-	return ft.Protos[proto]
+
+	return nil
 }
 
 type FirewallCA struct {
@@ -263,30 +320,20 @@ func NewFirewallFromConfig(l *slog.Logger, cs *CertState, c *config.C) (*Firewal
 
 // AddRule properly creates the in memory rule structure for a firewall table.
 func (f *Firewall) AddRule(incoming bool, proto uint8, startPort int32, endPort int32, groups []string, host string, cidr, localCidr, caName string, caSha string) error {
-	var (
-		ft *FirewallTable
-		fp firewallPort
-	)
-
+	var ft *FirewallTable
 	if incoming {
 		ft = f.InRules
 	} else {
 		ft = f.OutRules
 	}
 
-	switch proto {
-	case iputil.IPProtocolICMP, iputil.IPProtocolICMPv6:
+	if proto == iputil.IPProtocolICMP || proto == iputil.IPProtocolICMPv6 {
 		//ICMP traffic doesn't have ports, so we always coerce to "any", even if a value is provided
 		if startPort != firewall.PortAny {
 			f.l.Warn("ignoring port specification for ICMP firewall rule", "startPort", startPort)
 		}
 		startPort = firewall.PortAny
 		endPort = firewall.PortAny
-		fp = ft.protoPorts(proto)
-	case firewall.ProtoAny:
-		fp = ft.AnyProto
-	default:
-		fp = ft.protoPorts(proto)
 	}
 
 	// We need this rule string because we generate a hash. Removing this will break firewall reload.
@@ -304,7 +351,16 @@ func (f *Firewall) AddRule(incoming bool, proto uint8, startPort int32, endPort
 		"firewallRule", m{"direction": direction, "proto": proto, "startPort": startPort, "endPort": endPort, "groups": groups, "host": host, "cidr": cidr, "localCidr": localCidr, "caName": caName, "caSha": caSha},
 	)
 
-	return fp.addRule(f, startPort, endPort, groups, host, cidr, localCidr, caName, caSha)
+	return ft.addRule(f, proto, firewallPortRule{
+		startPort: startPort,
+		endPort:   endPort,
+		groups:    groups,
+		host:      host,
+		cidr:      cidr,
+		localCidr: localCidr,
+		caName:    caName,
+		caSha:     caSha,
+	})
 }
 
 // GetRuleHash returns a hash representation of all inbound and outbound rules
@@ -652,10 +708,6 @@ func (f *Firewall) evict(p firewall.Packet) {
 }
 
 func (ft *FirewallTable) match(p *firewall.Packet, incoming bool, c *cert.CachedCertificate, caPool *cert.CAPool) bool {
-	if ft.AnyProto.match(p, incoming, c, caPool) {
-		return true
-	}
-
 	return ft.Protos[p.Protocol].match(p, incoming, c, caPool)
 }
 
@@ -686,7 +738,7 @@ func (fp firewallPort) match(p *firewall.Packet, incoming bool, c *cert.CachedCe
 		return false
 	}
 
-	// this branch is here to catch ICMP traffic from FirewallTable.AnyProto.match and the ICMP rules
+	// ICMP has no ports, only the port `any` rules apply, including the ones copied in from proto `any` rules
 	if p.Protocol == iputil.IPProtocolICMP || p.Protocol == iputil.IPProtocolICMPv6 {
 		// port numbers are re-used for connection tracking of ICMP,
 		// but we don't want to actually filter on them.

firewall_test.go

@@ -348,15 +348,14 @@ func TestFirewall_DropV6(t *testing.T) {
 func BenchmarkFirewallTable_match(b *testing.B) {
 	f := &Firewall{}
 	ft := FirewallTable{}
-	tcp := ft.protoPorts(iputil.IPProtocolTCP)
 
 	pfix := netip.MustParsePrefix("172.1.1.1/32")
-	_ = tcp.addRule(f, 10, 10, []string{"good-group"}, "good-host", pfix.String(), "", "", "")
-	_ = tcp.addRule(f, 100, 100, []string{"good-group"}, "good-host", "", pfix.String(), "", "")
+	_ = ft.addRule(f, iputil.IPProtocolTCP, firewallPortRule{startPort: 10, endPort: 10, groups: []string{"good-group"}, host: "good-host", cidr: pfix.String()})
+	_ = ft.addRule(f, iputil.IPProtocolTCP, firewallPortRule{startPort: 100, endPort: 100, groups: []string{"good-group"}, host: "good-host", localCidr: pfix.String()})
 
 	pfix6 := netip.MustParsePrefix("fd11::11/128")
-	_ = tcp.addRule(f, 10, 10, []string{"good-group"}, "good-host", pfix6.String(), "", "", "")
-	_ = tcp.addRule(f, 100, 100, []string{"good-group"}, "good-host", "", pfix6.String(), "", "")
+	_ = ft.addRule(f, iputil.IPProtocolTCP, firewallPortRule{startPort: 10, endPort: 10, groups: []string{"good-group"}, host: "good-host", cidr: pfix6.String()})
+	_ = ft.addRule(f, iputil.IPProtocolTCP, firewallPortRule{startPort: 100, endPort: 100, groups: []string{"good-group"}, host: "good-host", localCidr: pfix6.String()})
 	cp := cert.NewCAPool()
 
 	b.Run("fail on proto", func(b *testing.B) {
@@ -958,6 +957,64 @@ func TestFirewall_ProtocolRules(t *testing.T) {
 		require.NoError(t, fw.Drop(pkt(iputil.IPProtocolICMPv6), true, &h, cp, nil), "rule %d", ruleProto)
 		assert.Equal(t, ErrNoMatchingRule, fw.Drop(pkt(47), true, &h, cp, nil), "rule %d", ruleProto)
 	}
+
+	// proto `any` rules apply to protocols with rules of their own, whichever rule was added first
+	for _, anyFirst := range []bool{true, false} {
+		fw := NewFirewall(l, time.Second, time.Minute, time.Hour, c.Certificate)
+		addAny := func() {
+			require.NoError(t, fw.AddRule(true, firewall.ProtoAny, 80, 80, []string{"any"}, "", "", "", "", ""))
+		}
+		if anyFirst {
+			addAny()
+		}
+		require.NoError(t, fw.AddRule(true, iputil.IPProtocolTCP, 22, 22, []string{"any"}, "", "", "", "", ""))
+		if !anyFirst {
+			addAny()
+		}
+
+		for _, tc := range []struct {
+			proto uint8
+			port  uint16
+			want  error
+		}{
+			{iputil.IPProtocolTCP, 22, nil},
+			{iputil.IPProtocolTCP, 80, nil},
+			{iputil.IPProtocolTCP, 23, ErrNoMatchingRule},
+			{iputil.IPProtocolUDP, 80, nil},
+			{iputil.IPProtocolUDP, 22, ErrNoMatchingRule},
+		} {
+			p := pkt(tc.proto)
+			p.LocalPort = tc.port
+			assert.Equal(t, tc.want, fw.Drop(p, true, &h, cp, nil), "anyFirst %v, proto %d, port %d", anyFirst, tc.proto, tc.port)
+		}
+	}
+
+	// ICMP and ICMPv6 share one copy of each proto `any` rule, whichever rule was added first
+	for _, anyFirst := range []bool{true, false} {
+		fw := NewFirewall(l, time.Second, time.Minute, time.Hour, c.Certificate)
+		addAny := func() {
+			require.NoError(t, fw.AddRule(true, firewall.ProtoAny, firewall.PortAny, firewall.PortAny, []string{"default-group"}, "", "", "", "", ""))
+		}
+		if anyFirst {
+			addAny()
+		}
+		require.NoError(t, fw.AddRule(true, iputil.IPProtocolICMPv6, firewall.PortAny, firewall.PortAny, []string{"nope"}, "", "", "", "", ""))
+		if !anyFirst {
+			addAny()
+		}
+
+		require.NoError(t, fw.Drop(pkt(iputil.IPProtocolICMP), true, &h, cp, nil), "anyFirst %v", anyFirst)
+		require.NoError(t, fw.Drop(pkt(iputil.IPProtocolICMPv6), true, &h, cp, nil), "anyFirst %v", anyFirst)
+		assert.Len(t, fw.InRules.Protos[iputil.IPProtocolICMP][firewall.PortAny].Any.Groups, 2, "anyFirst %v", anyFirst)
+	}
+
+	// Protocols without rules of their own have no rules to check until there is a proto `any` rule to share
+	fw = NewFirewall(l, time.Second, time.Minute, time.Hour, c.Certificate)
+	require.NoError(t, fw.AddRule(true, iputil.IPProtocolTCP, 22, 22, []string{"any"}, "", "", "", "", ""))
+	assert.Nil(t, fw.InRules.Protos[47])
+	assert.Equal(t, ErrNoMatchingRule, fw.Drop(pkt(47), true, &h, cp, nil))
+	require.NoError(t, fw.AddRule(true, firewall.ProtoAny, firewall.PortAny, firewall.PortAny, []string{"any"}, "", "", "", "", ""))
+	require.NoError(t, fw.Drop(pkt(47), true, &h, cp, nil))
 }
 
 func TestFirewall_DropIPSpoofing(t *testing.T) {