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 {