From 427ac0fddba12aa8845c140df74df3c98d66d421 Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Mon, 21 Feb 2022 09:49:25 -0600
Subject: [PATCH] Fix goroutine leak for closing while gathering

The "gatherCandidates()" function could previously leak after
close, because the WaitGroup was left unchecked.
---
 agent.go  | 4 ++++
 gather.go | 3 +++
 2 files changed, 7 insertions(+)

diff --git a/agent.go b/agent.go
index aa2a1d6..951103b 100644
--- a/agent.go
+++ b/agent.go
@@ -114,6 +114,7 @@ type Agent struct {
 	err          atomicError
 
 	gatherCandidateCancel func()
+	gatherCandidateDone   chan struct{}
 
 	chanCandidate     chan Candidate
 	chanCandidatePair chan *CandidatePair
@@ -907,6 +908,9 @@ func (a *Agent) Close() error {
 
 	a.afterRun(func(context.Context) {
 		a.gatherCandidateCancel()
+		if a.gatherCandidateDone != nil {
+			<-a.gatherCandidateDone
+		}
 	})
 	a.err.Store(ErrClosed)
 
diff --git a/gather.go b/gather.go
index a1dc437..bbe416d 100644
--- a/gather.go
+++ b/gather.go
@@ -71,6 +71,7 @@ func (a *Agent) GatherCandidates() error {
 		a.gatherCandidateCancel() // Cancel previous gathering routine
 		ctx, cancel := context.WithCancel(ctx)
 		a.gatherCandidateCancel = cancel
+		a.gatherCandidateDone = make(chan struct{})
 
 		go a.gatherCandidates(ctx)
 	}); runErr != nil {
@@ -80,6 +81,7 @@ func (a *Agent) GatherCandidates() error {
 }
 
 func (a *Agent) gatherCandidates(ctx context.Context) {
+	defer close(a.gatherCandidateDone)
 	if err := a.setGatheringState(GatheringStateGathering); err != nil {
 		a.log.Warnf("failed to set gatheringState to GatheringStateGathering: %v", err)
 		return
@@ -120,6 +122,7 @@ func (a *Agent) gatherCandidates(ctx context.Context) {
 		case CandidateTypePeerReflexive, CandidateTypeUnspecified:
 		}
 	}
+
 	// Block until all STUN and TURN URLs have been gathered (or timed out)
 	wg.Wait()