diff --git a/firewall/rules_builder_test.go b/firewall/rules_builder_test.go index 0207f91d..cd9c0249 100644 --- a/firewall/rules_builder_test.go +++ b/firewall/rules_builder_test.go @@ -268,7 +268,7 @@ func TestRulesBuilder_protoIndex(t *testing.T) { assert.Nil(t, rules.Out.protos[iputil.IPProtocolTCP]) } -// TestRulesBuilder_manyRules checks that rule ids beyond the first 64 are indexed correctly. +// TestRulesBuilder_manyRules checks that rule ids beyond the first byte of a rule set are indexed correctly. func TestRulesBuilder_manyRules(t *testing.T) { rb := NewRulesBuilder(test.NewLogger()) const n = 130 @@ -278,7 +278,7 @@ func TestRulesBuilder_manyRules(t *testing.T) { require.NoError(t, rb.AddRule(true, ProtoAny, PortAny, PortAny, nil, "any", "", "", "", "")) rules := rb.Build(nil, nil) - assert.Equal(t, 3, rules.In.protos[iputil.IPProtocolTCP].byPort.setLen) + assert.Equal(t, ruleSetLen(n+1), rules.In.protos[iputil.IPProtocolTCP].byPort.setLen) for i := range n { rs := rules.In.rulesAt(iputil.IPProtocolTCP, int32(i+1)) require.Len(t, rs, 2, "port %d", i+1) diff --git a/firewall/ruleset.go b/firewall/ruleset.go index 9e2c7961..a51e50eb 100644 --- a/firewall/ruleset.go +++ b/firewall/ruleset.go @@ -5,12 +5,14 @@ import ( "math/bits" ) -// ruleSet is a set of rule ids, stored as a bitset. s[i] holds ids i*64 through i*64+63, one per bit. -type ruleSet []uint64 +// ruleSet is a set of rule ids, stored as a bitset. s[i] holds ids i*8 through i*8+7, one per bit. +// Bytes rather than words keep a set to the size the rule count needs, which matters when there is a set for +// every port number. +type ruleSet []uint8 // ruleSetLen returns the length of a ruleSet that can hold ids below n. func ruleSetLen(n int) int { - return (n + 63) / 64 + return (n + 7) / 8 } // newRuleSet returns an empty set that can hold ids below idLimit. @@ -20,21 +22,21 @@ func newRuleSet(idLimit int) ruleSet { // add puts id in s. func (s ruleSet) add(id int) { - s[id/64] |= 1 << (id % 64) + s[id/8] |= 1 << (id % 8) } // ruleSets holds several ruleSet objects of the same length in a single allocation. type ruleSets struct { - // setLen is the amount of uint64s needed to represent a full set + // setLen is the number of bytes needed to represent one set setLen int // bits holds the sets in order. Set i starts at i*setLen. - bits []uint64 + bits []uint8 } // newRuleSets returns n empty sets, each of which can hold ids below idLimit. func newRuleSets(n, idLimit int) ruleSets { setLen := ruleSetLen(idLimit) - return ruleSets{setLen: setLen, bits: make([]uint64, n*setLen)} + return ruleSets{setLen: setLen, bits: make([]uint8, n*setLen)} } // at returns set i. @@ -50,7 +52,7 @@ func (s ruleSets) at(i int) ruleSet { func (s ruleSet) any(pred func(id int) bool) bool { for i, ids := range s { for ids != 0 { - id := i*64 + bits.TrailingZeros64(ids) + id := i*8 + bits.TrailingZeros8(ids) ids &= ids - 1 // Clear the lowest set bit, which is id's. if pred(id) { return true @@ -65,7 +67,7 @@ func (s ruleSet) all() iter.Seq[int] { return func(yield func(int) bool) { for i, ids := range s { for ids != 0 { - id := i*64 + bits.TrailingZeros64(ids) + id := i*8 + bits.TrailingZeros8(ids) ids &= ids - 1 // Clear the lowest set bit, which is id's. if !yield(id) { return diff --git a/firewall_test.go b/firewall_test.go index 41cd81e7..ab129759 100644 --- a/firewall_test.go +++ b/firewall_test.go @@ -2,6 +2,7 @@ package nebula import ( "bytes" + "fmt" "log/slog" "net/netip" "testing" @@ -419,6 +420,40 @@ func BenchmarkFirewallTable_match(b *testing.B) { var benchMatchSink bool +// BenchmarkFirewallTable_matchManyRules measures a table with more rules than fit in one word of a rule set. +// Rule 0 is a proto `any` rule that fails on host, and rules 1 through 130 are tcp ports 1 through 130 for hosts +// h1 through h130, so the last rule id is the one that passes. +func BenchmarkFirewallTable_matchManyRules(b *testing.B) { + rb := firewall.NewRulesBuilder(test.NewLogger()) + require.NoError(b, rb.AddRule(true, firewall.ProtoAny, firewall.PortAny, firewall.PortAny, nil, "nope", "", "", "", "")) + for i := 1; i <= 130; i++ { + require.NoError(b, rb.AddRule(true, iputil.IPProtocolTCP, int32(i), int32(i), nil, fmt.Sprintf("h%d", i), "", "", "", "")) + } + ft := rb.Build(nil, nil).In + cp := cert.NewCAPool() + c := &cert.CachedCertificate{Certificate: &dummyCert{name: "h130"}, InvertedGroups: map[string]struct{}{}} + + for _, tc := range []struct { + name string + port uint16 + want bool + }{ + {"fail on port", 5000, false}, + {"pass on name, last id", 130, true}, + {"fail on name, middle id", 65, false}, + } { + p := &firewall.Packet{Protocol: iputil.IPProtocolTCP, LocalPort: tc.port} + b.Run(tc.name, func(b *testing.B) { + if ft.Match(p, true, c, cp) != tc.want { + b.Fatal("wrong verdict") + } + for b.Loop() { + benchMatchSink = ft.Match(p, true, c, cp) + } + }) + } +} + func TestFirewall_Drop2(t *testing.T) { ob := &bytes.Buffer{} l := test.NewLoggerWithOutput(ob)