From f8abacf2e5b14a6e1d52adf8798a2f68ba62587b Mon Sep 17 00:00:00 2001
From: MiMoCode <mimocode@localhost>
Date: Fri, 10 Jul 2026 19:43:41 +0800
Subject: [PATCH] fix: close hijacks racing with shutdown

---
 proxy/proxy.go                |  9 ++++
 tests/proxy/websocket_test.go | 80 +++++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+)

diff --git a/proxy/proxy.go b/proxy/proxy.go
index 1730924..40f1b67 100644
--- a/proxy/proxy.go
+++ b/proxy/proxy.go
@@ -33,6 +33,8 @@ type Handler struct {
 	connMu       sync.Mutex
 	conns        map[net.Conn]struct{}
 	connChanged  chan struct{}
+	closing      bool
+	closed       bool
 }
 
 // Options 控制反代连接上游时的网络行为。
@@ -253,7 +255,9 @@ func (h *Handler) serveRequestBodyError(cw *captureWriter, r *http.Request, st *
 func (h *Handler) Shutdown(ctx context.Context) error {
 	for {
 		h.connMu.Lock()
+		h.closing = true
 		if len(h.conns) == 0 {
+			h.closed = true
 			h.connMu.Unlock()
 			return nil
 		}
@@ -284,6 +288,11 @@ func (h *Handler) trackConn(conn net.Conn) net.Conn {
 		h.connMu.Unlock()
 	}
 	h.connMu.Lock()
+	if h.closing || h.closed {
+		h.connMu.Unlock()
+		_ = conn.Close()
+		return conn
+	}
 	h.conns[tracked] = struct{}{}
 	close(h.connChanged)
 	h.connChanged = make(chan struct{})
diff --git a/tests/proxy/websocket_test.go b/tests/proxy/websocket_test.go
index caf42a1..5bcf54b 100644
--- a/tests/proxy/websocket_test.go
+++ b/tests/proxy/websocket_test.go
@@ -11,6 +11,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -45,6 +46,30 @@ type failingWriteConn struct{ net.Conn }
 
 func (failingWriteConn) Write([]byte) (int, error) { return 0, errors.New("handshake write failed") }
 
+type closeNotifyConn struct {
+	net.Conn
+	closed chan struct{}
+	once   sync.Once
+}
+
+func (c *closeNotifyConn) Close() error {
+	err := c.Conn.Close()
+	c.once.Do(func() { close(c.closed) })
+	return err
+}
+
+type blockingHijackResponseWriter struct {
+	*hijackResponseWriter
+	hijacked chan struct{}
+	release  chan struct{}
+}
+
+func (w *blockingHijackResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	close(w.hijacked)
+	<-w.release
+	return w.hijackResponseWriter.Hijack()
+}
+
 // fakeUpstream 模拟一个最简 WebSocket 升级：
 // 收到 GET + Upgrade: websocket 后回 101，然后做字节回声直到对端关闭。
 func fakeUpstream(t *testing.T) *httptest.Server {
@@ -161,6 +186,61 @@ func TestWebSocketHandshakeCapturedAndShutdownClosesConnection(t *testing.T) {
 	}
 }
 
+func TestShutdownClosesConnectionHijackedBeforeRegistration(t *testing.T) {
+	upstream := fakeUpstream(t)
+	defer upstream.Close()
+	u, _ := url.Parse(upstream.URL)
+	filter, err := config.NewFilter(config.FilterDisabled, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	h := proxy.New(u, filter, noopSubmitter{}, 1024)
+	downstream, peer := net.Pipe()
+	defer peer.Close()
+	closed := make(chan struct{})
+	conn := &closeNotifyConn{Conn: downstream, closed: closed}
+	w := &blockingHijackResponseWriter{
+		hijackResponseWriter: &hijackResponseWriter{header: make(http.Header), conn: conn},
+		hijacked:             make(chan struct{}),
+		release:              make(chan struct{}),
+	}
+	req := httptest.NewRequest(http.MethodGet, "http://proxy.test/v1/realtime", nil)
+	req.Header.Set("Upgrade", "websocket")
+	req.Header.Set("Connection", "Upgrade")
+	serveDone := make(chan struct{})
+	var releaseOnce sync.Once
+	releaseHijack := func() { releaseOnce.Do(func() { close(w.release) }) }
+	go func() {
+		h.ServeHTTP(w, req)
+		close(serveDone)
+	}()
+	defer func() {
+		releaseHijack()
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+		defer cancel()
+		_ = h.Shutdown(ctx)
+		<-serveDone
+	}()
+
+	select {
+	case <-w.hijacked:
+	case <-time.After(time.Second):
+		t.Fatal("downstream connection was not hijacked")
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	defer cancel()
+	if err := h.Shutdown(ctx); err != nil {
+		t.Fatal(err)
+	}
+	releaseHijack()
+
+	select {
+	case <-closed:
+	case <-time.After(100 * time.Millisecond):
+		t.Fatal("connection hijacked during Shutdown remained open")
+	}
+}
+
 func TestWebSocketHandshakeWriteFailureIsNotCaptured(t *testing.T) {
 	upstream := fakeUpstream(t)
 	defer upstream.Close()