fix: resolve goroutine hang/leak on shutdown

fatheradvisor committed Mar 6, 2026 at 01:27 UTC 0a9772617622aaf8b3fc34500e8d12cb173ec165
5 files changed +60 -1
cmd/portal-tunnel/main.go
+4
@@ -243,6 +243,10 @@ func proxyConnection(ctx context.Context, localAddr string, relayConn net.Conn)
243 if err != nil {
244 log.Debug().Err(err).Msg("local->relay copy ended")
245 }
246 + // TLS does not support half-close; full close unblocks the relay→local goroutine.
247 + if closeErr := relayConn.Close(); closeErr != nil {
248 + log.Debug().Err(closeErr).Msg("failed to close relay conn after local->relay copy")
249 + }
250 errCh <- err
251 }()
252
cmd/relay-server/main.go
+4 -1
@@ -176,13 +176,16 @@ func runServer(cfg relayServerConfig) error {
176 if err := serv.Start(); err != nil {
177 return fmt.Errorf("start relay server: %w", err)
178 }
179 - defer serv.Stop()
179
180 apiServ := serveAPI(fmt.Sprintf(":%d", cfg.AdminPort), serv, admin, frontend, cfg, stop)
181
182 <-ctx.Done()
183 log.Info().Msg("[server] shutting down...")
184
185 + // Stop relay first: drain idle reverse conns + close active SNI conns
186 + // so that all HTTP handlers blocked on HandleConnect.Wait() can return.
187 + serv.Stop()
188 +
189 shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
190 defer cancel()
191 if apiServ != nil {
portal/relay.go
+3
@@ -215,6 +215,9 @@ func (g *RelayServer) Stop() {
215 if g.leaseManager != nil {
216 g.leaseManager.Stop()
217 }
218 + if g.reverseHub != nil {
219 + g.reverseHub.Shutdown()
220 + }
221 if g.sniRouter != nil {
222 if err := g.sniRouter.Stop(); err != nil {
223 log.Warn().Err(err).Msg("[RelayServer] Failed to stop SNI router")
portal/reverse_hub.go
+36
@@ -105,7 +105,9 @@ type ReverseHub struct {
105 authorizer func(leaseID, token string) bool
106 ipBanChecker func(ip string) bool
107 onAccepted func(leaseID, ip string)
108 + stopCh chan struct{}
109 mu sync.RWMutex
110 + stopOnce sync.Once
111 }
112
113 // NewReverseHub creates a new reverse connection hub.
@@ -113,6 +115,7 @@ func NewReverseHub() *ReverseHub {
115 return &ReverseHub{
116 pools: make(map[string]chan *ReverseConn),
117 dropped: make(map[string]struct{}),
118 + stopCh: make(chan struct{}),
119 }
120 }
121
@@ -377,12 +380,45 @@ func (h *ReverseHub) closeConn(conn net.Conn, debugCloseMessage string) {
380 }
381 }
382
383 +// Shutdown closes the stop channel and drains all pools, causing idle
384 +// HandleConnect goroutines to unblock and return.
385 +func (h *ReverseHub) Shutdown() {
386 + h.stopOnce.Do(func() {
387 + close(h.stopCh)
388 +
389 + h.mu.Lock()
390 + pools := h.pools
391 + h.pools = make(map[string]chan *ReverseConn)
392 + h.mu.Unlock()
393 +
394 + for _, pool := range pools {
395 + drainPool(pool)
396 + }
397 + })
398 +}
399 +
400 +func drainPool(pool chan *ReverseConn) {
401 + for {
402 + select {
403 + case conn := <-pool:
404 + if conn != nil {
405 + conn.Close()
406 + }
407 + default:
408 + return
409 + }
410 + }
411 +}
412 +
413 func (h *ReverseHub) keepAliveWhileIdle(conn *ReverseConn, leaseID string) {
414 ticker := time.NewTicker(ReverseIdleKeepaliveInterval)
415 defer ticker.Stop()
416
417 for {
418 select {
419 + case <-h.stopCh:
420 + conn.Close()
421 + return
422 case <-conn.done:
423 return
424 case <-conn.active:
portal/sni/router.go
+13
@@ -39,6 +39,7 @@ type Router struct {
39 listener net.Listener
40 routes map[string]*Route
41 leases map[string]*Route
42 + conns map[net.Conn]struct{}
43 onConnection func(conn net.Conn, route *Route)
44 onNoRoute func(conn net.Conn, sni string) bool
45 stopCh chan struct{}
@@ -54,6 +55,7 @@ func NewRouter(addr string) *Router {
55 addr: addr,
56 routes: make(map[string]*Route),
57 leases: make(map[string]*Route),
58 + conns: make(map[net.Conn]struct{}),
59 stopCh: make(chan struct{}),
60 }
61 }
@@ -237,6 +239,9 @@ func (r *Router) Stop() error {
239 if r.listener != nil {
240 closeWithDebugLog(r.listener, "[SNI] failed to close listener")
241 }
242 + for conn := range r.conns {
243 + closeWithDebugLog(conn, "[SNI] failed to close tracked connection")
244 + }
245 r.mu.Unlock()
246 })
247
@@ -334,7 +339,15 @@ func (r *Router) handleConnection(clientConn net.Conn) {
339 onConnection := r.getConnectionHandler()
340
341 if onConnection != nil {
342 + r.mu.Lock()
343 + r.conns[wrappedConn] = struct{}{}
344 + r.mu.Unlock()
345 +
346 onConnection(wrappedConn, route)
347 +
348 + r.mu.Lock()
349 + delete(r.conns, wrappedConn)
350 + r.mu.Unlock()
351 return
352 }
353