Fix context return early

kylecarbs · kylecarbs · commit ae36c63fd4a2 · 2022-02-18T15:22:45.000Z
diff --git a/agent/agent.go b/agent/agent.go
@@ -260,9 +260,17 @@ func (s *server) run(ctx context.Context) {
 			s.options.Logger.Warn(context.Background(), "failed to dial", slog.Error(err))
 			continue
 		}
+		defer func() {
+			_ = peerListener.Close()
+		}()
 		s.options.Logger.Debug(context.Background(), "connected")
 		break
 	}
+	select {
+	case <-ctx.Done():
+		return
+	default:
+	}
 
 	for {
 		conn, err := peerListener.Accept()
