TUN-4602: Added UDP resolves to Edge discovery

Author: sudarshan-reddy

edgediscovery/allregions/discovery.go

@@ -30,6 +30,12 @@ var (
 	netLookupIP  = net.LookupIP
 )
 
+// EdgeAddr is a representation of possible ways to refer an edge location.
+type EdgeAddr struct {
+	TCP *net.TCPAddr
+	UDP *net.UDPAddr
+}
+
 // If the call to net.LookupSRV fails, try to fall back to DoT from Cloudflare directly.
 //
 // Note: Instead of DoT, we could also have used DoH. Either of these:
@@ -53,7 +59,7 @@ var friendlyDNSErrorLines = []string{
 }
 
 // EdgeDiscovery implements HA service discovery lookup.
-func edgeDiscovery(log *zerolog.Logger) ([][]*net.TCPAddr, error) {
+func edgeDiscovery(log *zerolog.Logger) ([][]*EdgeAddr, error) {
 	_, addrs, err := netLookupSRV(srvService, srvProto, srvName)
 	if err != nil {
 		_, fallbackAddrs, fallbackErr := fallbackLookupSRV(srvService, srvProto, srvName)
@@ -69,16 +75,16 @@ func edgeDiscovery(log *zerolog.Logger) ([][]*net.TCPAddr, error) {
 		addrs = fallbackAddrs
 	}
 
-	var resolvedIPsPerCNAME [][]*net.TCPAddr
+	var resolvedAddrPerCNAME [][]*EdgeAddr
 	for _, addr := range addrs {
-		ips, err := resolveSRVToTCP(addr)
+		edgeAddrs, err := resolveSRV(addr)
 		if err != nil {
 			return nil, err
 		}
-		resolvedIPsPerCNAME = append(resolvedIPsPerCNAME, ips)
+		resolvedAddrPerCNAME = append(resolvedAddrPerCNAME, edgeAddrs)
 	}
 
-	return resolvedIPsPerCNAME, nil
+	return resolvedAddrPerCNAME, nil
 }
 
 func lookupSRVWithDOT(service, proto, name string) (cname string, addrs []*net.SRV, err error) {
@@ -100,31 +106,44 @@ func lookupSRVWithDOT(service, proto, name string) (cname string, addrs []*net.S
 	return r.LookupSRV(ctx, srvService, srvProto, srvName)
 }
 
-func resolveSRVToTCP(srv *net.SRV) ([]*net.TCPAddr, error) {
+func resolveSRV(srv *net.SRV) ([]*EdgeAddr, error) {
 	ips, err := netLookupIP(srv.Target)
 	if err != nil {
 		return nil, errors.Wrapf(err, "Couldn't resolve SRV record %v", srv)
 	}
 	if len(ips) == 0 {
 		return nil, fmt.Errorf("SRV record %v had no IPs", srv)
 	}
-	addrs := make([]*net.TCPAddr, len(ips))
+	addrs := make([]*EdgeAddr, len(ips))
 	for i, ip := range ips {
-		addrs[i] = &net.TCPAddr{IP: ip, Port: int(srv.Port)}
+		addrs[i] = &EdgeAddr{
+			TCP: &net.TCPAddr{IP: ip, Port: int(srv.Port)},
+			UDP: &net.UDPAddr{IP: ip, Port: int(srv.Port)},
+		}
 	}
 	return addrs, nil
 }
 
 // ResolveAddrs resolves TCP address given a list of addresses. Address can be a hostname, however, it will return at most one
 // of the hostname's IP addresses.
-func ResolveAddrs(addrs []string, log *zerolog.Logger) (resolved []*net.TCPAddr) {
+func ResolveAddrs(addrs []string, log *zerolog.Logger) (resolved []*EdgeAddr) {
 	for _, addr := range addrs {
 		tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
 		if err != nil {
-			log.Err(err).Msgf("Failed to resolve %s", addr)
-		} else {
-			resolved = append(resolved, tcpAddr)
+			log.Error().Msgf("Failed to resolve %s to TCP address, err: %v", addr, err)
+			continue
 		}
+
+		udpAddr, err := net.ResolveUDPAddr("udp", addr)
+		if err != nil {
+			log.Error().Msgf("Failed to resolve %s to UDP address, err: %v", addr, err)
+			continue
+		}
+		resolved = append(resolved, &EdgeAddr{
+			TCP: tcpAddr,
+			UDP: udpAddr,
+		})
+
 	}
 	return
 }

edgediscovery/allregions/discovery_test.go

@@ -1,12 +1,17 @@
 package allregions
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/rs/zerolog"
 	"github.com/stretchr/testify/assert"
 )
 
+func (ea *EdgeAddr) String() string {
+	return fmt.Sprintf("%s-%s", ea.TCP, ea.UDP)
+}
+
 func TestEdgeDiscovery(t *testing.T) {
 	mockAddrs := newMockAddrs(19, 2, 5)
 	netLookupSRV = mockNetLookupSRV(mockAddrs)

edgediscovery/allregions/mocks_for_test.go

@@ -11,27 +11,32 @@ import (
 
 type mockAddrs struct {
 	// a set of synthetic SRV records
-	addrMap map[net.SRV][]*net.TCPAddr
+	addrMap map[net.SRV][]*EdgeAddr
 	// the total number of addresses, aggregated across addrMap.
 	// For the convenience of test code that would otherwise have to compute
 	// this by hand every time.
 	numAddrs int
 }
 
 func newMockAddrs(port uint16, numRegions uint8, numAddrsPerRegion uint8) mockAddrs {
-	addrMap := make(map[net.SRV][]*net.TCPAddr)
+	addrMap := make(map[net.SRV][]*EdgeAddr)
 	numAddrs := 0
 
 	for r := uint8(0); r < numRegions; r++ {
 		var (
 			srv   = net.SRV{Target: fmt.Sprintf("test-region-%v.example.com", r), Port: port}
-			addrs []*net.TCPAddr
+			addrs []*EdgeAddr
 		)
 		for a := uint8(0); a < numAddrsPerRegion; a++ {
-			addrs = append(addrs, &net.TCPAddr{
+			tcpAddr := &net.TCPAddr{
 				IP:   net.ParseIP(fmt.Sprintf("10.0.%v.%v", r, a)),
 				Port: int(port),
-			})
+			}
+			udpAddr := &net.UDPAddr{
+				IP:   net.ParseIP(fmt.Sprintf("10.0.%v.%v", r, a)),
+				Port: int(port),
+			}
+			addrs = append(addrs, &EdgeAddr{tcpAddr, udpAddr})
 		}
 		addrMap[srv] = addrs
 		numAddrs += len(addrs)
@@ -74,13 +79,13 @@ func mockNetLookupIP(
 	m mockAddrs,
 ) func(host string) ([]net.IP, error) {
 	return func(host string) ([]net.IP, error) {
-		for srv, tcpAddrs := range m.addrMap {
+		for srv, addrs := range m.addrMap {
 			if srv.Target != host {
 				continue
 			}
-			result := make([]net.IP, len(tcpAddrs))
-			for i, tcpAddr := range tcpAddrs {
-				result[i] = tcpAddr.IP
+			result := make([]net.IP, len(addrs))
+			for i, addr := range addrs {
+				result[i] = addr.TCP.IP
 			}
 			return result, nil
 		}

edgediscovery/allregions/region.go

@@ -1,29 +1,27 @@
 package allregions
 
-import (
-	"net"
-)
-
 // Region contains cloudflared edge addresses. The edge is partitioned into several regions for
 // redundancy purposes.
 type Region struct {
-	connFor map[*net.TCPAddr]UsedBy
+	connFor map[*EdgeAddr]UsedBy
 }
 
 // NewRegion creates a region with the given addresses, which are all unused.
-func NewRegion(addrs []*net.TCPAddr) Region {
+func NewRegion(addrs []*EdgeAddr) Region {
 	// The zero value of UsedBy is Unused(), so we can just initialize the map's values with their
 	// zero values.
-	m := make(map[*net.TCPAddr]UsedBy)
+	connFor := make(map[*EdgeAddr]UsedBy)
 	for _, addr := range addrs {
-		m[addr] = Unused()
+		connFor[addr] = Unused()
+	}
+	return Region{
+		connFor: connFor,
 	}
-	return Region{connFor: m}
 }
 
 // AddrUsedBy finds the address used by the given connection in this region.
 // Returns nil if the connection isn't using any IP.
-func (r *Region) AddrUsedBy(connID int) *net.TCPAddr {
+func (r *Region) AddrUsedBy(connID int) *EdgeAddr {
 	for addr, used := range r.connFor {
 		if used.Used && used.ConnID == connID {
 			return addr
@@ -45,7 +43,7 @@ func (r Region) AvailableAddrs() int {
 
 // GetUnusedIP returns a random unused address in this region.
 // Returns nil if all addresses are in use.
-func (r Region) GetUnusedIP(excluding *net.TCPAddr) *net.TCPAddr {
+func (r Region) GetUnusedIP(excluding *EdgeAddr) *EdgeAddr {
 	for addr, usedby := range r.connFor {
 		if !usedby.Used && addr != excluding {
 			return addr
@@ -55,15 +53,15 @@ func (r Region) GetUnusedIP(excluding *net.TCPAddr) *net.TCPAddr {
 }
 
 // Use the address, assigning it to a proxy connection.
-func (r Region) Use(addr *net.TCPAddr, connID int) {
+func (r Region) Use(addr *EdgeAddr, connID int) {
 	if addr == nil {
 		return
 	}
 	r.connFor[addr] = InUse(connID)
 }
 
 // GetAnyAddress returns an arbitrary address from the region.
-func (r Region) GetAnyAddress() *net.TCPAddr {
+func (r Region) GetAnyAddress() *EdgeAddr {
 	for addr := range r.connFor {
 		return addr
 	}
@@ -72,7 +70,7 @@ func (r Region) GetAnyAddress() *net.TCPAddr {
 
 // GiveBack the address, ensuring it is no longer assigned to an IP.
 // Returns true if the address is in this region.
-func (r Region) GiveBack(addr *net.TCPAddr) (ok bool) {
+func (r Region) GiveBack(addr *EdgeAddr) (ok bool) {
 	if _, ok := r.connFor[addr]; !ok {
 		return false
 	}