From ae5505bc74a77fc08701f658adb19eef552cbea0 Mon Sep 17 00:00:00 2001
From: Wade Simmons <wadey@slack-corp.com>
Date: Tue, 19 Oct 2021 10:53:55 -0400
Subject: [PATCH] handshake: update to preferred remote (#532)

If we receive a handshake packet for a tunnel that has already been
completed, check to see if the new remote is preferred. If so, update to
the preferred remote and send a test packet to influence the other side
to do the same.
---
 handshake_ix.go | 21 ++++++++++++++++++++-
 hostmap.go      | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/handshake_ix.go b/handshake_ix.go
index 770e19b..7fcdfb7 100644
--- a/handshake_ix.go
+++ b/handshake_ix.go
@@ -200,6 +200,20 @@ func ixHandshakeStage1(f *Interface, addr *udpAddr, packet []byte, h *Header) {
 	if err != nil {
 		switch err {
 		case ErrAlreadySeen:
+			// Update remote if preferred (Note we have to switch to locking
+			// the existing hostinfo, and then switch back so the defer Unlock
+			// higher in this function still works)
+			hostinfo.Unlock()
+			existing.Lock()
+			// Update remote if preferred
+			if existing.SetRemoteIfPreferred(f.hostMap, addr) {
+				// Send a test packet to ensure the other side has also switched to
+				// the preferred remote
+				f.SendMessageToVpnIp(test, testRequest, vpnIP, []byte(""), make([]byte, 12, 12), make([]byte, mtu))
+			}
+			existing.Unlock()
+			hostinfo.Lock()
+
 			msg = existing.HandshakePacket[2]
 			f.messageMetrics.Tx(handshake, NebulaMessageSubType(msg[1]), 1)
 			err := f.outside.WriteTo(msg, addr)
@@ -305,7 +319,12 @@ func ixHandshakeStage2(f *Interface, addr *udpAddr, hostinfo *HostInfo, packet [
 			WithField("handshake", m{"stage": 2, "style": "ix_psk0"}).WithField("header", h).
 			Info("Handshake is already complete")
 
-		//TODO: evaluate addr for preference, if we handshook with a less preferred addr we can correct quickly here
+		// Update remote if preferred
+		if hostinfo.SetRemoteIfPreferred(f.hostMap, addr) {
+			// Send a test packet to ensure the other side has also switched to
+			// the preferred remote
+			f.SendMessageToVpnIp(test, testRequest, hostinfo.hostId, []byte(""), make([]byte, 12, 12), make([]byte, mtu))
+		}
 
 		// We already have a complete tunnel, there is nothing that can be done by processing further stage 1 packets
 		return false
diff --git a/hostmap.go b/hostmap.go
index 281b6e2..25982c0 100644
--- a/hostmap.go
+++ b/hostmap.go
@@ -507,6 +507,42 @@ func (i *HostInfo) SetRemote(remote *udpAddr) {
 	}
 }
 
+// SetRemoteIfPreferred returns true if the remote was changed. The lastRoam
+// time on the HostInfo will also be updated.
+func (i *HostInfo) SetRemoteIfPreferred(hm *HostMap, newRemote *udpAddr) bool {
+	currentRemote := i.remote
+	if currentRemote == nil {
+		i.SetRemote(newRemote)
+		return true
+	}
+
+	// NOTE: We do this loop here instead of calling `isPreferred` in
+	// remote_list.go so that we only have to loop over preferredRanges once.
+	newIsPreferred := false
+	for _, l := range hm.preferredRanges {
+		// return early if we are already on a preferred remote
+		if l.Contains(currentRemote.IP) {
+			return false
+		}
+
+		if l.Contains(newRemote.IP) {
+			newIsPreferred = true
+		}
+	}
+
+	if newIsPreferred {
+		// Consider this a roaming event
+		i.lastRoam = time.Now()
+		i.lastRoamRemote = currentRemote.Copy()
+
+		i.SetRemote(newRemote)
+
+		return true
+	}
+
+	return false
+}
+
 func (i *HostInfo) ClearConnectionState() {
 	i.ConnectionState = nil
 }