sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-09-17 04:09:07 +02:00
Code Issues Packages Projects Releases Wiki Activity

[management] cleanup resources when ws-grpc proxy connection goes away (#7484)

Browse Source 
* ws to grpc connection adapter

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* support for timeouts on reading h2 stream headers

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* cleanups

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* we can't always expect a DATA frame, as not all http methods send it

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* set default headers read timeout to 10s

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fix a race in tests

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* remove frame interceptor

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* cleanup test cleanup

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* make linter happy

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* removed unused consts

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* set 5s ReadTimeout

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* making linter happy

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* making linter happy

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* updated comments

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* fix spelling

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* disabled all http server read timeouts

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* Revert "disabled all http server read timeouts"

This reverts commit adf5005ba4.

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

* clarify comment re: ReadTimeout/WriteTimeout issues

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

---------

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>
This commit is contained in:
dmitri-netbird
2026-09-10 16:37:43 +02:00
committed by GitHub
parent 9615d2ab16
commit 0fac1ee638
3 changed files with 373 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files
util/wsproxy/server/proxy.go
+43 -120
View File
@@ -1,12 +1,8 @@
package server
import (
"context"
"io"
"net"
"net/http"
"sync"
"time"
"sync/atomic"
"github.com/coder/websocket"
log "github.com/sirupsen/logrus"
@@ -15,11 +11,6 @@ import (
"github.com/netbirdio/netbird/util/wsproxy"
)
const (
bufferSize = 32 * 1024
ioTimeout = 5 * time.Second
)
// Config contains the configuration for the WebSocket proxy.
type Config struct {
Handler http.Handler
@@ -53,14 +44,23 @@ func New(handler http.Handler, opts ...Option) *Proxy {
// Handler returns an http.Handler that proxies WebSocket connections to the local gRPC server.
func (p *Proxy) Handler() http.Handler {
return http.HandlerFunc(p.handleWebSocket)
return &proxyHandler{
metrics: p.config.MetricsRecorder,
handler: p.config.Handler,
}
}
func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) {
type proxyHandler struct {
metrics MetricsRecorder
handler http.Handler
conn atomic.Pointer[wsConnAdapter]
}
func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
p.metrics.RecordConnection(ctx)
defer p.metrics.RecordDisconnection(ctx)
ph.metrics.RecordConnection(ctx)
defer ph.metrics.RecordDisconnection(ctx)
log.Debugf("WebSocket proxy handling connection from %s, forwarding to internal gRPC handler", r.RemoteAddr)
acceptOptions := &websocket.AcceptOptions{
@@ -69,121 +69,44 @@ func (p *Proxy) handleWebSocket(w http.ResponseWriter, r *http.Request) {
wsConn, err := websocket.Accept(w, r, acceptOptions)
if err != nil {
p.metrics.RecordError(ctx, "websocket_accept_failed")
ph.metrics.RecordError(ctx, "websocket_accept_failed")
log.Errorf("WebSocket upgrade failed from %s: %v", r.RemoteAddr, err)
return
}
defer func() {
_ = wsConn.Close(websocket.StatusNormalClosure, "")
}()
serverConn := (&wsConnAdapter{
ctx: ctx,
conn: wsConn,
metrics: ph.metrics,
clientAddr: r.RemoteAddr,
})
clientConn, serverConn := net.Pipe()
defer func() {
_ = clientConn.Close()
_ = serverConn.Close()
}()
ph.conn.Store(serverConn) // used in tests only
log.Debugf("WebSocket proxy established: %s -> gRPC handler", r.RemoteAddr)
go func() {
(&http2.Server{}).ServeConn(serverConn, &http2.ServeConnOpts{
Context: ctx,
Handler: p.config.Handler,
})
}()
(&http2.Server{
// TODO (dmitri) we should limit the number of concurrent streams per connection (peer)
// and idle timeouts
// MaxConcurrentStreams: 20,
// IdleTimeout: 10 * time.Second,
}).ServeConn(serverConn, &http2.ServeConnOpts{
Context: ctx,
Handler: ph.handler,
BaseConfig: &http.Server{
// b/c we are wrapping a ws connection, read and write connection deadlines normally set
// via ReadTimeout and WriteTimeout http.Server fields aren't available to us. The ws
// library doesn't expose connection deadline timer config, and we ignore these calls in "wsConnAdapter".
//
// Another issue is that Server.ServeConn() call bypasses setting of connection deadlines altogether,
// ReadTimeout and Writetimeout set here would only apply to h2 streams, i.e. after a HEADERS frame
// arrival and processing, turning ReadTimeout into a request body read deadline, and WriteTimeout into
// a response deadline (the latter not useful for streaming requests).
},
})
p.proxyData(ctx, wsConn, clientConn, r.RemoteAddr)
}
func (p *Proxy) proxyData(ctx context.Context, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
proxyCtx, cancel := context.WithCancel(ctx)
defer cancel()
var wg sync.WaitGroup
wg.Add(2)
go p.wsToPipe(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr)
go p.pipeToWS(proxyCtx, cancel, &wg, wsConn, pipeConn, clientAddr)
wg.Wait()
}
func (p *Proxy) wsToPipe(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
defer wg.Done()
defer cancel()
for {
msgType, data, err := wsConn.Read(ctx)
if err != nil {
switch {
case ctx.Err() != nil:
log.Debugf("WebSocket from %s terminating due to context cancellation", clientAddr)
case websocket.CloseStatus(err) != -1:
log.Debugf("WebSocket from %s disconnected", clientAddr)
default:
p.metrics.RecordError(ctx, "websocket_read_error")
log.Debugf("WebSocket read error from %s: %v", clientAddr, err)
}
return
}
if msgType != websocket.MessageBinary {
log.Warnf("Unexpected WebSocket message type from %s: %v", clientAddr, msgType)
continue
}
if ctx.Err() != nil {
log.Tracef("wsToPipe goroutine terminating due to context cancellation before pipe write")
return
}
if err := pipeConn.SetWriteDeadline(time.Now().Add(ioTimeout)); err != nil {
log.Debugf("Failed to set pipe write deadline: %v", err)
}
n, err := pipeConn.Write(data)
if err != nil {
p.metrics.RecordError(ctx, "pipe_write_error")
log.Warnf("Pipe write error for %s: %v", clientAddr, err)
return
}
p.metrics.RecordBytesTransferred(ctx, "ws_to_grpc", int64(n))
}
}
func (p *Proxy) pipeToWS(ctx context.Context, cancel context.CancelFunc, wg *sync.WaitGroup, wsConn *websocket.Conn, pipeConn net.Conn, clientAddr string) {
defer wg.Done()
defer cancel()
buf := make([]byte, bufferSize)
for {
n, err := pipeConn.Read(buf)
if err != nil {
if ctx.Err() != nil {
log.Tracef("pipeToWS goroutine terminating due to context cancellation")
return
}
if err != io.EOF {
log.Debugf("Pipe read error for %s: %v", clientAddr, err)
}
return
}
if ctx.Err() != nil {
log.Tracef("pipeToWS goroutine terminating due to context cancellation before WebSocket write")
return
}
if n > 0 {
if err := wsConn.Write(ctx, websocket.MessageBinary, buf[:n]); err != nil {
p.metrics.RecordError(ctx, "websocket_write_error")
log.Warnf("WebSocket write error for %s: %v", clientAddr, err)
return
}
p.metrics.RecordBytesTransferred(ctx, "grpc_to_ws", int64(n))
}
}
log.Debugf("WebSocket proxy closing: %s -> gRPC handler", r.RemoteAddr)
}