mirror of
https://github.com/netbirdio/netbird.git
synced 2026-04-18 00:06:38 +00:00
1bd7190954
* Fix WebSocket support by implementing Hijacker interface Add responsewriter.PassthroughWriter to preserve optional HTTP interfaces (Hijacker, Flusher, Pusher) when wrapping http.ResponseWriter in middleware. Without this delegation: - WebSocket connections fail (can't hijack the connection) - Streaming breaks (can't flush buffers) - HTTP/2 push doesn't work * Add HijackTracker to manage hijacked connections during graceful shutdown * Refactor HijackTracker to use middleware for tracking hijacked connections * Refactor server handler chain setup for improved readability and maintainability
50 lines
1.1 KiB
Go
50 lines
1.1 KiB
Go
package conntrack
|
|
|
|
import (
|
|
"bufio"
|
|
"net"
|
|
"net/http"
|
|
)
|
|
|
|
// trackedConn wraps a net.Conn and removes itself from the tracker on Close.
|
|
type trackedConn struct {
|
|
net.Conn
|
|
tracker *HijackTracker
|
|
}
|
|
|
|
func (c *trackedConn) Close() error {
|
|
c.tracker.conns.Delete(c)
|
|
return c.Conn.Close()
|
|
}
|
|
|
|
// trackingWriter wraps an http.ResponseWriter and intercepts Hijack calls
|
|
// to replace the raw connection with a trackedConn that auto-deregisters.
|
|
type trackingWriter struct {
|
|
http.ResponseWriter
|
|
tracker *HijackTracker
|
|
}
|
|
|
|
func (w *trackingWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
|
|
hijacker, ok := w.ResponseWriter.(http.Hijacker)
|
|
if !ok {
|
|
return nil, nil, http.ErrNotSupported
|
|
}
|
|
conn, buf, err := hijacker.Hijack()
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
tc := &trackedConn{Conn: conn, tracker: w.tracker}
|
|
w.tracker.conns.Store(tc, struct{}{})
|
|
return tc, buf, nil
|
|
}
|
|
|
|
func (w *trackingWriter) Flush() {
|
|
if flusher, ok := w.ResponseWriter.(http.Flusher); ok {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
|
|
func (w *trackingWriter) Unwrap() http.ResponseWriter {
|
|
return w.ResponseWriter
|
|
}
|