From 1bbdb40185e908176e6f8d91d29f550199287d2b Mon Sep 17 00:00:00 2001
From: Kyle Carberry <kyle@carberry.com>
Date: Thu, 14 Apr 2022 09:39:40 -0500
Subject: [PATCH] Fix close delay collecting STUN candidates

This was causing a 5s delay when closing the agent.
---
 gather.go | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/gather.go b/gather.go
index 7cf9783..36fed81 100644
--- a/gather.go
+++ b/gather.go
@@ -402,7 +402,7 @@ func (a *Agent) gatherCandidatesSrflxUDPMux(ctx context.Context, urls []*URL, ne
 	}
 }
 
-func (a *Agent) gatherCandidatesSrflx(ctx context.Context, urls []*URL, networkTypes []NetworkType) {
+func (a *Agent) gatherCandidatesSrflx(ctx context.Context, urls []*URL, networkTypes []NetworkType) { //nolint:gocognit
 	var wg sync.WaitGroup
 	defer wg.Wait()
 
@@ -428,6 +428,18 @@ func (a *Agent) gatherCandidatesSrflx(ctx context.Context, urls []*URL, networkT
 					closeConnAndLog(conn, a.log, fmt.Sprintf("Failed to listen for %s: %v\n", serverAddr.String(), err))
 					return
 				}
+				// If the agent closes midway through the connection
+				// we end it early to prevent close delay.
+				cancelCtx, cancelFunc := context.WithCancel(ctx)
+				defer cancelFunc()
+				go func() {
+					select {
+					case <-cancelCtx.Done():
+						return
+					case <-a.done:
+						_ = conn.Close()
+					}
+				}()
 
 				xoraddr, err := getXORMappedAddr(conn, serverAddr, stunGatherTimeout)
 				if err != nil {