From a6ad3c46e2916a633a7f4a50740ce690880ca063 Mon Sep 17 00:00:00 2001
From: David Crawshaw <crawshaw@tailscale.com>
Date: Fri, 21 Feb 2020 22:20:31 -0500
Subject: [PATCH] magicsock: spray some normal packets after a handshake

In particular, this is designed to catch the case where a
HandshakeInitiation packet is sent out but the intermediate NATs
have not been primed, so the packet passes over DERP.
In that case, the HandshakeResponse also comes back over DERP,
and the connection proceeds via DERP without ever trying to punch
through the NAT.

With this change, the HandshakeResponse (which was sprayed out
and so primed one NAT) triggers an UpdateDst, which triggers
the extra spray logic.

(For this to work, there has to be an initial supply of packets
to send on to a peer for the three seconds following a handshake.
The source of these packets is left as a future exercise.)

Signed-off-by: David Crawshaw <crawshaw@tailscale.com>
---
 wgengine/magicsock/magicsock.go | 36 +++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/wgengine/magicsock/magicsock.go b/wgengine/magicsock/magicsock.go
index 387148ca8..6258943d5 100644
--- a/wgengine/magicsock/magicsock.go
+++ b/wgengine/magicsock/magicsock.go
@@ -387,11 +387,37 @@ var logPacketDests, _ = strconv.ParseBool(os.Getenv("DEBUG_LOG_PACKET_DESTS"))
 //
 // It also returns as's current roamAddr, if any.
 func appendDests(dsts []*net.UDPAddr, as *AddrSet, b []byte) (_ []*net.UDPAddr, roamAddr *net.UDPAddr) {
-	spray := shouldSprayPacket(b)
+	spray := shouldSprayPacket(b) // true for handshakes
+	now := time.Now()
 
 	as.mu.Lock()
 	defer as.mu.Unlock()
 
+	// Spray logic.
+	//
+	// After exchanging a handshake with a peer, we send some outbound
+	// packets to every endpoint of that peer. These packets are spaced out
+	// over several seconds to make sure that our peer has an opportunity to
+	// send its own spray packet to us before we are done spraying.
+	//
+	// Multiple packets are necessary because we have to both establish the
+	// NAT mappings between two peers *and use* the mappings to switch away
+	// from DERP to a higher-priority UDP endpoint.
+	const sprayPeriod = 3 * time.Second
+	const sprayFreq = 250 * time.Millisecond
+	if spray {
+		as.lastSpray = now
+		as.stopSpray = now.Add(sprayPeriod)
+	} else if now.Before(as.stopSpray) {
+		// We are in the spray window. If it has been sprayFreq since we
+		// last sprayed a packet, spray this packet.
+		if now.Sub(as.lastSpray) >= sprayFreq {
+			spray = true
+			as.lastSpray = now
+		}
+	}
+
+	// Pick our destination address(es).
 	roamAddr = as.roamAddr
 	if roamAddr != nil {
 		dsts = append(dsts, roamAddr)
@@ -775,7 +801,7 @@ type AddrSet struct {
 	publicKey key.Public    // peer public key used for DERP communication
 	addrs     []net.UDPAddr // ordered priority list (low to high) provided by wgengine
 
-	mu sync.Mutex // guards roamAddr and curAddr
+	mu sync.Mutex // guards following fields
 
 	// roamAddr is non-nil if/when we receive a correctly signed
 	// WireGuard packet from an unexpected address. If so, we
@@ -789,6 +815,12 @@ type AddrSet struct {
 	// address a valid packet has been received from so far.
 	// If no valid packet from addrs has been received, curAddr is -1.
 	curAddr int
+
+	// stopSpray is the time after which we stop spraying packets.
+	stopSpray time.Time
+
+	// lastSpray is the lsat time we sprayed a packet.
+	lastSpray time.Time
 }
 
 var noAddr = &net.UDPAddr{