From b9fa7ab3e6f5aa9f3abb2a033490f0719e6fc581 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Thu, 27 Jan 2022 09:40:11 -0600
Subject: [PATCH] Fix goroutine leak after close

It was possible for the Close() function to exit without all goroutines
existing. By closing these channels before the after function, all
goroutines are forced to exit before.
---
 agent.go | 18 ++++++++----------
 1 file changed, 8 insertions(+), 10 deletions(-)

diff --git a/agent.go b/agent.go
index 7d449f3..b4aba19 100644
--- a/agent.go
+++ b/agent.go
@@ -109,8 +109,9 @@ type Agent struct {
 	extIPMapper *externalIPMapper
 
 	// State for closing
-	done chan struct{}
-	err  atomicError
+	done         chan struct{}
+	taskLoopDone chan struct{}
+	err          atomicError
 
 	gatherCandidateCancel func()
 
@@ -213,6 +214,7 @@ func (a *Agent) taskLoop() {
 		close(a.chanState)
 		close(a.chanCandidate)
 		close(a.chanCandidatePair)
+		close(a.taskLoopDone)
 	}()
 
 	for {
@@ -289,6 +291,7 @@ func NewAgent(config *AgentConfig) (*Agent, error) { //nolint:gocognit
 		onConnected:       make(chan struct{}),
 		buffer:            packetio.NewBuffer(),
 		done:              make(chan struct{}),
+		taskLoopDone:      make(chan struct{}),
 		startedCh:         startedCtx.Done(),
 		startedFn:         startedFn,
 		portmin:           config.PortMin,
@@ -421,12 +424,12 @@ func (a *Agent) startOnConnectionStateChangeRoutine() {
 					}
 					return
 				}
-				a.onConnectionStateChange(s)
+				go a.onConnectionStateChange(s)
 
 			case c, isOpen := <-a.chanCandidate:
 				if !isOpen {
 					for s := range a.chanState {
-						a.onConnectionStateChange(s)
+						go a.onConnectionStateChange(s)
 					}
 					return
 				}
@@ -897,20 +900,15 @@ func (a *Agent) Close() error {
 		return err
 	}
 
-	done := make(chan struct{})
-
 	a.afterRun(func(context.Context) {
 		a.gatherCandidateCancel()
-		close(done)
 	})
-
 	a.err.Store(ErrClosed)
 
 	a.removeUfragFromMux()
 
 	close(a.done)
-
-	<-done
+	<-a.taskLoopDone
 	return nil
 }