From 8a374a35a0fe62dfd86df2c16166d2bb84115b93 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Tue, 2 Feb 2021 10:46:34 -0800
Subject: [PATCH] device: tie encryption queue lifetime to the peers that write
 to it

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
---
 device/device.go | 6 ++++--
 device/peer.go   | 2 ++
 device/send.go   | 2 --
 3 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/device/device.go b/device/device.go
index 5f36036..d860ed3 100644
--- a/device/device.go
+++ b/device/device.go
@@ -397,6 +397,10 @@ func (device *Device) Close() {
 
 	device.isUp.Set(false)
 
+	// Remove peers before closing queues,
+	// because peers assume that queues are active.
+	device.RemoveAllPeers()
+
 	// We kept a reference to the encryption and decryption queues,
 	// in case we started any new peers that might write to them.
 	// No new peers are coming; we are done with these queues.
@@ -405,8 +409,6 @@ func (device *Device) Close() {
 	device.queue.handshake.wg.Done()
 	device.state.stopping.Wait()
 
-	device.RemoveAllPeers()
-
 	device.rate.limiter.Close()
 
 	device.state.changing.Set(false)
diff --git a/device/peer.go b/device/peer.go
index 76f9a96..65f73a9 100644
--- a/device/peer.go
+++ b/device/peer.go
@@ -177,6 +177,7 @@ func (peer *Peer) Start() {
 	if peer.queue.staged == nil {
 		peer.queue.staged = make(chan *QueueOutboundElement, QueueStagedSize)
 	}
+	peer.device.queue.encryption.wg.Add(1) // keep encryption queue open for our writes
 
 	peer.timersInit()
 	peer.handshake.lastSentHandshake = time.Now().Add(-(RekeyTimeout + time.Second))
@@ -248,6 +249,7 @@ func (peer *Peer) Stop() {
 	close(peer.queue.inbound)
 	close(peer.queue.outbound)
 	peer.stopping.Wait()
+	peer.device.queue.encryption.wg.Done() // no more writes to encryption queue from us
 
 	peer.ZeroAndFlushAll()
 }
diff --git a/device/send.go b/device/send.go
index 9d63c4e..8def4ce 100644
--- a/device/send.go
+++ b/device/send.go
@@ -291,8 +291,6 @@ func (peer *Peer) StagePacket(elem *QueueOutboundElement) {
 }
 
 func (peer *Peer) SendStagedPackets() {
-	peer.device.queue.encryption.wg.Add(1)
-	defer peer.device.queue.encryption.wg.Done()
 top:
 	if len(peer.queue.staged) == 0 || !peer.device.isUp.Get() {
 		return