fix: close hijacks racing with shutdown
This commit is contained in:
@@ -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{})
|
||||
|
||||
Reference in New Issue
Block a user