Aggregate NSG destination addresses optimally

pkg/provider/loadbalancer/accesscontrol.go

@@ -281,8 +281,8 @@ func (ac *AccessControl) CleanSecurityGroup(
 	logger.V(10).Info("Start cleaning")
 
 	var (
-		ipv4Prefixes = fnutil.Map(func(addr netip.Addr) string { return addr.String() }, dstIPv4Addresses)
-		ipv6Prefixes = fnutil.Map(func(addr netip.Addr) string { return addr.String() }, dstIPv6Addresses)
+		ipv4Prefixes = fnutil.Map(fnutil.AsString, dstIPv4Addresses)
+		ipv6Prefixes = fnutil.Map(fnutil.AsString, dstIPv6Addresses)
 	)
 
 	protocols := []armnetwork.SecurityRuleProtocol{

pkg/provider/loadbalancer/accesscontrol_test.go

@@ -1234,14 +1234,14 @@ func TestAccessControl_CleanSecurityGroup(t *testing.T) {
 			{
 				Name: ptr.To("test-rule-2"),
 				Properties: &armnetwork.SecurityRulePropertiesFormat{
-					Protocol:                   to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
-					Access:                     to.Ptr(armnetwork.SecurityRuleAccessAllow),
-					Direction:                  to.Ptr(armnetwork.SecurityRuleDirectionInbound),
-					SourceAddressPrefixes:      to.SliceOfPtrs("*"),
-					SourcePortRange:            ptr.To("*"),
-					DestinationAddressPrefixes: to.SliceOfPtrs("8.8.8.8"),
-					DestinationPortRanges:      to.SliceOfPtrs("5000"),
-					Priority:                   ptr.To(int32(502)),
+					Protocol:                 to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
+					Access:                   to.Ptr(armnetwork.SecurityRuleAccessAllow),
+					Direction:                to.Ptr(armnetwork.SecurityRuleDirectionInbound),
+					SourceAddressPrefixes:    to.SliceOfPtrs("*"),
+					SourcePortRange:          ptr.To("*"),
+					DestinationAddressPrefix: to.Ptr("8.8.8.8"),
+					DestinationPortRanges:    to.SliceOfPtrs("5000"),
+					Priority:                 ptr.To(int32(502)),
 				},
 			},
 		}, outputSG.Properties.SecurityRules)
@@ -1336,14 +1336,14 @@ func TestAccessControl_CleanSecurityGroup(t *testing.T) {
 			{
 				Name: ptr.To("test-rule-2"),
 				Properties: &armnetwork.SecurityRulePropertiesFormat{
-					Protocol:                   to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
-					Access:                     to.Ptr(armnetwork.SecurityRuleAccessAllow),
-					Direction:                  to.Ptr(armnetwork.SecurityRuleDirectionInbound),
-					SourceAddressPrefixes:      to.SliceOfPtrs("*"),
-					SourcePortRange:            ptr.To("*"),
-					DestinationAddressPrefixes: to.SliceOfPtrs("8.8.8.8"),
-					DestinationPortRanges:      to.SliceOfPtrs("5000"),
-					Priority:                   ptr.To(int32(502)),
+					Protocol:                 to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
+					Access:                   to.Ptr(armnetwork.SecurityRuleAccessAllow),
+					Direction:                to.Ptr(armnetwork.SecurityRuleDirectionInbound),
+					SourceAddressPrefixes:    to.SliceOfPtrs("*"),
+					SourcePortRange:          ptr.To("*"),
+					DestinationAddressPrefix: to.Ptr("8.8.8.8"),
+					DestinationPortRanges:    to.SliceOfPtrs("5000"),
+					Priority:                 ptr.To(int32(502)),
 				},
 			},
 		}, outputSG.Properties.SecurityRules)
@@ -1440,14 +1440,14 @@ func TestAccessControl_CleanSecurityGroup(t *testing.T) {
 			{
 				Name: ptr.To("test-rule-2"),
 				Properties: &armnetwork.SecurityRulePropertiesFormat{
-					Protocol:                   to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
-					Access:                     to.Ptr(armnetwork.SecurityRuleAccessAllow),
-					Direction:                  to.Ptr(armnetwork.SecurityRuleDirectionInbound),
-					SourceAddressPrefixes:      to.SliceOfPtrs("*"),
-					SourcePortRange:            ptr.To("*"),
-					DestinationAddressPrefixes: to.SliceOfPtrs("8.8.8.8"),
-					DestinationPortRanges:      to.SliceOfPtrs("5000"),
-					Priority:                   ptr.To(int32(502)),
+					Protocol:                 to.Ptr(armnetwork.SecurityRuleProtocolAsterisk),
+					Access:                   to.Ptr(armnetwork.SecurityRuleAccessAllow),
+					Direction:                to.Ptr(armnetwork.SecurityRuleDirectionInbound),
+					SourceAddressPrefixes:    to.SliceOfPtrs("*"),
+					SourcePortRange:          ptr.To("*"),
+					DestinationAddressPrefix: to.Ptr("8.8.8.8"),
+					DestinationPortRanges:    to.SliceOfPtrs("5000"),
+					Priority:                 ptr.To(int32(502)),
 				},
 			},
 			{

pkg/provider/loadbalancer/fnutil/set.go

@@ -0,0 +1,35 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package fnutil
+
+type Set[T comparable] map[T]struct{}
+
+func SliceToSet[T comparable](xs []T) Set[T] {
+	rv := make(Set[T], len(xs))
+	for _, x := range xs {
+		rv[x] = struct{}{}
+	}
+	return rv
+}
+
+func SetToSlice[T comparable](s Set[T]) []T {
+	rv := make([]T, 0, len(s))
+	for x := range s {
+		rv = append(rv, x)
+	}
+	return rv
+}

pkg/provider/loadbalancer/fnutil/slice.go

@@ -102,6 +102,12 @@ func (xs *IndexSetWithComparableIndex[I, D]) SubtractedBy(ys []D) []D {
 	return rv
 }
 
+// Intersection returns the elements that are in both xs and ys.
 func Intersection[D comparable](xs, ys []D) []D {
 	return IndexSet(xs).Intersection(ys)
 }
+
+// Difference returns the elements in xs but not in ys.
+func Difference[D comparable](xs, ys []D) []D {
+	return IndexSet(ys).SubtractedBy(xs)
+}

pkg/provider/loadbalancer/fnutil/string.go

@@ -0,0 +1,9 @@
+package fnutil
+
+type Stringer interface {
+	String() string
+}
+
+func AsString[T Stringer](v T) string {
+	return v.String()
+}

pkg/provider/loadbalancer/iputil/family.go

@@ -29,6 +29,13 @@ const (
 	IPv6 Family = "IPv6"
 )
 
+func (f Family) MaxMask() int {
+	if f == IPv4 {
+		return 32
+	}
+	return 128
+}
+
 func FamilyOfAddr(addr netip.Addr) Family {
 	if addr.Is4() {
 		return IPv4

pkg/provider/loadbalancer/iputil/internal/bits.go

@@ -0,0 +1,37 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package internal
+
+// setBitAt sets the bit at the i-th position in the byte slice to the given value.
+// Panics if the index is out of bounds.
+// For example,
+// - setBitAt([0x00, 0x00], 8, 1) returns [0x00, 0b1000_0000].
+// - setBitAt([0xff, 0xff], 0, 0) returns [0b0111_1111, 0xff].
+func setBitAt(bytes []byte, i int, bit uint8) {
+	if bit == 1 {
+		bytes[i/8] |= 1 << (7 - i%8)
+	} else {
+		bytes[i/8] &^= 1 << (7 - i%8)
+	}
+}
+
+// bitAt returns the bit at the i-th position in the byte slice.
+// The return value is either 0 or 1 as uint8.
+// Panics if the index is out of bounds.
+func bitAt(bytes []byte, i int) uint8 {
+	return bytes[i/8] >> (7 - i%8) & 1
+}

pkg/provider/loadbalancer/iputil/internal/bits_test.go

@@ -0,0 +1,103 @@
+/*
+Copyright 2024 The Kubernetes Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package internal
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_bitAt(t *testing.T) {
+	bytes := []byte{0b1010_1010, 0b0101_0101}
+	assert.Equal(t, uint8(1), bitAt(bytes, 0))
+	assert.Equal(t, uint8(0), bitAt(bytes, 1))
+	assert.Equal(t, uint8(1), bitAt(bytes, 2))
+	assert.Equal(t, uint8(0), bitAt(bytes, 3))
+
+	assert.Equal(t, uint8(1), bitAt(bytes, 4))
+	assert.Equal(t, uint8(0), bitAt(bytes, 5))
+	assert.Equal(t, uint8(1), bitAt(bytes, 6))
+	assert.Equal(t, uint8(0), bitAt(bytes, 7))
+
+	assert.Equal(t, uint8(0), bitAt(bytes, 8))
+	assert.Equal(t, uint8(1), bitAt(bytes, 9))
+	assert.Equal(t, uint8(0), bitAt(bytes, 10))
+	assert.Equal(t, uint8(1), bitAt(bytes, 11))
+
+	assert.Equal(t, uint8(0), bitAt(bytes, 12))
+	assert.Equal(t, uint8(1), bitAt(bytes, 13))
+	assert.Equal(t, uint8(0), bitAt(bytes, 14))
+	assert.Equal(t, uint8(1), bitAt(bytes, 15))
+
+	assert.Panics(t, func() { bitAt(bytes, 16) })
+}
+
+func Test_setBitAt(t *testing.T) {
+	tests := []struct {
+		name     string
+		initial  []byte
+		index    int
+		bit      uint8
+		expected []byte
+	}{
+		{
+			name:     "Set first bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    0,
+			bit:      1,
+			expected: []byte{0b1000_0000},
+		},
+		{
+			name:     "Set last bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    7,
+			bit:      1,
+			expected: []byte{0b0000_0001},
+		},
+		{
+			name:     "Set middle bit to 1",
+			initial:  []byte{0b0000_0000},
+			index:    4,
+			bit:      1,
+			expected: []byte{0b0000_1000},
+		},
+		{
+			name:     "Set bit to 0",
+			initial:  []byte{0b1111_1111},
+			index:    3,
+			bit:      0,
+			expected: []byte{0b1110_1111},
+		},
+		{
+			name:     "Set bit in second byte",
+			initial:  []byte{0b0000_0000, 0b0000_0000},
+			index:    9,
+			bit:      1,
+			expected: []byte{0b0000_0000, 0b0100_0000},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			setBitAt(tt.initial, tt.index, tt.bit)
+			assert.Equal(t, tt.expected, tt.initial)
+		})
+	}
+
+	assert.Panics(t, func() { setBitAt([]byte{0x00}, 8, 1) })
+}