From 0cf6ce9aebaa8e1052662a686f928f39c2a53dc1 Mon Sep 17 00:00:00 2001
From: Adam Chalmers <achalmers@cloudflare.com>
Date: Mon, 16 Mar 2020 15:49:04 -0500
Subject: [PATCH] TUN-2812: Tunnel proxies and RPCs can share an edge address

---
 edgediscovery/allregions/regions.go |  4 ++--
 edgediscovery/edgediscovery_test.go | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/edgediscovery/allregions/regions.go b/edgediscovery/allregions/regions.go
index 8a883523..62d30130 100644
--- a/edgediscovery/allregions/regions.go
+++ b/edgediscovery/allregions/regions.go
@@ -68,8 +68,8 @@ func NewNoResolve(addrs []*net.TCPAddr) *Regions {
 
 // GetAnyAddress returns an arbitrary address from the larger region.
 func (rs *Regions) GetAnyAddress() *net.TCPAddr {
-	if rs.region1.AvailableAddrs() > rs.region2.AvailableAddrs() {
-		return rs.region1.GetAnyAddress()
+	if addr := rs.region1.GetAnyAddress(); addr != nil {
+		return addr
 	}
 	return rs.region2.GetAnyAddress()
 }
diff --git a/edgediscovery/edgediscovery_test.go b/edgediscovery/edgediscovery_test.go
index e25dc1bc..0fbc90db 100644
--- a/edgediscovery/edgediscovery_test.go
+++ b/edgediscovery/edgediscovery_test.go
@@ -47,6 +47,24 @@ func TestGiveBack(t *testing.T) {
 	edge.GiveBack(addr)
 	assert.Equal(t, 4, edge.AvailableAddrs())
 }
+
+func TestRPCAndProxyShareSingleEdgeIP(t *testing.T) {
+	l := logrus.New()
+
+	// Make an edge with a single IP
+	edge := MockEdge(l, []*net.TCPAddr{&addr0})
+	tunnelConnID := 0
+
+	// Use the IP for a tunnel
+	addrTunnel, err := edge.GetAddr(tunnelConnID)
+	assert.NoError(t, err)
+
+	// Ensure the IP can be used for RPC too
+	addrRPC, err := edge.GetAddrForRPC()
+	assert.NoError(t, err)
+	assert.Equal(t, addrTunnel, addrRPC)
+}
+
 func TestGetAddrForRPC(t *testing.T) {
 	l := logrus.New()
 	edge := MockEdge(l, []*net.TCPAddr{&addr0, &addr1, &addr2, &addr3})