From 5970c4671ccfe1cd6fe91254113c877767df006f Mon Sep 17 00:00:00 2001
From: Tim Hockin <thockin@google.com>
Date: Tue, 17 Nov 2020 20:41:55 -0800
Subject: [PATCH] Add an IPFamily() method to ipallocator

---
 .../core/service/ipallocator/allocator.go      | 18 +++++++++++++++---
 .../core/service/ipallocator/allocator_test.go |  7 +++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/pkg/registry/core/service/ipallocator/allocator.go b/pkg/registry/core/service/ipallocator/allocator.go
index 634569f767a..e42b3401ca4 100644
--- a/pkg/registry/core/service/ipallocator/allocator.go
+++ b/pkg/registry/core/service/ipallocator/allocator.go
@@ -35,6 +35,7 @@ type Interface interface {
 	Release(net.IP) error
 	ForEach(func(net.IP))
 	CIDR() net.IPNet
+	IPFamily() api.IPFamily
 
 	// For testing
 	Has(ip net.IP) bool
@@ -76,6 +77,8 @@ type Range struct {
 	base *big.Int
 	// max is the maximum size of the usable addresses in the range
 	max int
+	// family is the IP family of this range
+	family api.IPFamily
 
 	alloc allocator.Interface
 }
@@ -85,13 +88,16 @@ func NewAllocatorCIDRRange(cidr *net.IPNet, allocatorFactory allocator.Allocator
 	max := utilnet.RangeSize(cidr)
 	base := utilnet.BigForIP(cidr.IP)
 	rangeSpec := cidr.String()
+	var family api.IPFamily
 
 	if utilnet.IsIPv6CIDR(cidr) {
+		family = api.IPv6Protocol
 		// Limit the max size, since the allocator keeps a bitmap of that size.
 		if max > 65536 {
 			max = 65536
 		}
 	} else {
+		family = api.IPv4Protocol
 		// Don't use the IPv4 network's broadcast address.
 		max--
 	}
@@ -101,9 +107,10 @@ func NewAllocatorCIDRRange(cidr *net.IPNet, allocatorFactory allocator.Allocator
 	max--
 
 	r := Range{
-		net:  cidr,
-		base: base,
-		max:  maximum(0, int(max)),
+		net:    cidr,
+		base:   base,
+		max:    maximum(0, int(max)),
+		family: family,
 	}
 	var err error
 	r.alloc, err = allocatorFactory(r.max, rangeSpec)
@@ -219,6 +226,11 @@ func (r *Range) Has(ip net.IP) bool {
 	return r.alloc.Has(offset)
 }
 
+// IPFamily returns the IP family of this range.
+func (r *Range) IPFamily() api.IPFamily {
+	return r.family
+}
+
 // Snapshot saves the current state of the pool.
 func (r *Range) Snapshot(dst *api.RangeAllocation) error {
 	snapshottable, ok := r.alloc.(allocator.Snapshottable)
diff --git a/pkg/registry/core/service/ipallocator/allocator_test.go b/pkg/registry/core/service/ipallocator/allocator_test.go
index 63ab5725695..db43f65ef5f 100644
--- a/pkg/registry/core/service/ipallocator/allocator_test.go
+++ b/pkg/registry/core/service/ipallocator/allocator_test.go
@@ -28,6 +28,7 @@ func TestAllocate(t *testing.T) {
 	testCases := []struct {
 		name             string
 		cidr             string
+		family           api.IPFamily
 		free             int
 		released         string
 		outOfRange       []string
@@ -36,6 +37,7 @@ func TestAllocate(t *testing.T) {
 		{
 			name:     "IPv4",
 			cidr:     "192.168.1.0/24",
+			family:   api.IPv4Protocol,
 			free:     254,
 			released: "192.168.1.5",
 			outOfRange: []string{
@@ -49,6 +51,7 @@ func TestAllocate(t *testing.T) {
 		{
 			name:     "IPv6",
 			cidr:     "2001:db8:1::/48",
+			family:   api.IPv6Protocol,
 			free:     65535,
 			released: "2001:db8:1::5",
 			outOfRange: []string{
@@ -79,6 +82,10 @@ func TestAllocate(t *testing.T) {
 			t.Errorf("allocator returned a different cidr")
 		}
 
+		if r.IPFamily() != tc.family {
+			t.Errorf("allocator returned wrong IP family")
+		}
+
 		if f := r.Used(); f != 0 {
 			t.Errorf("Test %s unexpected used %d", tc.name, f)
 		}