Merge pull request #6995 from zarvd/feat/nsg/collapse-cidrs

Consolidate adjacent CIDR ranges for allowed sources
kubernetes-sigs · Sep 11, 2024 · 9b55875 · 9b55875
2 parents d61c51c + f14830e
commit 9b55875
Show file tree

Hide file tree

Showing 3 changed files with 128 additions and 10 deletions.
diff --git a/pkg/provider/loadbalancer/iputil/prefix_test.go b/pkg/provider/loadbalancer/iputil/prefix_test.go
@@ -203,18 +203,57 @@ func TestAggregatePrefixes(t *testing.T) {
  Name: "Overlap IPv4",
  Input: []netip.Prefix{
  netip.MustParsePrefix("192.168.0.0/16"),
- netip.MustParsePrefix("192.169.0.0/16"),
+ netip.MustParsePrefix("192.170.0.0/16"),
  netip.MustParsePrefix("10.10.0.1/32"),
 
  netip.MustParsePrefix("192.168.1.0/24"),
  netip.MustParsePrefix("192.168.1.1/32"),
  },
  Output: []netip.Prefix{
  netip.MustParsePrefix("192.168.0.0/16"),
- netip.MustParsePrefix("192.169.0.0/16"),
+ netip.MustParsePrefix("192.170.0.0/16"),
  netip.MustParsePrefix("10.10.0.1/32"),
  },
  },
+ {
+ Name: "Collapse IPv4",
+ Input: []netip.Prefix{
+ netip.MustParsePrefix("192.168.0.0/24"),
+ netip.MustParsePrefix("192.168.1.0/24"),
+ netip.MustParsePrefix("192.168.2.0/24"),
+ netip.MustParsePrefix("192.168.3.0/24"),
+ netip.MustParsePrefix("10.0.0.0/8"),
+ netip.MustParsePrefix("172.16.0.0/12"),
+ netip.MustParsePrefix("192.168.4.0/24"),
+ netip.MustParsePrefix("192.168.5.0/24"),
+ },
+ Output: []netip.Prefix{
+ netip.MustParsePrefix("10.0.0.0/8"),
+ netip.MustParsePrefix("172.16.0.0/12"),
+ netip.MustParsePrefix("192.168.0.0/22"),
+ netip.MustParsePrefix("192.168.4.0/23"),
+ },
+ },
+ {
+ Name: "Collapse IPv6",
+ Input: []netip.Prefix{
+ netip.MustParsePrefix("2001:db8::/32"),
+ netip.MustParsePrefix("2001:db8:1::/48"),
+ netip.MustParsePrefix("2001:db8:2::/48"),
+ netip.MustParsePrefix("2001:db8:3::/48"),
+ netip.MustParsePrefix("2001:db8:4::/48"),
+ netip.MustParsePrefix("2001:db8:5::/48"),
+ netip.MustParsePrefix("2001:db8:6::/48"),
+ netip.MustParsePrefix("2001:db8:7::/48"),
+ netip.MustParsePrefix("2001:dbf::/32"),
+ netip.MustParsePrefix("2001:dba::/32"),
+ },
+ Output: []netip.Prefix{
+ netip.MustParsePrefix("2001:db8::/32"),
+ netip.MustParsePrefix("2001:dbf::/32"),
+ netip.MustParsePrefix("2001:dba::/32"),
+ },
+ },
  }
 
  for _, tt := range tests {

diff --git a/pkg/provider/loadbalancer/iputil/prefix_tree.go b/pkg/provider/loadbalancer/iputil/prefix_tree.go
@@ -16,14 +16,37 @@ limitations under the License.
 
 package iputil
 
-import "net/netip"
+import (
+ "net/netip"
+)
 
 type prefixTreeNode struct {
  masked bool
  prefix netip.Prefix
 
- l *prefixTreeNode
- r *prefixTreeNode
+ p *prefixTreeNode // parent node
+ l *prefixTreeNode // left child node
+ r *prefixTreeNode // right child node
+}
+
+// pruneToRoot prunes the tree to the root.
+// If a node's left and right children are both masked,
+// it is masked and its children are pruned.
+// This is done recursively up to the root.
+func (n *prefixTreeNode) pruneToRoot() {
+ var node = n
+ for node.p != nil {
+ p := node.p
+ if p.l == nil || !p.l.masked {
+ break
+ }
+ if p.r == nil || !p.r.masked {
+ break
+ }
+ p.masked = true
+ p.l, p.r = nil, nil
+ node = p
+ }
 }
 
 type prefixTree struct {
@@ -70,6 +93,7 @@ func (t *prefixTree) Add(prefix netip.Prefix) {
  }
  n.l = &prefixTreeNode{
  prefix: next,
+ p: n,
  }
  }
  n = n.l
@@ -81,6 +105,7 @@ func (t *prefixTree) Add(prefix netip.Prefix) {
  }
  n.r = &prefixTreeNode{
  prefix: next,
+ p: n,
  }
  }
  n = n.r
@@ -90,10 +115,18 @@ func (t *prefixTree) Add(prefix netip.Prefix) {
  }
 
  n.masked = true
+ n.l, n.r = nil, nil
+ n.pruneToRoot()
 }
 
 // List returns all prefixes in the tree.
 // Overlapping prefixes are merged.
+// It will also collapse the neighboring prefixes.
+// The order of the prefixes in the output is guaranteed.
+//
+// Example:
+// - [192.168.0.0/16, 192.168.1.0/24, 192.168.0.1/32] -> [192.168.0.0/16]
+// - [192.168.0.0/32, 192.168.0.1/32] -> [192.168.0.0/31]
 func (t *prefixTree) List() []netip.Prefix {
  var (
  rv []netip.Prefix

diff --git a/pkg/provider/loadbalancer/iputil/prefix_tree_test.go b/pkg/provider/loadbalancer/iputil/prefix_tree_test.go
@@ -51,18 +51,37 @@ func TestPrefixTreeIPv4(t *testing.T) {
  "Overlap",
  []string{
  "192.168.0.0/16",
- "192.169.0.0/16",
+ "192.170.0.0/16",
  "10.10.0.1/32",
 
  "192.168.1.0/24",
  "192.168.1.1/32",
  },
  []string{
  "192.168.0.0/16",
- "192.169.0.0/16",
+ "192.170.0.0/16",
  "10.10.0.1/32",
  },
  },
+ {
+ "Collapse",
+ []string{
+ "192.168.0.0/24",
+ "192.168.1.0/24",
+ "192.168.2.0/24",
+ "192.168.3.0/24",
+ "10.0.0.0/8",
+ "172.16.0.0/12",
+ "192.168.4.0/24",
+ "192.168.5.0/24",
+ },
+ []string{
+ "10.0.0.0/8",
+ "172.16.0.0/12",
+ "192.168.0.0/22",
+ "192.168.4.0/23",
+ },
+ },
  }
 
  for _, tt := range tests {
@@ -101,27 +120,54 @@ func TestPrefixTreeIPv6(t *testing.T) {
  "NoOverlap",
  []string{
  "2001:db8:0:1::/64",
- "2001:db8:0:2::/64",
  "2001:db8:0:3::/64",
+ "2001:db8:0:5::/64",
  },
  []string{
  "2001:db8:0:1::/64",
- "2001:db8:0:2::/64",
  "2001:db8:0:3::/64",
+ "2001:db8:0:5::/64",
  },
  },
  {
  "Overlap",
  []string{
  "2001:db8::/32",
  "2001:db8:0:1::/64",
- "2001:db8:0:2::/64",
  "2001:db8:0:3::/64",
  },
  []string{
  "2001:db8::/32",
  },
  },
+ {
+ "Collapse",
+ []string{
+ "2001:db8::/32",
+ "2001:db8:1::/48",
+ "2001:db8:2::/48",
+ "2001:db8:3::/48",
+ "2001:db8:4::/48",
+ "2001:db8:5::/48",
+ "2001:db8:6::/48",
+ "2001:db8:7::/48",
+ "2001:db8:8::/48",
+ "2001:db8:9::/48",
+ "2001:db8:a::/48",
+ "2001:db8:b::/48",
+ "2001:db8:c::/48",
+ "2001:db8:d::/48",
+ "2001:db8:e::/48",
+ "2001:db8:f::/48",
+ "2001:dbf::/32", // Noise data
+ "2001:dba::/32", // Noise data
+ },
+ []string{
+ "2001:db8::/32",
+ "2001:dbf::/32",
+ "2001:dba::/32",
+ },
+ },
  }
 
  for _, tt := range tests {