From 6f37280e8e122092abe40a05861fe33039f4531d Mon Sep 17 00:00:00 2001
From: Nathan Brown <nbrown.us@gmail.com>
Date: Mon, 26 Apr 2021 10:42:24 -0500
Subject: [PATCH] Fully close tunnels when CloseAllTunnels is called (#448)

---
 control.go |  4 +++-
 outside.go | 11 ++++++++---
 ssh.go     |  2 +-
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/control.go b/control.go
index c00a958..39e1979 100644
--- a/control.go
+++ b/control.go
@@ -142,7 +142,7 @@ func (c *Control) CloseTunnel(vpnIP uint32, localOnly bool) bool {
 		)
 	}
 
-	c.f.closeTunnel(hostInfo)
+	c.f.closeTunnel(hostInfo, false)
 	return true
 }
 
@@ -160,6 +160,8 @@ func (c *Control) CloseAllTunnels(excludeLighthouses bool) (closed int) {
 
 		if h.ConnectionState.ready {
 			c.f.send(closeTunnel, 0, h.ConnectionState, h, h.remote, []byte{}, make([]byte, 12, 12), make([]byte, mtu))
+			c.f.closeTunnel(h, true)
+
 			c.l.WithField("vpnIp", IntIp(h.hostId)).WithField("udpAddr", h.remote).
 				Debug("Sending close tunnel message")
 			closed++
diff --git a/outside.go b/outside.go
index b2fd6e2..3fc2170 100644
--- a/outside.go
+++ b/outside.go
@@ -118,7 +118,7 @@ func (f *Interface) readOutsidePackets(addr *udpAddr, out []byte, packet []byte,
 		hostinfo.logger(f.l).WithField("udpAddr", addr).
 			Info("Close tunnel received, tearing down.")
 
-		f.closeTunnel(hostinfo)
+		f.closeTunnel(hostinfo, false)
 		return
 
 	default:
@@ -133,12 +133,17 @@ func (f *Interface) readOutsidePackets(addr *udpAddr, out []byte, packet []byte,
 }
 
 // closeTunnel closes a tunnel locally, it does not send a closeTunnel packet to the remote
-func (f *Interface) closeTunnel(hostInfo *HostInfo) {
+func (f *Interface) closeTunnel(hostInfo *HostInfo, hasHostMapLock bool) {
 	//TODO: this would be better as a single function in ConnectionManager that handled locks appropriately
 	f.connectionManager.ClearIP(hostInfo.hostId)
 	f.connectionManager.ClearPendingDeletion(hostInfo.hostId)
 	f.lightHouse.DeleteVpnIP(hostInfo.hostId)
-	f.hostMap.DeleteHostInfo(hostInfo)
+
+	if hasHostMapLock {
+		f.hostMap.unlockedDeleteHostInfo(hostInfo)
+	} else {
+		f.hostMap.DeleteHostInfo(hostInfo)
+	}
 }
 
 // sendCloseTunnel is a helper function to send a proper close tunnel packet to a remote
diff --git a/ssh.go b/ssh.go
index 3714671..4cfaf1d 100644
--- a/ssh.go
+++ b/ssh.go
@@ -520,7 +520,7 @@ func sshCloseTunnel(ifce *Interface, fs interface{}, a []string, w sshd.StringWr
 		)
 	}
 
-	ifce.closeTunnel(hostInfo)
+	ifce.closeTunnel(hostInfo, false)
 	return w.WriteLine("Closed")
 }