diff --git a/.github/instructions/ARCHITECTURE.instructions.md b/.github/instructions/ARCHITECTURE.instructions.md index fc85e3238..e6f66cd70 100644 --- a/.github/instructions/ARCHITECTURE.instructions.md +++ b/.github/instructions/ARCHITECTURE.instructions.md @@ -127,7 +127,7 @@ graph TB | **Database** | SQLite | 3.x | Embedded database | | **ORM** | GORM | Latest | Database abstraction layer | | **Reverse Proxy** | Caddy Server | 2.11.2 | Embedded HTTP/HTTPS proxy | -| **WebSocket** | gorilla/websocket | Latest | Real-time log streaming | +| **WebSocket** | coder/websocket | Latest | Real-time log streaming | | **Crypto** | golang.org/x/crypto | Latest | Password hashing, encryption | | **Metrics** | Prometheus Client | Latest | Application metrics | | **Notifications** | Notify | Latest | Multi-platform alerts | diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 84049ec18..53d100a7c 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -128,7 +128,7 @@ graph TB | **Database** | SQLite | 3.x | Embedded database | | **ORM** | GORM | Latest | Database abstraction layer | | **Reverse Proxy** | Caddy Server | 2.11.2 | Embedded HTTP/HTTPS proxy | -| **WebSocket** | gorilla/websocket | Latest | Real-time log streaming | +| **WebSocket** | coder/websocket | Latest | Real-time log streaming | | **Crypto** | golang.org/x/crypto | Latest | Password hashing, encryption | | **Metrics** | Prometheus Client | Latest | Application metrics | | **Notifications** | github.com/Wikid82/Hermes | Current | External delivery-engine module (Discord, Slack, Gotify, Pushover, Ntfy, Telegram, generic webhook, Web Push, and email) consumed via Charon-supplied SSRF/SMTP/template adapters — see Service Layer below | diff --git a/agent/go.mod b/agent/go.mod index aed56d190..7e25ebfb0 100644 --- a/agent/go.mod +++ b/agent/go.mod @@ -3,7 +3,7 @@ module github.com/Wikid82/charon/agent go 1.27.4 require ( - github.com/gorilla/websocket v1.5.3 + github.com/coder/websocket v1.8.15 github.com/hashicorp/yamux v0.1.2 github.com/sirupsen/logrus v1.10.2 github.com/stretchr/testify v1.12.1 diff --git a/agent/go.sum b/agent/go.sum index 5badf0cc2..d44f2ba1a 100644 --- a/agent/go.sum +++ b/agent/go.sum @@ -1,5 +1,5 @@ -github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= -github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= +github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA= +github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/hashicorp/yamux v0.1.2 h1:XtB8kyFOyHXYVFnwT5C3+Bdo8gArse7j2AQ0DA0Uey8= github.com/hashicorp/yamux v0.1.2/go.mod h1:C+zze2n6e/7wshOZep2A70/aQU6QBRWJO/G6FT1wIns= github.com/sirupsen/logrus v1.10.2 h1:G2SED73/qrAu6YwbdxOD6peLkCBI3z7L+ykJFTXJBBo= diff --git a/agent/leash/leash.go b/agent/leash/leash.go index b8e5f89cc..c18d3107b 100644 --- a/agent/leash/leash.go +++ b/agent/leash/leash.go @@ -14,7 +14,7 @@ import ( "strings" "time" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/hashicorp/yamux" "github.com/sirupsen/logrus" @@ -117,24 +117,27 @@ func (l *Leash) Run(ctx context.Context) error { // connect performs a single WebSocket dial and yamux session. It blocks until // the session is closed or ctx is cancelled. func (l *Leash) connect(ctx context.Context) error { - dialer := websocket.Dialer{ - HandshakeTimeout: wsDialTimeout, - } - l.log.WithField("server_url", l.serverURL).Info("leash: connecting to server") - wsConn, resp, err := dialer.DialContext(ctx, l.serverURL, http.Header{ - "Authorization": {"Bearer " + l.authKey}, - "X-Orthrus-Version": {"1.0.0"}, - "X-Orthrus-Name": {l.agentID}, + // The dial context bounds only the handshake; the established connection + // outlives it. + dialCtx, cancelDial := context.WithTimeout(ctx, wsDialTimeout) + defer cancelDial() + + wsConn, resp, err := websocket.Dial(dialCtx, l.serverURL, &websocket.DialOptions{ //nolint:bodyclose // coder/websocket closes the handshake body itself; resp.Body is nil on success + HTTPClient: newDialHTTPClient(), + HTTPHeader: http.Header{ + "Authorization": {"Bearer " + l.authKey}, + "X-Orthrus-Version": {"1.0.0"}, + "X-Orthrus-Name": {l.agentID}, + }, + // Compression has never been negotiated by this agent; keep it off. + CompressionMode: websocket.CompressionDisabled, }) if err != nil { return fmt.Errorf("leash: websocket dial: %w", err) } - writeEnabled := resp != nil && resp.Header.Get("X-Orthrus-Write-Enabled") == "true" - if resp != nil && resp.Body != nil { - _ = resp.Body.Close() - } + writeEnabled := resp.Header.Get("X-Orthrus-Write-Enabled") == "true" // X-Orthrus-Write-Enabled is delivered atomically in the 101 Switching // Protocols response that completes the handshake (server.go sets it via @@ -145,9 +148,10 @@ func (l *Leash) connect(ctx context.Context) error { cfg := yamux.DefaultConfig() cfg.LogOutput = io.Discard - session, err := yamux.Client(NewWSNetConn(wsConn), cfg) + netConn := NewWSNetConn(wsConn) + session, err := yamux.Client(netConn, cfg) if err != nil { - _ = wsConn.Close() + _ = netConn.Close() return fmt.Errorf("leash: yamux client: %w", err) } defer func() { _ = session.Close() }() diff --git a/agent/leash/leash_test.go b/agent/leash/leash_test.go index a096a500b..1eae74953 100644 --- a/agent/leash/leash_test.go +++ b/agent/leash/leash_test.go @@ -7,10 +7,11 @@ import ( "net/http" "net/http/httptest" "path/filepath" + "sync/atomic" "testing" "time" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/hashicorp/yamux" "github.com/sirupsen/logrus" "github.com/stretchr/testify/assert" @@ -47,20 +48,21 @@ func newYamuxServerSession(t *testing.T, wsConn *websocket.Conn) *yamux.Session return session } -var testUpgrader = websocket.Upgrader{ - CheckOrigin: func(r *http.Request) bool { return true }, +// acceptWS upgrades a test request to a WebSocket connection, accepting any origin. +func acceptWS(w http.ResponseWriter, r *http.Request) (*websocket.Conn, error) { + return websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true}) } // TestLeash_Reconnect verifies that the Leash attempts to reconnect after the server // closes the connection immediately. func TestLeash_Reconnect(t *testing.T) { - connectCount := 0 + var connectCount atomic.Int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - connectCount++ - conn, err := testUpgrader.Upgrade(w, r, nil) + connectCount.Add(1) + conn, err := acceptWS(w, r) require.NoError(t, err) - _ = conn.Close() + _ = conn.CloseNow() })) defer srv.Close() @@ -83,28 +85,26 @@ func TestLeash_Reconnect(t *testing.T) { _ = l.Run(ctx) - assert.GreaterOrEqual(t, connectCount, 2, "expected at least 2 connection attempts within the timeout") + assert.GreaterOrEqual(t, connectCount.Load(), int32(2), "expected at least 2 connection attempts within the timeout") } // TestWsNetConn_ReadWrite verifies that the net.Conn wrapper correctly transfers // binary data through a WebSocket echo server. func TestWsNetConn_ReadWrite(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := testUpgrader.Upgrade(w, r, nil) + conn, err := acceptWS(w, r) require.NoError(t, err) - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() - _, msg, err := conn.ReadMessage() + _, msg, err := conn.Read(context.Background()) require.NoError(t, err) - require.NoError(t, conn.WriteMessage(websocket.BinaryMessage, msg)) + require.NoError(t, conn.Write(context.Background(), websocket.MessageBinary, msg)) })) defer srv.Close() - dialer := websocket.Dialer{} - wsConn, resp, err := dialer.Dial("ws"+srv.URL[4:], nil) + wsConn, _, err := websocket.Dial(context.Background(), "ws"+srv.URL[4:], nil) //nolint:bodyclose // coder/websocket closes the handshake body itself; resp.Body is nil on success require.NoError(t, err) - defer func() { _ = resp.Body.Close() }() - defer func() { _ = wsConn.Close() }() + defer func() { _ = wsConn.CloseNow() }() netConn := leash.NewWSNetConn(wsConn) @@ -119,6 +119,36 @@ func TestWsNetConn_ReadWrite(t *testing.T) { assert.Equal(t, sent, buf) } +// TestWsNetConn_CloseIsAbrupt verifies that closing the net.Conn tears the +// socket down without a WebSocket close frame, matching the agent's previous +// raw-socket close. The peer must observe an error that carries no close status. +func TestWsNetConn_CloseIsAbrupt(t *testing.T) { + peerErr := make(chan error, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := acceptWS(w, r) + require.NoError(t, err) + defer func() { _ = conn.CloseNow() }() + + _, _, readErr := conn.Read(context.Background()) + peerErr <- readErr + })) + defer srv.Close() + + wsConn, _, err := websocket.Dial(context.Background(), "ws"+srv.URL[4:], nil) //nolint:bodyclose // coder/websocket closes the handshake body itself; resp.Body is nil on success + require.NoError(t, err) + + netConn := leash.NewWSNetConn(wsConn) + require.NoError(t, netConn.Close()) + + select { + case readErr := <-peerErr: + require.Error(t, readErr) + assert.Equal(t, websocket.StatusCode(-1), websocket.CloseStatus(readErr), "peer must not receive a close frame") + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for the peer to observe the close") + } +} + // TestLeash_Connect_DockerStreamDispatchesThroughFilter proves the full // write-mode wiring chain actually executes for a real accepted stream: // connect()'s AcceptStream loop -> handleStream's streamTypeDocker dispatch @@ -160,9 +190,9 @@ func TestLeash_Connect_DockerStreamDispatchesThroughFilter(t *testing.T) { defer close(serverDone) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - wsConn, upgradeErr := testUpgrader.Upgrade(w, r, nil) + wsConn, upgradeErr := acceptWS(w, r) require.NoError(t, upgradeErr) - defer func() { _ = wsConn.Close() }() + defer func() { _ = wsConn.CloseNow() }() session := newYamuxServerSession(t, wsConn) defer func() { _ = session.Close() }() @@ -234,9 +264,9 @@ func TestLeash_Connect_PortForwardStreamDialsTarget(t *testing.T) { defer close(serverDone) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - wsConn, upgradeErr := testUpgrader.Upgrade(w, r, nil) + wsConn, upgradeErr := acceptWS(w, r) require.NoError(t, upgradeErr) - defer func() { _ = wsConn.Close() }() + defer func() { _ = wsConn.CloseNow() }() session := newYamuxServerSession(t, wsConn) defer func() { _ = session.Close() }() diff --git a/agent/leash/wsconn.go b/agent/leash/wsconn.go index 5ea4d14ed..2cd34c9b4 100644 --- a/agent/leash/wsconn.go +++ b/agent/leash/wsconn.go @@ -1,80 +1,46 @@ package leash import ( - "io" + "context" + "crypto/tls" "net" - "sync" - "time" + "net/http" - "github.com/gorilla/websocket" + "github.com/coder/websocket" ) -// wsNetConn adapts a gorilla WebSocket connection to the net.Conn interface -// so that yamux can use it as a byte-stream transport. -// gorilla websocket supports one concurrent reader and one concurrent writer; -// rMu serialises reads and wMu serialises writes accordingly. +// wsNetConn exposes a WebSocket connection as a net.Conn so that yamux can use +// it as a byte-stream transport. Each Write is sent as one binary message and +// reads stream the payload of incoming data messages. type wsNetConn struct { - conn *websocket.Conn - rMu sync.Mutex - wMu sync.Mutex - reader io.Reader + net.Conn + ws *websocket.Conn } // NewWSNetConn wraps conn as a net.Conn suitable for use with yamux. func NewWSNetConn(conn *websocket.Conn) net.Conn { - return &wsNetConn{conn: conn} -} - -func (c *wsNetConn) Read(b []byte) (int, error) { - c.rMu.Lock() - defer c.rMu.Unlock() - - for { - if c.reader != nil { - n, err := c.reader.Read(b) - if err == io.EOF { - c.reader = nil - continue - } - return n, err - } - - _, msg, err := c.conn.NextReader() - if err != nil { - return 0, err - } - c.reader = msg - } -} - -func (c *wsNetConn) Write(b []byte) (int, error) { - c.wMu.Lock() - defer c.wMu.Unlock() - - if err := c.conn.WriteMessage(websocket.BinaryMessage, b); err != nil { - return 0, err + return &wsNetConn{ + Conn: websocket.NetConn(context.Background(), conn, websocket.MessageBinary), + ws: conn, } - return len(b), nil } +// Close tears the connection down immediately without the WebSocket closing +// handshake. This matches the agent's previous raw-socket close and avoids +// blocking session teardown on a peer that has stopped responding. func (c *wsNetConn) Close() error { - return c.conn.Close() -} - -func (c *wsNetConn) LocalAddr() net.Addr { return c.conn.NetConn().LocalAddr() } -func (c *wsNetConn) RemoteAddr() net.Addr { return c.conn.NetConn().RemoteAddr() } - -func (c *wsNetConn) SetDeadline(t time.Time) error { - if err := c.conn.SetReadDeadline(t); err != nil { - return err + return c.ws.CloseNow() +} + +// newDialHTTPClient returns the HTTP client used for the WebSocket handshake. +// It dials directly, ignoring proxy environment variables, as the previous +// dialer did, and disables HTTP/2 negotiation because the Upgrade handshake +// requires HTTP/1.1. +func newDialHTTPClient() *http.Client { + return &http.Client{ + Transport: &http.Transport{ + DialContext: (&net.Dialer{}).DialContext, + TLSNextProto: map[string]func(string, *tls.Conn) http.RoundTripper{}, + }, } - return c.conn.SetWriteDeadline(t) -} - -func (c *wsNetConn) SetReadDeadline(t time.Time) error { - return c.conn.SetReadDeadline(t) -} - -func (c *wsNetConn) SetWriteDeadline(t time.Time) error { - return c.conn.SetWriteDeadline(t) } diff --git a/agent/leash/wsconn_internal_test.go b/agent/leash/wsconn_internal_test.go new file mode 100644 index 000000000..dc650caac --- /dev/null +++ b/agent/leash/wsconn_internal_test.go @@ -0,0 +1,20 @@ +package leash + +import ( + "net/http" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// TestNewDialHTTPClient_BypassesProxyAndDisablesHTTP2 verifies the handshake +// transport dials directly and never negotiates HTTP/2, which would break the +// WebSocket Upgrade. +func TestNewDialHTTPClient_BypassesProxyAndDisablesHTTP2(t *testing.T) { + tr, ok := newDialHTTPClient().Transport.(*http.Transport) + require.True(t, ok, "dial client must use an *http.Transport") + assert.Nil(t, tr.Proxy, "handshake must not honor proxy environment variables") + assert.NotNil(t, tr.TLSNextProto, "HTTP/2 must be disabled via an empty TLSNextProto map") + assert.Empty(t, tr.TLSNextProto) +} diff --git a/backend/go.mod b/backend/go.mod index c31147787..f73b854fb 100644 --- a/backend/go.mod +++ b/backend/go.mod @@ -5,12 +5,12 @@ go 1.27.4 require ( filippo.io/age v1.3.2 github.com/Wikid82/Hermes v0.3.7 + github.com/coder/websocket v1.8.15 github.com/gin-contrib/gzip v1.2.10 github.com/gin-gonic/gin v1.12.0 github.com/glebarez/sqlite v1.11.0 github.com/golang-jwt/jwt/v5 v5.3.1 github.com/google/uuid v1.6.0 - github.com/gorilla/websocket v1.5.3 github.com/hashicorp/golang-lru/v2 v2.0.7 github.com/hashicorp/yamux v0.1.2 github.com/minio/minio-go/v7 v7.4.1 diff --git a/backend/go.sum b/backend/go.sum index 11cd1ce3b..71e09822d 100644 --- a/backend/go.sum +++ b/backend/go.sum @@ -20,6 +20,8 @@ github.com/cespare/xxhash/v2 v2.3.0 h1:UL815xU9SqsFlibzuggzjXhog7bL6oX9BbNZnL2UF github.com/cespare/xxhash/v2 v2.3.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs= github.com/cloudwego/base64x v0.1.7 h1:NppS+Fgzg5ovhn4NkUXaDT3x9jldgH5ToMCqzBSi2zI= github.com/cloudwego/base64x v0.1.7/go.mod h1:Cu1PV9zfrSf7ET2tIbWbbEy7jO7HHJ13q4X2SQ8aWYg= +github.com/coder/websocket v1.8.15 h1:6B2JPeOGlpff2Uz6vOEH1Vzpi0iUz20A+lPVhPHtNUA= +github.com/coder/websocket v1.8.15/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/containerd/errdefs v1.0.0 h1:tg5yIfIlQIrxYtu9ajqY42W3lpS19XqdxRQeEwYG8PI= github.com/containerd/errdefs v1.0.0/go.mod h1:+YBYIdtsnF4Iw6nWZhJcqGSg/dwvV7tyJ/kCkyJ2k+M= github.com/containerd/errdefs/pkg v0.3.0 h1:9IKJ06FvyNlexW690DXuQNx2KA2cUJXx151Xdx3ZPPE= @@ -75,8 +77,6 @@ github.com/google/pprof v0.0.0-20261016150608-1facf8f91b50 h1:PlNWJRxu5APSV0pWqf github.com/google/pprof v0.0.0-20261016150608-1facf8f91b50/go.mod h1:yUo1AMMpsPIDFGuwDK8nLcf8+vvAd5l2O6JDZ6hpwrM= github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0= github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= -github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= -github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k= github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= github.com/hashicorp/yamux v0.1.2 h1:XtB8kyFOyHXYVFnwT5C3+Bdo8gArse7j2AQ0DA0Uey8= diff --git a/backend/internal/api/handlers/cerberus_logs_ws.go b/backend/internal/api/handlers/cerberus_logs_ws.go index 3eb3fac3d..26a595232 100644 --- a/backend/internal/api/handlers/cerberus_logs_ws.go +++ b/backend/internal/api/handlers/cerberus_logs_ws.go @@ -2,9 +2,11 @@ package handlers import ( + "context" "strings" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/google/uuid" @@ -45,8 +47,10 @@ func (h *CerberusLogsHandler) LiveLogs(c *gin.Context) { logger.Log().WithError(err).Error("Failed to upgrade Cerberus logs WebSocket") return } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() defer func() { - if err := conn.Close(); err != nil { + if err := conn.Close(websocket.StatusNormalClosure, ""); err != nil { logger.Log().WithError(err).Debug("Failed to close Cerberus logs WebSocket connection") } }() @@ -83,20 +87,8 @@ func (h *CerberusLogsHandler) LiveLogs(c *gin.Context) { logChan := h.watcher.Subscribe() defer h.watcher.Unsubscribe(logChan) - // Channel to detect client disconnect - done := make(chan struct{}) - go func() { - defer close(done) - for { - if _, _, err := conn.ReadMessage(); err != nil { - return - } - } - }() - - // Keep-alive ticker - ticker := time.NewTicker(logSocketPingPeriod) - defer ticker.Stop() + // Closes when the client disconnects or stops answering pings + done := watchLogSocket(ctx, conn) for { select { @@ -142,13 +134,6 @@ func (h *CerberusLogsHandler) LiveLogs(c *gin.Context) { h.tracker.UpdateActivity(subscriberID) } - case <-ticker.C: - // Send ping to keep connection alive - if err := writeLogSocketPing(conn); err != nil { - logger.Log().WithError(err).WithField("subscriber_id", subscriberID).Debug("Failed to send ping to Cerberus logs WebSocket") - return - } - case <-done: // Client disconnected logger.Log().WithField("subscriber_id", subscriberID).Info("Cerberus logs WebSocket client disconnected") diff --git a/backend/internal/api/handlers/cerberus_logs_ws_test.go b/backend/internal/api/handlers/cerberus_logs_ws_test.go index e96928f3d..f4732421b 100644 --- a/backend/internal/api/handlers/cerberus_logs_ws_test.go +++ b/backend/internal/api/handlers/cerberus_logs_ws_test.go @@ -12,9 +12,9 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/google/uuid" - "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -67,10 +67,9 @@ func TestCerberusLogsHandler_SuccessfulConnection(t *testing.T) { wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws" // Connect WebSocket - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + conn, resp, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - defer func() { _ = resp.Body.Close() }() - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() assert.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) } @@ -103,9 +102,9 @@ func TestCerberusLogsHandler_ReceiveLogEntries(t *testing.T) { // Connect WebSocket wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() // Give the subscription time to register and watcher to seek to end time.Sleep(300 * time.Millisecond) @@ -130,8 +129,7 @@ func TestCerberusLogsHandler_ReceiveLogEntries(t *testing.T) { _ = file.Sync() // Read the entry from WebSocket - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - _, msg, err := conn.ReadMessage() + msg, err := readWithin(t, conn, 2*time.Second) require.NoError(t, err) var entry models.SecurityLogEntry @@ -172,9 +170,9 @@ func TestCerberusLogsHandler_SourceFilter(t *testing.T) { // Connect with WAF source filter wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws?source=waf" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() time.Sleep(300 * time.Millisecond) @@ -213,8 +211,7 @@ func TestCerberusLogsHandler_SourceFilter(t *testing.T) { _ = file.Sync() // Read from WebSocket - should only get WAF entry - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - _, msg, err := conn.ReadMessage() + msg, err := readWithin(t, conn, 2*time.Second) require.NoError(t, err) var entry models.SecurityLogEntry @@ -252,9 +249,9 @@ func TestCerberusLogsHandler_BlockedOnlyFilter(t *testing.T) { // Connect with blocked_only filter wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws?blocked_only=true" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() time.Sleep(300 * time.Millisecond) @@ -292,8 +289,7 @@ func TestCerberusLogsHandler_BlockedOnlyFilter(t *testing.T) { _ = file.Sync() // Read from WebSocket - should only get blocked entry - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - _, msg, err := conn.ReadMessage() + msg, err := readWithin(t, conn, 2*time.Second) require.NoError(t, err) var entry models.SecurityLogEntry @@ -329,9 +325,9 @@ func TestCerberusLogsHandler_IPFilter(t *testing.T) { // Connect with IP filter wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws?ip=192.168" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() time.Sleep(300 * time.Millisecond) @@ -369,8 +365,7 @@ func TestCerberusLogsHandler_IPFilter(t *testing.T) { _ = file.Sync() // Read from WebSocket - should only get matching IP entry - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - _, msg, err := conn.ReadMessage() + msg, err := readWithin(t, conn, 2*time.Second) require.NoError(t, err) var entry models.SecurityLogEntry @@ -403,11 +398,11 @@ func TestCerberusLogsHandler_ClientDisconnect(t *testing.T) { defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) // Close the connection - _ = conn.Close() + _ = conn.CloseNow() // Give time for cleanup time.Sleep(100 * time.Millisecond) @@ -446,12 +441,12 @@ func TestCerberusLogsHandler_MultipleClients(t *testing.T) { // Close all connections after test for _, conn := range conns { if conn != nil { - _ = conn.Close() + _ = conn.CloseNow() } } }() for i := 0; i < 3; i++ { - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) conns[i] = conn } @@ -477,8 +472,7 @@ func TestCerberusLogsHandler_MultipleClients(t *testing.T) { // All clients should receive the entry for i, conn := range conns { - _ = conn.SetReadDeadline(time.Now().Add(2 * time.Second)) - _, msg, err := conn.ReadMessage() + msg, err := readWithin(t, conn, 2*time.Second) require.NoError(t, err, "Client %d should receive message", i) var entry models.SecurityLogEntry @@ -503,8 +497,8 @@ func TestCerberusLogsHandler_UpgradeFailure(t *testing.T) { w := httptest.NewRecorder() router.ServeHTTP(w, req) - // Should fail upgrade (400 Bad Request) - assert.Equal(t, http.StatusBadRequest, w.Code) + // Should fail upgrade (426 Upgrade Required: the Upgrade header is missing) + assert.Equal(t, http.StatusUpgradeRequired, w.Code) } // wsTestToken is a credential-like value used to check that it never reaches connection info or streamed entries. @@ -533,12 +527,9 @@ func TestCerberusLogsHandler_ConnectionListOmitsToken(t *testing.T) { defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws?source=waf&token=" + wsTestToken - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - defer func() { _ = conn.Close() }() - if resp != nil { - defer func() { _ = resp.Body.Close() }() - } + defer func() { _ = conn.CloseNow() }() waitFor(t, func() bool { return tracker.GetCount() == 1 }) @@ -572,12 +563,9 @@ func TestCerberusLogsHandler_StreamedEntryOmitsToken(t *testing.T) { defer server.Close() wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/ws" - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - defer func() { _ = conn.Close() }() - if resp != nil { - defer func() { _ = resp.Body.Close() }() - } + defer func() { _ = conn.CloseNow() }() time.Sleep(300 * time.Millisecond) @@ -593,8 +581,7 @@ func TestCerberusLogsHandler_StreamedEntryOmitsToken(t *testing.T) { require.NoError(t, err) require.NoError(t, file.Sync()) - require.NoError(t, conn.SetReadDeadline(time.Now().Add(2*time.Second))) - _, msg, err := conn.ReadMessage() + msg, err := readWithin(t, conn, 2*time.Second) require.NoError(t, err) assert.NotContains(t, string(msg), wsTestToken) diff --git a/backend/internal/api/handlers/hecate_handler_test.go b/backend/internal/api/handlers/hecate_handler_test.go index f7e802c15..07d8d537b 100644 --- a/backend/internal/api/handlers/hecate_handler_test.go +++ b/backend/internal/api/handlers/hecate_handler_test.go @@ -12,8 +12,8 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "gorm.io/driver/sqlite" @@ -890,10 +890,7 @@ func TestHecateWSHandler_StreamLogs_UpgradeAndStream(t *testing.T) { t.Cleanup(srv.Close) wsURL := toWebSocketURL(srv.URL) + "/ws/tunnels/" + tunnelUUID + "/logs" - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) - if resp != nil { - defer resp.Body.Close() - } + conn, resp, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) @@ -901,7 +898,7 @@ func TestHecateWSHandler_StreamLogs_UpgradeAndStream(t *testing.T) { return tracker.GetCount() == 1 }) - require.NoError(t, conn.Close()) + require.NoError(t, conn.CloseNow()) waitFor(t, func() bool { return tracker.GetCount() == 0 @@ -1174,7 +1171,7 @@ func TestHecateWSHandler_StreamLogs_UpgradeError(t *testing.T) { srv := httptest.NewServer(r) t.Cleanup(srv.Close) - // Plain HTTP GET — no WebSocket upgrade headers — causes upgrader.Upgrade to fail. + // Plain HTTP GET — no WebSocket upgrade headers — causes websocket.Accept to fail. resp, err := http.Get(srv.URL + "/ws/tunnels/" + tunnelUUID + "/logs") //nolint:noctx require.NoError(t, err) _ = resp.Body.Close() @@ -1205,18 +1202,17 @@ func TestHecateWSHandler_StreamLogs_SubClose(t *testing.T) { t.Cleanup(srv.Close) wsURL := toWebSocketURL(srv.URL) + "/ws/tunnels/" + tunnelUUID + "/logs" - conn, resp, dialErr := websocket.DefaultDialer.Dial(wsURL, nil) + conn, resp, dialErr := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, dialErr) - defer resp.Body.Close() require.Equal(t, http.StatusSwitchingProtocols, resp.StatusCode) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) // Write a live line so the sub case fires with ok=true (covers lines 90-96 assignment). buf.Write("live line") // Give the handler goroutine a moment to process the live line. waitFor(t, func() bool { - _, _, readErr := conn.ReadMessage() + _, readErr := readWithin(t, conn, 3*time.Second) return readErr == nil }) @@ -1225,8 +1221,7 @@ func TestHecateWSHandler_StreamLogs_SubClose(t *testing.T) { // Wait for the WS connection to be terminated by the handler. waitFor(t, func() bool { - _ = conn.SetReadDeadline(time.Now().Add(50 * time.Millisecond)) - _, _, readErr := conn.ReadMessage() + _, readErr := readWithin(t, conn, 50*time.Millisecond) return readErr != nil }) } diff --git a/backend/internal/api/handlers/hecate_ws_handler.go b/backend/internal/api/handlers/hecate_ws_handler.go index 541b5b0fe..291b6fb90 100644 --- a/backend/internal/api/handlers/hecate_ws_handler.go +++ b/backend/internal/api/handlers/hecate_ws_handler.go @@ -1,12 +1,13 @@ package handlers import ( + "context" "net/http" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/google/uuid" - "github.com/gorilla/websocket" "github.com/Wikid82/charon/backend/internal/logger" "github.com/Wikid82/charon/backend/internal/services" @@ -25,7 +26,7 @@ func NewHecateWSHandler(svc *services.HecateService, tracker *services.WebSocket // StreamLogs upgrades the request to WebSocket and streams log lines from // the tunnel's ring buffer. It replays buffered lines first, then delivers -// new lines as they arrive. A 30 s ping keeps the connection alive. +// new lines as they arrive. The connection is kept alive with pings. func (h *HecateWSHandler) StreamLogs(c *gin.Context) { tunnelUUID := c.Param("uuid") @@ -42,8 +43,11 @@ func (h *HecateWSHandler) StreamLogs(c *gin.Context) { logger.Log().WithError(upgradeErr).Error("hecate ws: upgrade failed") return } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() defer func() { - if closeErr := conn.Close(); closeErr != nil { + if closeErr := conn.Close(websocket.StatusNormalClosure, ""); closeErr != nil { logger.Log().WithError(closeErr).Debug("hecate ws: close error") } }() @@ -64,25 +68,15 @@ func (h *HecateWSHandler) StreamLogs(c *gin.Context) { // Replay buffered lines. for _, line := range buf.ReadAll() { - if writeErr := conn.WriteMessage(websocket.TextMessage, []byte(line)); writeErr != nil { + if writeErr := writeLogSocketText(conn, []byte(line)); writeErr != nil { return } } sub := buf.Subscribe() - done := make(chan struct{}) - go func() { - defer close(done) - for { - if _, _, readErr := conn.ReadMessage(); readErr != nil { - return - } - } - }() - - ticker := time.NewTicker(30 * time.Second) - defer ticker.Stop() + // Closes when the client disconnects or stops answering pings + done := watchLogSocket(ctx, conn) for { select { @@ -92,11 +86,7 @@ func (h *HecateWSHandler) StreamLogs(c *gin.Context) { if !ok { return } - if writeErr := conn.WriteMessage(websocket.TextMessage, []byte(line)); writeErr != nil { - return - } - case <-ticker.C: - if pingErr := conn.WriteMessage(websocket.PingMessage, nil); pingErr != nil { + if writeErr := writeLogSocketText(conn, []byte(line)); writeErr != nil { return } } diff --git a/backend/internal/api/handlers/log_socket.go b/backend/internal/api/handlers/log_socket.go new file mode 100644 index 000000000..0b65d03e7 --- /dev/null +++ b/backend/internal/api/handlers/log_socket.go @@ -0,0 +1,120 @@ +package handlers + +import ( + "context" + "encoding/json" + "fmt" + "sync/atomic" + "time" + + "github.com/coder/websocket" + + "github.com/Wikid82/charon/backend/internal/logger" +) + +// logSocketReadLimit caps the size of a single client message. +const logSocketReadLimit = 4096 + +// logSocketTimers holds the log socket timers. Tests shorten them, and a +// handler goroutine from an earlier test may still be reading them, so the +// whole snapshot is swapped atomically. +type logSocketTimers struct { + // pongWait is how long a ping may go unanswered before the connection is closed. + // The pong deadline restarts at each answered ping. + pongWait time.Duration + // pingPeriod is how often the server pings the client. + pingPeriod time.Duration + // writeWait bounds every data write. + writeWait time.Duration +} + +var logSocketTiming atomic.Pointer[logSocketTimers] + +func init() { + logSocketTiming.Store(&logSocketTimers{ + pongWait: 60 * time.Second, + pingPeriod: 30 * time.Second, + writeWait: 10 * time.Second, + }) +} + +// currentLogSocketTiming returns the timers in effect for a log socket. +func currentLogSocketTiming() logSocketTimers { + return *logSocketTiming.Load() +} + +// configureLogSocket applies the read limit shared by every log WebSocket. +// Call it right after the upgrade, before watchLogSocket. +func configureLogSocket(conn *websocket.Conn) { + conn.SetReadLimit(logSocketReadLimit) +} + +// watchLogSocket keeps a log WebSocket alive and detects client disconnects. +// It drains client messages so pings and close frames are processed, and it +// pings the client every pingPeriod. The returned channel closes when +// the client disconnects or a ping goes unanswered past its pong deadline. In +// both cases the connection is closed. Cancel ctx when the handler returns. +func watchLogSocket(ctx context.Context, conn *websocket.Conn) <-chan struct{} { + go keepLogSocketAlive(ctx, conn) + + done := make(chan struct{}) + go func() { + defer close(done) + for { + if _, _, err := conn.Read(ctx); err != nil { + return + } + } + }() + return done +} + +// keepLogSocketAlive sends a ping every pingPeriod. Each ping must be +// answered before the pong deadline, which is pongWait after the last +// answered ping, or after the upgrade if none was answered yet. A missed +// deadline closes the connection. Pongs are read by the reader started in +// watchLogSocket, which is why this runs alongside it. +func keepLogSocketAlive(ctx context.Context, conn *websocket.Conn) { + ticker := time.NewTicker(currentLogSocketTiming().pingPeriod) + defer ticker.Stop() + + lastPong := time.Now() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + + pingCtx, cancel := context.WithDeadline(ctx, lastPong.Add(currentLogSocketTiming().pongWait)) + err := conn.Ping(pingCtx) + cancel() + if err != nil { + if ctx.Err() == nil { + logger.Log().WithError(err).Debug("Log socket keep-alive failed, closing connection") + _ = conn.CloseNow() + } + return + } + lastPong = time.Now() + } +} + +// writeLogSocketText sends data as one text message under the write wait. +func writeLogSocketText(conn *websocket.Conn, data []byte) error { + ctx, cancel := context.WithTimeout(context.Background(), currentLogSocketTiming().writeWait) + defer cancel() + if err := conn.Write(ctx, websocket.MessageText, data); err != nil { + return fmt.Errorf("write log socket message: %w", err) + } + return nil +} + +// writeLogSocketJSON sends v as a JSON text message under the write wait. +func writeLogSocketJSON(conn *websocket.Conn, v any) error { + data, err := json.Marshal(v) + if err != nil { + return fmt.Errorf("marshal log socket message: %w", err) + } + return writeLogSocketText(conn, data) +} diff --git a/backend/internal/api/handlers/log_socket_deadlines_test.go b/backend/internal/api/handlers/log_socket_deadlines_test.go index 80fcc18a5..2e5d3c585 100644 --- a/backend/internal/api/handlers/log_socket_deadlines_test.go +++ b/backend/internal/api/handlers/log_socket_deadlines_test.go @@ -3,41 +3,37 @@ package handlers import ( "context" "errors" - "fmt" - "net" "net/http" "net/http/httptest" "strings" "testing" "time" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/stretchr/testify/require" ) -// withLogSocketTimings overrides the log socket deadlines for one test and -// restores the production values on cleanup. Tests in this file do not run in -// parallel, so the package-level values are never read concurrently with a write. +// withLogSocketTimings overrides the log socket timers for one test and restores +// the previous values on cleanup. Handlers read the timers when they run, so a +// test that expects a handler to see these timers must wait for that handler +// before its cleanup restores the values. func withLogSocketTimings(t *testing.T, pongWait, writeWait, pingPeriod time.Duration) { t.Helper() - prevPong, prevWrite, prevPing := logSocketPongWait, logSocketWriteWait, logSocketPingPeriod - logSocketPongWait, logSocketWriteWait, logSocketPingPeriod = pongWait, writeWait, pingPeriod - t.Cleanup(func() { - logSocketPongWait, logSocketWriteWait, logSocketPingPeriod = prevPong, prevWrite, prevPing - }) + prev := logSocketTiming.Load() + logSocketTiming.Store(&logSocketTimers{pongWait: pongWait, pingPeriod: pingPeriod, writeWait: writeWait}) + t.Cleanup(func() { logSocketTiming.Store(prev) }) } // newLogSocketTestServer serves one WebSocket connection, configures it with // configureLogSocket and hands it to handle. func newLogSocketTestServer(t *testing.T, handle func(conn *websocket.Conn)) *httptest.Server { t.Helper() - testUpgrader := websocket.Upgrader{CheckOrigin: func(r *http.Request) bool { return r.Header.Get("Origin") == "" }} srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := testUpgrader.Upgrade(w, r, nil) + conn, err := websocket.Accept(w, r, nil) if err != nil { return } - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() configureLogSocket(conn) handle(conn) })) @@ -45,50 +41,56 @@ func newLogSocketTestServer(t *testing.T, handle func(conn *websocket.Conn)) *ht return srv } +// dialLogSocketTestServer dials srv and closes the connection at test cleanup. func dialLogSocketTestServer(t *testing.T, srv *httptest.Server) *websocket.Conn { t.Helper() - return dialWithDialer(t, srv, websocket.DefaultDialer) -} - -func dialWithDialer(t *testing.T, srv *httptest.Server, dialer *websocket.Dialer) *websocket.Conn { - t.Helper() - conn, resp, err := dialer.Dial(toWebSocketURL(srv.URL), nil) - if resp != nil && resp.Body != nil { - _ = resp.Body.Close() - } + conn, _, err := websocket.Dial(context.Background(), toWebSocketURL(srv.URL), nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) return conn } +// watchLogSocketUntilDone runs watchLogSocket on conn and returns a channel that +// closes once the watch ends, which happens when the connection drops. +func watchLogSocketUntilDone(conn *websocket.Conn) <-chan struct{} { + ctx, cancel := context.WithCancel(context.Background()) + done := watchLogSocket(ctx, conn) + stopped := make(chan struct{}) + go func() { + <-done + cancel() + close(stopped) + }() + return stopped +} + func TestConfigureLogSocket_OversizeClientMessageClosesConnection(t *testing.T) { serverErr := make(chan error, 1) srv := newLogSocketTestServer(t, func(conn *websocket.Conn) { - _, _, err := conn.ReadMessage() + _, _, err := conn.Read(context.Background()) serverErr <- err }) client := dialLogSocketTestServer(t, srv) oversize := strings.Repeat("a", logSocketReadLimit+1) - require.NoError(t, client.WriteMessage(websocket.TextMessage, []byte(oversize))) + require.NoError(t, client.Write(context.Background(), websocket.MessageText, []byte(oversize))) select { - case err := <-serverErr: - require.ErrorIs(t, err, websocket.ErrReadLimit) + case srvErr := <-serverErr: + require.Error(t, srvErr, "server read should fail on an oversize message") case <-time.After(5 * time.Second): t.Fatal("server did not reject the oversize message") } - // The server closed the connection, so the client's next read must fail. - require.NoError(t, client.SetReadDeadline(time.Now().Add(5*time.Second))) - _, _, err := client.ReadMessage() - require.Error(t, err) + // The server closes with 1009 (message too big), which the client sees as its next read error. + _, err := readWithin(t, client, 5*time.Second) + require.Equal(t, websocket.StatusMessageTooBig, websocket.CloseStatus(err), "expected message-too-big close, got %v", err) } func TestConfigureLogSocket_MessageWithinLimitIsAccepted(t *testing.T) { serverMsg := make(chan string, 1) srv := newLogSocketTestServer(t, func(conn *websocket.Conn) { - _, data, err := conn.ReadMessage() + _, data, err := conn.Read(context.Background()) if err == nil { serverMsg <- string(data) } @@ -96,7 +98,7 @@ func TestConfigureLogSocket_MessageWithinLimitIsAccepted(t *testing.T) { client := dialLogSocketTestServer(t, srv) payload := strings.Repeat("b", logSocketReadLimit) - require.NoError(t, client.WriteMessage(websocket.TextMessage, []byte(payload))) + require.NoError(t, client.Write(context.Background(), websocket.MessageText, []byte(payload))) select { case got := <-serverMsg: @@ -106,148 +108,70 @@ func TestConfigureLogSocket_MessageWithinLimitIsAccepted(t *testing.T) { } } -func TestConfigureLogSocket_SilentPeerHitsReadDeadline(t *testing.T) { - withLogSocketTimings(t, 200*time.Millisecond, 10*time.Second, 30*time.Second) +// TestWatchLogSocket_SilentPeerIsDisconnected covers a peer that never answers +// pings. The client never reads, so no pong is sent and the watch must end once +// the pong deadline passes. +func TestWatchLogSocket_SilentPeerIsDisconnected(t *testing.T) { + withLogSocketTimings(t, 200*time.Millisecond, 10*time.Second, 20*time.Millisecond) - serverErr := make(chan error, 1) + dropped := make(chan struct{}) srv := newLogSocketTestServer(t, func(conn *websocket.Conn) { - // The ping is never answered because the client does not read. - if err := writeLogSocketPing(conn); err != nil { - serverErr <- err - return - } - _, _, err := conn.ReadMessage() - serverErr <- err + <-watchLogSocketUntilDone(conn) + close(dropped) }) _ = dialLogSocketTestServer(t, srv) select { - case err := <-serverErr: - var netErr net.Error - require.True(t, errors.As(err, &netErr) && netErr.Timeout(), "expected read timeout, got %v", err) + case <-dropped: case <-time.After(5 * time.Second): - t.Fatal("read deadline did not fire without a pong") + t.Fatal("keep-alive did not drop a peer that never answered pings") } } -func TestConfigureLogSocket_PongRefreshesReadDeadline(t *testing.T) { +// TestWatchLogSocket_PongsKeepConnectionAlive covers a responsive peer. Each +// answered ping must move the pong deadline forward, so the connection survives +// several pong windows. +func TestWatchLogSocket_PongsKeepConnectionAlive(t *testing.T) { const pongWait = 300 * time.Millisecond withLogSocketTimings(t, pongWait, 10*time.Second, 50*time.Millisecond) - // The pinger reads the package timers, so the cleanup must wait for it to - // exit before withLogSocketTimings restores the production values. - stop := make(chan struct{}) - pingerDone := make(chan struct{}) - t.Cleanup(func() { - close(stop) - <-pingerDone - }) - - serverErr := make(chan error, 1) + dropped := make(chan struct{}) srv := newLogSocketTestServer(t, func(conn *websocket.Conn) { - // Keep pinging for longer than one pong wait. Each pong must extend the deadline. - go func() { - defer close(pingerDone) - ticker := time.NewTicker(logSocketPingPeriod) - defer ticker.Stop() - for { - select { - case <-stop: - return - case <-ticker.C: - if err := writeLogSocketPing(conn); err != nil { - return - } - } - } - }() - _, _, err := conn.ReadMessage() - serverErr <- err + <-watchLogSocketUntilDone(conn) + close(dropped) }) client := dialLogSocketTestServer(t, srv) - // The default ping handler answers every ping with a pong while the client reads. + // Reading makes the client answer every ping with a pong. + readDone := make(chan struct{}) go func() { + defer close(readDone) for { - if _, _, err := client.ReadMessage(); err != nil { + if _, _, err := client.Read(context.Background()); err != nil { return } } }() select { - case err := <-serverErr: - t.Fatalf("connection dropped although the peer answered pings: %v", err) + case <-dropped: + t.Fatal("connection dropped although the peer answered pings") case <-time.After(3 * pongWait): } - require.NoError(t, client.Close()) + + require.NoError(t, client.CloseNow()) select { - case err := <-serverErr: - require.Error(t, err, "server read must fail once the client closes") + case <-dropped: case <-time.After(5 * time.Second): - t.Fatal("server read did not end after the client closed") + t.Fatal("keep-alive did not end after the client closed") } + <-readDone } -// stalledSocketBufferBytes is the send buffer on the server and the receive buffer -// on the client of a stalled-write test. Setting SO_SNDBUF and SO_RCVBUF locks both -// sizes, so the kernel cannot autotune them upward and the queued bytes stay bounded. -const stalledSocketBufferBytes = 4096 - // stalledWriteBudget caps how long a stalled-write test keeps writing before it // gives up. A write that hits its deadline ends the loop well before this budget. const stalledWriteBudget = 10 * time.Second -// shrinkServerSendBuffer bounds the server-side send buffer of conn. -func shrinkServerSendBuffer(conn *websocket.Conn) error { - tcpConn, ok := conn.NetConn().(*net.TCPConn) - if !ok { - return fmt.Errorf("log socket test connection is %T, want *net.TCPConn", conn.NetConn()) - } - if err := tcpConn.SetWriteBuffer(stalledSocketBufferBytes); err != nil { - return fmt.Errorf("shrink server send buffer: %w", err) - } - return nil -} - -// dialStalledReaderClient dials srv with a small receive buffer on the client side. -// The buffer is set inside NetDialContext, before the HTTP upgrade request is sent. -func dialStalledReaderClient(t *testing.T, srv *httptest.Server) *websocket.Conn { - t.Helper() - dialer := &websocket.Dialer{ - NetDialContext: func(ctx context.Context, network, addr string) (net.Conn, error) { - var d net.Dialer - conn, err := d.DialContext(ctx, network, addr) - if err != nil { - return nil, fmt.Errorf("dial stalled reader client: %w", err) - } - tcpConn, ok := conn.(*net.TCPConn) - if !ok { - _ = conn.Close() - return nil, fmt.Errorf("stalled reader connection is %T, want *net.TCPConn", conn) - } - if err := tcpConn.SetReadBuffer(stalledSocketBufferBytes); err != nil { - _ = conn.Close() - return nil, fmt.Errorf("shrink client receive buffer: %w", err) - } - return conn, nil - }, - } - return dialWithDialer(t, srv, dialer) -} - -// fillSocketUntilWriteError calls write until it fails. The client never reads, -// so the bounded buffers fill and a write fails on its deadline. -func fillSocketUntilWriteError(write func() error) error { - deadline := time.Now().Add(stalledWriteBudget) - for time.Now().Before(deadline) { - if err := write(); err != nil { - return err - } - } - return fmt.Errorf("writes never failed within %s", stalledWriteBudget) -} - // runStalledWriteTest runs write on the server side of a connection whose client // never reads, and returns the first write error. The caller must have applied // withLogSocketTimings first: the cleanup registered here runs before that restore, @@ -265,14 +189,10 @@ func runStalledWriteTest(t *testing.T, write func(conn *websocket.Conn) error) e srv := newLogSocketTestServer(t, func(conn *websocket.Conn) { defer close(resultReady) - if err := shrinkServerSendBuffer(conn); err != nil { - writeErr = err - return - } writeErr = fillSocketUntilWriteError(func() error { return write(conn) }) }) - // Registered after the wait cleanup, so the client closes before the handler is awaited. - _ = dialStalledReaderClient(t, srv) + // The client is dialed but never reads, so the server's writes back up. + _ = dialLogSocketTestServer(t, srv) select { case <-resultReady: @@ -283,11 +203,23 @@ func runStalledWriteTest(t *testing.T, write func(conn *websocket.Conn) error) e } } -// requireWriteTimeout asserts err is a network timeout from a write deadline. +// fillSocketUntilWriteError calls write until it fails. The client never reads, +// so the kernel buffers fill and a write fails on its deadline. +func fillSocketUntilWriteError(write func() error) error { + deadline := time.Now().Add(stalledWriteBudget) + for time.Now().Before(deadline) { + if err := write(); err != nil { + return err + } + } + return errors.New("writes never failed within " + stalledWriteBudget.String()) +} + +// requireWriteTimeout asserts err comes from a write deadline. A coder write that +// hits its context deadline reports context.DeadlineExceeded. func requireWriteTimeout(t *testing.T, err error) { t.Helper() - var netErr net.Error - require.True(t, errors.As(err, &netErr) && netErr.Timeout(), "expected write timeout, got %v", err) + require.ErrorIs(t, err, context.DeadlineExceeded, "expected write timeout, got %v", err) } func TestWriteLogSocketJSON_SetsWriteDeadline(t *testing.T) { @@ -300,9 +232,12 @@ func TestWriteLogSocketJSON_SetsWriteDeadline(t *testing.T) { requireWriteTimeout(t, err) } -func TestWriteLogSocketPing_SetsWriteDeadline(t *testing.T) { +func TestWriteLogSocketText_SetsWriteDeadline(t *testing.T) { withLogSocketTimings(t, 60*time.Second, 100*time.Millisecond, 30*time.Second) - err := runStalledWriteTest(t, writeLogSocketPing) + payload := []byte(strings.Repeat("d", 64*1024)) + err := runStalledWriteTest(t, func(conn *websocket.Conn) error { + return writeLogSocketText(conn, payload) + }) requireWriteTimeout(t, err) } diff --git a/backend/internal/api/handlers/log_socket_failures_test.go b/backend/internal/api/handlers/log_socket_failures_test.go index 8102ba6a6..fd42e822a 100644 --- a/backend/internal/api/handlers/log_socket_failures_test.go +++ b/backend/internal/api/handlers/log_socket_failures_test.go @@ -9,8 +9,8 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -18,14 +18,14 @@ import ( ) // newClosedLogSocketConn serves one WebSocket connection, closes it on the server -// side, and hands the closed connection to check. Deadline setters on a closed +// side, and hands the closed connection to check. Writes and pings on a closed // connection fail, which lets tests reach the error branches of the log socket helpers. func newClosedLogSocketConn(t *testing.T, check func(conn *websocket.Conn)) { t.Helper() done := make(chan struct{}) srv := newLogSocketTestServer(t, func(conn *websocket.Conn) { defer close(done) - require.NoError(t, conn.Close()) + require.NoError(t, conn.CloseNow()) check(conn) }) client := dialLogSocketTestServer(t, srv) @@ -34,17 +34,27 @@ func newClosedLogSocketConn(t *testing.T, check func(conn *websocket.Conn)) { case <-time.After(5 * time.Second): t.Fatal("server handler did not run") } - _ = client.Close() + _ = client.CloseNow() } -func TestConfigureLogSocket_ClosedConnectionKeepsGoingWithoutDeadlines(t *testing.T) { +func TestConfigureLogSocket_ClosedConnectionDoesNotPanic(t *testing.T) { newClosedLogSocketConn(t, func(conn *websocket.Conn) { assert.NotPanics(t, func() { configureLogSocket(conn) }) + }) +} + +// TestWatchLogSocket_ClosedConnectionEndsImmediately covers a connection that is +// already closed when the watch starts. The watch must end without waiting for a +// ping period. +func TestWatchLogSocket_ClosedConnectionEndsImmediately(t *testing.T) { + withLogSocketTimings(t, 60*time.Second, 10*time.Second, time.Hour) - // The pong handler refreshes the read deadline, which fails on a closed connection. - pong := conn.PongHandler() - require.NotNil(t, pong) - assert.Error(t, pong("")) + newClosedLogSocketConn(t, func(conn *websocket.Conn) { + select { + case <-watchLogSocketUntilDone(conn): + case <-time.After(5 * time.Second): + t.Fatal("watch did not end on a closed connection") + } }) } @@ -53,15 +63,23 @@ func TestWriteLogSocketHelpers_ClosedConnectionReturnsWriteError(t *testing.T) { newClosedLogSocketConn(t, func(conn *websocket.Conn) { require.Error(t, writeLogSocketJSON(conn, map[string]string{"data": "x"})) - require.Error(t, writeLogSocketPing(conn)) + require.Error(t, writeLogSocketText(conn, []byte("x"))) + }) +} + +// TestWriteLogSocketJSON_UnmarshalableValueFails covers a value encoding/json +// cannot marshal. The connection is never touched. +func TestWriteLogSocketJSON_UnmarshalableValueFails(t *testing.T) { + newClosedLogSocketConn(t, func(conn *websocket.Conn) { + require.Error(t, writeLogSocketJSON(conn, map[string]any{"bad": make(chan int)})) }) } -// TestLogsWSHandler_FailedKeepAlivePingEndsStream runs the stream with a write -// deadline that has already passed, so the keep-alive ping fails while the client -// stays connected. The handler must end the stream instead of looping. -func TestLogsWSHandler_FailedKeepAlivePingEndsStream(t *testing.T) { - withLogSocketTimings(t, 60*time.Second, time.Nanosecond, 10*time.Millisecond) +// TestLogsWSHandler_UnansweredPingEndsStream runs the stream with a pong deadline +// that passes before any pong can arrive. The client never reads, so no pong is +// sent. The handler must close the stream instead of looping. +func TestLogsWSHandler_UnansweredPingEndsStream(t *testing.T) { + withLogSocketTimings(t, 50*time.Millisecond, time.Second, 10*time.Millisecond) // The handler reads the package timers, so the cleanup waits for it to return // before withLogSocketTimings restores the production values. @@ -75,23 +93,22 @@ func TestLogsWSHandler_FailedKeepAlivePingEndsStream(t *testing.T) { srv := httptest.NewServer(router) t.Cleanup(srv.Close) - conn, resp, err := websocket.DefaultDialer.Dial(toWebSocketURL(srv.URL)+"/logs", nil) - if resp != nil { - _ = resp.Body.Close() - } + conn, _, err := websocket.Dial(context.Background(), toWebSocketURL(srv.URL)+"/logs", nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) t.Cleanup(func() { <-handlerDone }) - require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second))) - _, _, err = conn.ReadMessage() - require.Error(t, err, "the server must close the stream after the failed ping") + // Wait past the pong deadline before reading. Reading answers pings, so the + // client must stay silent until the server has closed the stream. + time.Sleep(200 * time.Millisecond) + err = readUntilStreamEnds(t, conn) + require.Error(t, err, "the server must close the stream after the unanswered ping") } -// TestCerberusLogsHandler_FailedKeepAlivePingEndsStream mirrors the logs stream -// test for the Cerberus log stream, whose keep-alive failure is logged before return. -func TestCerberusLogsHandler_FailedKeepAlivePingEndsStream(t *testing.T) { - withLogSocketTimings(t, 60*time.Second, time.Nanosecond, 10*time.Millisecond) +// TestCerberusLogsHandler_UnansweredPingEndsStream mirrors the logs stream test +// for the Cerberus log stream, whose keep-alive failure ends the stream. +func TestCerberusLogsHandler_UnansweredPingEndsStream(t *testing.T) { + withLogSocketTimings(t, 50*time.Millisecond, time.Second, 10*time.Millisecond) logPath := filepath.Join(t.TempDir(), "access.log") // #nosec G304 -- Test fixture uses controlled path from t.TempDir() @@ -113,15 +130,14 @@ func TestCerberusLogsHandler_FailedKeepAlivePingEndsStream(t *testing.T) { t.Cleanup(srv.Close) wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) t.Cleanup(func() { <-handlerDone }) - if resp != nil { - _ = resp.Body.Close() - } - require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second))) - _, _, err = conn.ReadMessage() - require.Error(t, err, "the server must close the stream after the failed ping") + // Wait past the pong deadline before reading. Reading answers pings, so the + // client must stay silent until the server has closed the stream. + time.Sleep(200 * time.Millisecond) + err = readUntilStreamEnds(t, conn) + require.Error(t, err, "the server must close the stream after the unanswered ping") } diff --git a/backend/internal/api/handlers/log_socket_stream_write_test.go b/backend/internal/api/handlers/log_socket_stream_write_test.go index 65fd2b5c1..491746e6a 100644 --- a/backend/internal/api/handlers/log_socket_stream_write_test.go +++ b/backend/internal/api/handlers/log_socket_stream_write_test.go @@ -10,9 +10,9 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/google/uuid" - "github.com/gorilla/websocket" "github.com/stretchr/testify/require" charonlogger "github.com/Wikid82/charon/backend/internal/logger" @@ -36,12 +36,9 @@ func TestLogsWSHandler_FailedEntryWriteEndsStream(t *testing.T) { srv := httptest.NewServer(router) t.Cleanup(srv.Close) - conn, resp, err := websocket.DefaultDialer.Dial(toWebSocketURL(srv.URL)+"/logs", nil) - if resp != nil { - _ = resp.Body.Close() - } + conn, _, err := websocket.Dial(context.Background(), toWebSocketURL(srv.URL)+"/logs", nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) t.Cleanup(func() { <-handlerDone }) // Log until the handler has subscribed and written, then wait for it to close the stream. @@ -59,8 +56,9 @@ func TestLogsWSHandler_FailedEntryWriteEndsStream(t *testing.T) { } }() - require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second))) - _, _, err = conn.ReadMessage() + // A write that starts before its deadline can still reach the peer, so the + // stream may deliver a line first. It must still end on a later failed write. + err = readUntilStreamEnds(t, conn) require.Error(t, err, "the server must close the stream after a failed entry write") } @@ -90,13 +88,10 @@ func TestCerberusLogsHandler_FailedEntryWriteEndsStream(t *testing.T) { t.Cleanup(srv.Close) wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // WebSocket Dial response body is consumed by the dial require.NoError(t, err) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) t.Cleanup(func() { <-handlerDone }) - if resp != nil { - _ = resp.Body.Close() - } // Append lines until the stream ends. Every line is a valid access entry. stop := make(chan struct{}) @@ -113,7 +108,8 @@ func TestCerberusLogsHandler_FailedEntryWriteEndsStream(t *testing.T) { } }() - require.NoError(t, conn.SetReadDeadline(time.Now().Add(5*time.Second))) - _, _, err = conn.ReadMessage() + // A write that starts before its deadline can still reach the peer, so the + // stream may deliver a line first. It must still end on a later failed write. + err = readUntilStreamEnds(t, conn) require.Error(t, err, "the server must close the stream after a failed entry write") } diff --git a/backend/internal/api/handlers/logs_ws.go b/backend/internal/api/handlers/logs_ws.go index 21c86f5ca..02a40dabb 100644 --- a/backend/internal/api/handlers/logs_ws.go +++ b/backend/internal/api/handlers/logs_ws.go @@ -1,94 +1,18 @@ package handlers import ( - "net/http" - "net/url" + "context" "strings" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" "github.com/google/uuid" - "github.com/gorilla/websocket" "github.com/Wikid82/charon/backend/internal/logger" "github.com/Wikid82/charon/backend/internal/services" ) -var upgrader = websocket.Upgrader{ - ReadBufferSize: 1024, - WriteBufferSize: 1024, - CheckOrigin: func(r *http.Request) bool { - origin := r.Header.Get("Origin") - if origin == "" { - // No Origin header — non-browser client or same-origin request. - return true - } - originURL, err := url.Parse(origin) - if err != nil { - return false - } - requestHost := r.Host - if forwardedHost := r.Header.Get("X-Forwarded-Host"); forwardedHost != "" { - requestHost = forwardedHost - } - return originURL.Host == requestHost - }, -} - -// upgradeWebSocket upgrades an HTTP connection to WebSocket using the -// shared, origin-checked upgrader above. All WS handlers in this package -// MUST go through this helper rather than calling upgrader.Upgrade -// directly, so the origin check stays visible to single-file static -// analysis at every call site. -func upgradeWebSocket(c *gin.Context) (*websocket.Conn, error) { - return upgrader.Upgrade(c.Writer, c.Request, nil) -} - -// logSocketReadLimit caps the size of a single client message. -const logSocketReadLimit = 4096 - -// Log socket timers. They are variables only so tests can shorten them. -var ( - // logSocketPongWait is how long a client may stay silent before its read fails. - logSocketPongWait = 60 * time.Second - // logSocketPingPeriod is how often the server pings the client. - logSocketPingPeriod = 30 * time.Second - // logSocketWriteWait bounds every write, including pings. - logSocketWriteWait = 10 * time.Second -) - -// configureLogSocket applies the read and write limits shared by every log -// WebSocket. Call it right after the upgrade. Each pong from the client -// extends the read deadline, so an idle but responsive client stays open. -func configureLogSocket(conn *websocket.Conn) { - conn.SetReadLimit(logSocketReadLimit) - if err := conn.SetReadDeadline(time.Now().Add(logSocketPongWait)); err != nil { - logger.Log().WithError(err).Debug("Failed to set initial log socket read deadline") - } - conn.SetPongHandler(func(string) error { - return conn.SetReadDeadline(time.Now().Add(logSocketPongWait)) - }) -} - -// armLogSocketWriteDeadline refreshes the write deadline before every write, so a -// stalled client cannot block the handler. gorilla's SetWriteDeadline only records -// the time and never fails, so there is no error to report here. -func armLogSocketWriteDeadline(conn *websocket.Conn) { - _ = conn.SetWriteDeadline(time.Now().Add(logSocketWriteWait)) -} - -// writeLogSocketJSON sends v as a JSON text message under the write deadline. -func writeLogSocketJSON(conn *websocket.Conn, v any) error { - armLogSocketWriteDeadline(conn) - return conn.WriteJSON(v) -} - -// writeLogSocketPing sends a ping frame under the same write deadline as data writes. -func writeLogSocketPing(conn *websocket.Conn) error { - armLogSocketWriteDeadline(conn) - return conn.WriteMessage(websocket.PingMessage, []byte{}) -} - // LogEntry represents a structured log entry sent over WebSocket. type LogEntry struct { Level string `json:"level"` @@ -127,9 +51,11 @@ func (h *LogsWSHandler) HandleWebSocket(c *gin.Context) { logger.Log().WithError(err).Error("Failed to upgrade WebSocket connection") return } + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() defer func() { - if err := conn.Close(); err != nil { - logger.Log().WithError(err).Error("Failed to close WebSocket connection") + if err := conn.Close(websocket.StatusNormalClosure, ""); err != nil { + logger.Log().WithError(err).Debug("Failed to close WebSocket connection") } }() @@ -165,23 +91,10 @@ func (h *LogsWSHandler) HandleWebSocket(c *gin.Context) { logChan := hook.Subscribe(subscriberID) defer hook.Unsubscribe(subscriberID) - // Channel to signal when client disconnects - done := make(chan struct{}) - - // Goroutine to read from WebSocket (detect client disconnect) - go func() { - defer close(done) - for { - if _, _, err := conn.ReadMessage(); err != nil { - return - } - } - }() + // Closes when the client disconnects or stops answering pings + done := watchLogSocket(ctx, conn) // Main loop: stream logs to client - ticker := time.NewTicker(logSocketPingPeriod) - defer ticker.Stop() - for { select { case entry, ok := <-logChan: @@ -224,12 +137,6 @@ func (h *LogsWSHandler) HandleWebSocket(c *gin.Context) { h.tracker.UpdateActivity(subscriberID) } - case <-ticker.C: - // Send ping to keep connection alive - if err := writeLogSocketPing(conn); err != nil { - return - } - case <-done: // Client disconnected return diff --git a/backend/internal/api/handlers/logs_ws_test.go b/backend/internal/api/handlers/logs_ws_test.go index 768b81832..c43ff5c59 100644 --- a/backend/internal/api/handlers/logs_ws_test.go +++ b/backend/internal/api/handlers/logs_ws_test.go @@ -1,6 +1,7 @@ package handlers import ( + "context" "encoding/json" "io" "net/http" @@ -11,8 +12,8 @@ import ( charonlogger "github.com/Wikid82/charon/backend/internal/logger" "github.com/Wikid82/charon/backend/internal/services" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -21,6 +22,17 @@ func toWebSocketURL(httpURL string) string { return "ws" + strings.TrimPrefix(httpURL, "http") } +// readWithin reads one message, giving up after timeout. A coder Read that hits its +// deadline closes the connection, so callers must not expect the connection to +// survive a timed-out read. +func readWithin(t *testing.T, conn *websocket.Conn, timeout time.Duration) ([]byte, error) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + _, data, err := conn.Read(ctx) + return data, err +} + func waitFor(t *testing.T, condition func() bool) { t.Helper() deadline := time.Now().Add(2 * time.Second) @@ -33,7 +45,7 @@ func waitFor(t *testing.T, condition func() bool) { t.Fatalf("condition not met within %s", 2*time.Second) } -func TestUpgraderCheckOrigin(t *testing.T) { +func TestIsAllowedWebSocketOrigin(t *testing.T) { t.Parallel() tests := []struct { @@ -64,7 +76,7 @@ func TestUpgraderCheckOrigin(t *testing.T) { if tc.xForwardedHost != "" { req.Header.Set("X-Forwarded-Host", tc.xForwardedHost) } - got := upgrader.CheckOrigin(req) + got := isAllowedWebSocketOrigin(req) assert.Equal(t, tc.want, got, "origin=%q host=%q xfh=%q", tc.origin, tc.host, tc.xForwardedHost) }) } @@ -96,10 +108,7 @@ func TestLogsWSHandler_StreamWithFiltersAndTracker(t *testing.T) { defer srv.Close() wsURL := toWebSocketURL(srv.URL) + "/logs?level=error&source=api" - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) - if resp != nil { - defer resp.Body.Close() - } + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) waitFor(t, func() bool { @@ -110,8 +119,7 @@ func TestLogsWSHandler_StreamWithFiltersAndTracker(t *testing.T) { charonlogger.WithFields(map[string]any{"source": "worker"}).Error("should-be-filtered-by-source") charonlogger.WithFields(map[string]any{"source": "api"}).Error("should-pass-filters") - require.NoError(t, conn.SetReadDeadline(time.Now().Add(3*time.Second))) - _, payload, err := conn.ReadMessage() + payload, err := readWithin(t, conn, 3*time.Second) require.NoError(t, err) var entry LogEntry @@ -123,7 +131,7 @@ func TestLogsWSHandler_StreamWithFiltersAndTracker(t *testing.T) { require.NotNil(t, entry.Fields) assert.Equal(t, "api", entry.Fields["source"]) - require.NoError(t, conn.Close()) + require.NoError(t, conn.CloseNow()) waitFor(t, func() bool { return tracker.GetCount() == 0 @@ -143,12 +151,9 @@ func TestLogsWSHandler_ConnectionListOmitsToken(t *testing.T) { defer srv.Close() wsURL := toWebSocketURL(srv.URL) + "/logs?level=error&source=api&token=" + wsTestToken - conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) - if resp != nil { - defer resp.Body.Close() - } + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() waitFor(t, func() bool { return tracker.GetCount() == 1 }) @@ -161,3 +166,21 @@ func TestLogsWSHandler_ConnectionListOmitsToken(t *testing.T) { assert.NotContains(t, string(body), wsTestToken) assert.Contains(t, string(body), "source=api") } + +// streamEndTimeout bounds how long readUntilStreamEnds waits for a stream to end. +const streamEndTimeout = 5 * time.Second + +// readUntilStreamEnds reads messages until a read fails and returns that error. +// Frames queued before the close, such as log lines, are read and discarded. It +// fails the test if the stream is still open after timeout. +func readUntilStreamEnds(t *testing.T, conn *websocket.Conn) error { + t.Helper() + deadline := time.Now().Add(streamEndTimeout) + for time.Now().Before(deadline) { + if _, err := readWithin(t, conn, time.Until(deadline)); err != nil { + return err + } + } + t.Fatalf("stream still open after %s", streamEndTimeout) + return nil +} diff --git a/backend/internal/api/handlers/stats_handler.go b/backend/internal/api/handlers/stats_handler.go index 23c52c6f0..a0e514654 100644 --- a/backend/internal/api/handlers/stats_handler.go +++ b/backend/internal/api/handlers/stats_handler.go @@ -179,13 +179,9 @@ func (h *StatsHandler) StatsWS(c *gin.Context) { if h.hub == nil { // Hub not wired — close immediately. - _ = conn.Close() + _ = conn.CloseNow() return } - client := &statsWSClient{ - conn: conn, - send: make(chan []byte, statsWSWriteBuffer), - } - h.hub.ServeClient(client) + h.hub.ServeClient(newStatsWSClient(conn)) } diff --git a/backend/internal/api/handlers/stats_ws_hub.go b/backend/internal/api/handlers/stats_ws_hub.go index 6d08a39df..7e45ee316 100644 --- a/backend/internal/api/handlers/stats_ws_hub.go +++ b/backend/internal/api/handlers/stats_ws_hub.go @@ -5,7 +5,7 @@ import ( "encoding/json" "sync" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/Wikid82/charon/backend/internal/logger" "github.com/Wikid82/charon/backend/internal/services" @@ -14,16 +14,32 @@ import ( const statsWSWriteBuffer = 64 // statsWSClient represents a single connected WebSocket client. +// +// ctx is cancelled when the client is closed. Cancelling it closes the +// connection, which unblocks any write that is stalled on a slow peer. type statsWSClient struct { conn *websocket.Conn send chan []byte + ctx context.Context + cancel context.CancelFunc closeOnce sync.Once } +// newStatsWSClient wraps conn for registration with a StatsWSHub. +func newStatsWSClient(conn *websocket.Conn) *statsWSClient { + ctx, cancel := context.WithCancel(context.Background()) + return &statsWSClient{ + conn: conn, + send: make(chan []byte, statsWSWriteBuffer), + ctx: ctx, + cancel: cancel, + } +} + func (c *statsWSClient) close() { c.closeOnce.Do(func() { close(c.send) - _ = c.conn.Close() + c.cancel() }) } @@ -105,7 +121,7 @@ func (h *StatsWSHub) Run(ctx context.Context) { } } -// ServeClient upgrades c to WebSocket and registers it with the hub. +// ServeClient registers an upgraded client with the hub. // It blocks until the client disconnects. Call from a Gin handler. func (h *StatsWSHub) ServeClient(c *statsWSClient) { h.register <- c @@ -113,16 +129,16 @@ func (h *StatsWSHub) ServeClient(c *statsWSClient) { // Writer goroutine: forwards hub messages to the WebSocket connection. go func() { for payload := range c.send { - if err := c.conn.WriteMessage(websocket.TextMessage, payload); err != nil { + if err := c.conn.Write(c.ctx, websocket.MessageText, payload); err != nil { break } } - _ = c.conn.Close() + _ = c.conn.Close(websocket.StatusNormalClosure, "") }() // Reader loop: detects client disconnect. for { - if _, _, err := c.conn.ReadMessage(); err != nil { + if _, _, err := c.conn.Read(c.ctx); err != nil { break } } diff --git a/backend/internal/api/handlers/stats_ws_hub_test.go b/backend/internal/api/handlers/stats_ws_hub_test.go index b5bbcfcc7..a345986f1 100644 --- a/backend/internal/api/handlers/stats_ws_hub_test.go +++ b/backend/internal/api/handlers/stats_ws_hub_test.go @@ -9,8 +9,8 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -25,15 +25,11 @@ func newHubServer(t *testing.T, hub *StatsWSHub) *httptest.Server { gin.SetMode(gin.TestMode) r := gin.New() r.GET("/ws", func(c *gin.Context) { - conn, err := upgrader.Upgrade(c.Writer, c.Request, nil) // nosemgrep: go.gorilla.security.audit.websocket-missing-origin-check.websocket-missing-origin-check + conn, err := websocket.Accept(c.Writer, c.Request, nil) if err != nil { return } - client := &statsWSClient{ - conn: conn, - send: make(chan []byte, statsWSWriteBuffer), - } - hub.ServeClient(client) + hub.ServeClient(newStatsWSClient(conn)) }) return httptest.NewServer(r) } @@ -42,7 +38,7 @@ func newHubServer(t *testing.T, hub *StatsWSHub) *httptest.Server { func dialHub(t *testing.T, srv *httptest.Server) *websocket.Conn { t.Helper() wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose require.NoError(t, err) return conn } @@ -120,7 +116,7 @@ func TestStatsWSHub_BroadcastToConnectedClient(t *testing.T) { defer srv.Close() conn := dialHub(t, srv) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) // Allow the client to register. time.Sleep(50 * time.Millisecond) @@ -132,8 +128,7 @@ func TestStatsWSHub_BroadcastToConnectedClient(t *testing.T) { } hub.Broadcast(services.StatsPushMessage{Type: "stats_update", Data: summary}) - require.NoError(t, conn.SetReadDeadline(time.Now().Add(3*time.Second))) - _, data, err := conn.ReadMessage() + data, err := readWithin(t, conn, 3*time.Second) require.NoError(t, err) var got services.StatsPushMessage @@ -155,7 +150,7 @@ func TestStatsWSHub_Run_ClosesClientsOnContextCancel(t *testing.T) { defer srv.Close() conn := dialHub(t, srv) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) // Allow the client to register. time.Sleep(50 * time.Millisecond) @@ -163,8 +158,7 @@ func TestStatsWSHub_Run_ClosesClientsOnContextCancel(t *testing.T) { // Cancel hub — clients should be closed. cancel() - require.NoError(t, conn.SetReadDeadline(time.Now().Add(3*time.Second))) - _, _, err := conn.ReadMessage() + _, err := readWithin(t, conn, 3*time.Second) // Expect an error (close frame or EOF) after hub cancellation. assert.Error(t, err, "expected connection to be closed after hub cancel") } @@ -182,7 +176,7 @@ func TestStatsWSHub_Run_SlowClientDropsMessage(t *testing.T) { defer srv.Close() conn := dialHub(t, srv) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) // Allow the client to register. time.Sleep(50 * time.Millisecond) @@ -220,7 +214,7 @@ func TestStatsWSHub_StatsWSHandler_NonWebSocketRequest(t *testing.T) { srv := httptest.NewServer(r) defer srv.Close() - // Plain HTTP GET — no WebSocket upgrade headers — upgrader.Upgrade will fail. + // Plain HTTP GET — no WebSocket upgrade headers — websocket.Accept will fail. resp, err := http.Get(srv.URL + "/ws") //nolint:noctx require.NoError(t, err) _ = resp.Body.Close() @@ -245,13 +239,12 @@ func TestStatsWSHub_StatsWSHandler_NilHub(t *testing.T) { defer srv.Close() wsURL := "ws" + strings.TrimPrefix(srv.URL, "http") + "/ws" - conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) //nolint:bodyclose + conn, _, err := websocket.Dial(context.Background(), wsURL, nil) //nolint:bodyclose require.NoError(t, err) - t.Cleanup(func() { _ = conn.Close() }) + t.Cleanup(func() { _ = conn.CloseNow() }) // Hub is nil — handler closes the conn immediately after upgrade. - require.NoError(t, conn.SetReadDeadline(time.Now().Add(3*time.Second))) - _, _, err = conn.ReadMessage() + _, err = readWithin(t, conn, 3*time.Second) assert.Error(t, err, "expected immediate connection close when hub is nil") } @@ -279,7 +272,7 @@ func TestStatsWSHub_UnregisterRemovesClient(t *testing.T) { // Close the client-side connection — ServeClient's read loop will detect the close // and send to hub.unregister. - _ = conn.Close() + _ = conn.CloseNow() require.Eventually(t, func() bool { hub.mu.RLock() diff --git a/backend/internal/api/handlers/websocket_upgrade.go b/backend/internal/api/handlers/websocket_upgrade.go new file mode 100644 index 000000000..90c57fe76 --- /dev/null +++ b/backend/internal/api/handlers/websocket_upgrade.go @@ -0,0 +1,56 @@ +package handlers + +import ( + "errors" + "fmt" + "net/http" + "net/url" + + "github.com/coder/websocket" + "github.com/gin-gonic/gin" +) + +// errWebSocketOriginDenied is returned by upgradeWebSocket when the request +// Origin does not match the request host. +var errWebSocketOriginDenied = errors.New("websocket: request origin not allowed") + +// isAllowedWebSocketOrigin reports whether a browser WebSocket request may +// upgrade. Requests without an Origin header are non-browser or same-origin +// clients and are allowed. Otherwise the Origin host must equal the request +// host, or X-Forwarded-Host when a reverse proxy sets it. +func isAllowedWebSocketOrigin(r *http.Request) bool { + origin := r.Header.Get("Origin") + if origin == "" { + return true + } + originURL, err := url.Parse(origin) + if err != nil { + return false + } + requestHost := r.Host + if forwardedHost := r.Header.Get("X-Forwarded-Host"); forwardedHost != "" { + requestHost = forwardedHost + } + return originURL.Host == requestHost +} + +// upgradeWebSocket upgrades an HTTP connection to WebSocket after the origin +// check. All WS handlers in this package MUST go through this helper rather +// than calling websocket.Accept directly, so the origin check stays visible to +// single-file static analysis at every call site. +// +// The origin check is done here rather than by websocket.Accept, which only +// knows about the request Host and cannot honour X-Forwarded-Host. Accept's own +// check is therefore disabled with InsecureSkipVerify, and this function +// rejects disallowed origins with 403 before Accept runs. +func upgradeWebSocket(c *gin.Context) (*websocket.Conn, error) { + if !isAllowedWebSocketOrigin(c.Request) { + http.Error(c.Writer, http.StatusText(http.StatusForbidden), http.StatusForbidden) + return nil, errWebSocketOriginDenied + } + conn, err := websocket.Accept(c.Writer, c.Request, &websocket.AcceptOptions{InsecureSkipVerify: true}) + if err != nil { + return nil, fmt.Errorf("upgrade websocket: %w", err) + } + return conn, nil +} diff --git a/backend/internal/api/handlers/websocket_upgrade_test.go b/backend/internal/api/handlers/websocket_upgrade_test.go new file mode 100644 index 000000000..6d9659aa7 --- /dev/null +++ b/backend/internal/api/handlers/websocket_upgrade_test.go @@ -0,0 +1,92 @@ +package handlers + +import ( + "context" + "errors" + "net/http" + "net/http/httptest" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// newUpgradeTestServer serves upgradeWebSocket and closes each accepted +// connection once the test ends. It reports the upgrade error through upgradeErr. +func newUpgradeTestServer(t *testing.T, upgradeErr chan<- error) *httptest.Server { + t.Helper() + gin.SetMode(gin.TestMode) + router := gin.New() + router.GET("/ws", func(c *gin.Context) { + conn, err := upgradeWebSocket(c) + upgradeErr <- err + if err != nil { + return + } + defer func() { _ = conn.CloseNow() }() + _, _, _ = conn.Read(c.Request.Context()) + }) + srv := httptest.NewServer(router) + t.Cleanup(srv.Close) + return srv +} + +func TestUpgradeWebSocket_RejectsDisallowedOrigin(t *testing.T) { + upgradeErr := make(chan error, 1) + srv := newUpgradeTestServer(t, upgradeErr) + + header := http.Header{"Origin": []string{"http://evil.example"}} + conn, resp, err := websocket.Dial(context.Background(), toWebSocketURL(srv.URL)+"/ws", &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body + if conn != nil { + _ = conn.CloseNow() + } + require.Error(t, err, "a cross-origin upgrade must not succeed") + require.NotNil(t, resp) + assert.Equal(t, http.StatusForbidden, resp.StatusCode) + + select { + case got := <-upgradeErr: + assert.True(t, errors.Is(got, errWebSocketOriginDenied), "got %v", got) + case <-time.After(5 * time.Second): + t.Fatal("handler did not run") + } +} + +func TestUpgradeWebSocket_AcceptsSameOriginRequest(t *testing.T) { + upgradeErr := make(chan error, 1) + srv := newUpgradeTestServer(t, upgradeErr) + + // The Origin host matches the request host, so the upgrade is allowed. + header := http.Header{"Origin": []string{srv.URL}} + conn, _, err := websocket.Dial(context.Background(), toWebSocketURL(srv.URL)+"/ws", &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body + require.NoError(t, err) + t.Cleanup(func() { _ = conn.CloseNow() }) + + select { + case got := <-upgradeErr: + require.NoError(t, got) + case <-time.After(5 * time.Second): + t.Fatal("handler did not run") + } +} + +func TestUpgradeWebSocket_NonWebSocketRequestFails(t *testing.T) { + upgradeErr := make(chan error, 1) + srv := newUpgradeTestServer(t, upgradeErr) + + // A plain GET without upgrade headers passes the origin check but fails the handshake. + resp, err := http.Get(srv.URL + "/ws") //nolint:noctx // test-only request against httptest server + require.NoError(t, err) + _ = resp.Body.Close() + + select { + case got := <-upgradeErr: + require.Error(t, got) + assert.False(t, errors.Is(got, errWebSocketOriginDenied)) + case <-time.After(5 * time.Second): + t.Fatal("handler did not run") + } +} diff --git a/backend/internal/orthrus/server.go b/backend/internal/orthrus/server.go index 6e6f46544..1b3ffd0d0 100644 --- a/backend/internal/orthrus/server.go +++ b/backend/internal/orthrus/server.go @@ -12,18 +12,13 @@ import ( "github.com/Wikid82/charon/backend/internal/logger" "github.com/Wikid82/charon/backend/internal/models" "github.com/Wikid82/charon/backend/internal/util" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - "github.com/gorilla/websocket" "github.com/sirupsen/logrus" "golang.org/x/crypto/bcrypt" "gorm.io/gorm" ) -var wsUpgrader = websocket.Upgrader{ - // Allow all origins; the real trust boundary is the bearer token. - CheckOrigin: func(r *http.Request) bool { return true }, -} - // OrthrusServer manages incoming agent WebSocket connections and tracks // active sessions indexed by agent UUID. type OrthrusServer struct { @@ -103,10 +98,12 @@ func (s *OrthrusServer) HandleWebSocket(c *gin.Context) { // constant on top of the two (streamTypeDocker, streamTypePortForward) // that already have to be kept in sync between this package and // agent/leash/leash.go. - respHeader := http.Header{} - respHeader.Set("X-Orthrus-Write-Enabled", strconv.FormatBool(agent.WriteEnabled)) + // websocket.Accept writes the handshake response from the writer's headers, so + // the flag is set on the header before the upgrade. + c.Writer.Header().Set("X-Orthrus-Write-Enabled", strconv.FormatBool(agent.WriteEnabled)) - conn, err := wsUpgrader.Upgrade(c.Writer, c.Request, respHeader) + // Origin is not checked: the real trust boundary is the bearer token above. + conn, err := websocket.Accept(c.Writer, c.Request, &websocket.AcceptOptions{InsecureSkipVerify: true}) if err != nil { logger.Log().WithError(err).Error("orthrus: WebSocket upgrade failed") return @@ -115,7 +112,7 @@ func (s *OrthrusServer) HandleWebSocket(c *gin.Context) { session, err := NewAgentSession(agent.UUID, agent.Name, agent.WriteEnabled, s.auditLogger, conn) if err != nil { logger.Log().WithError(err).Error("orthrus: create agent session failed") - _ = conn.Close() + _ = conn.CloseNow() return } diff --git a/backend/internal/orthrus/server_coverage_test.go b/backend/internal/orthrus/server_coverage_test.go index 386d4bd24..790aa752d 100644 --- a/backend/internal/orthrus/server_coverage_test.go +++ b/backend/internal/orthrus/server_coverage_test.go @@ -1,6 +1,7 @@ package orthrus import ( + "context" "net" "net/http" "net/http/httptest" @@ -8,8 +9,8 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - gorillaws "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/crypto/bcrypt" @@ -254,12 +255,9 @@ func TestOrthrusServer_HandleWebSocket_ValidToken_UpgradesConnection(t *testing. wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/ws" header := http.Header{"Authorization": []string{"Bearer " + token}} - conn, dialResp, err := gorillaws.DefaultDialer.Dial(wsURL, header) + conn, _, err := websocket.Dial(context.Background(), wsURL, &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() var ok bool for i := 0; i < 20; i++ { @@ -281,7 +279,7 @@ func TestOrthrusServer_FindAgentByToken_DBError_ReturnsError(t *testing.T) { } // TestOrthrusServer_HandleWebSocket_UpgradeFailure covers server.go:82-85 — -// the error-log path when wsUpgrader.Upgrade rejects a non-WebSocket request. +// the error-log path when websocket.Accept rejects a non-WebSocket request. func TestOrthrusServer_HandleWebSocket_UpgradeFailure(t *testing.T) { gin.SetMode(gin.TestMode) db := setupServerTestDB(t) @@ -301,7 +299,7 @@ func TestOrthrusServer_HandleWebSocket_UpgradeFailure(t *testing.T) { require.NoError(t, db.Create(agent).Error) // Plain HTTP GET with valid auth — NOT a WebSocket upgrade. - // wsUpgrader.Upgrade writes 400 Bad Request and returns an error. + // websocket.Accept rejects the missing Upgrade header with 426 Upgrade Required. w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) req := httptest.NewRequest(http.MethodGet, "/ws", http.NoBody) @@ -310,7 +308,7 @@ func TestOrthrusServer_HandleWebSocket_UpgradeFailure(t *testing.T) { // Should not panic; handler logs the error and returns. srv.HandleWebSocket(c) - assert.Equal(t, http.StatusBadRequest, w.Code) + assert.Equal(t, http.StatusUpgradeRequired, w.Code) } // TestOrthrusServer_HandleWebSocket_ExternalProxyFails covers server.go:100-104 — @@ -349,12 +347,9 @@ func TestOrthrusServer_HandleWebSocket_ExternalProxyFails(t *testing.T) { wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/ws" header := http.Header{"Authorization": []string{"Bearer " + token}} - conn, dialResp, err := gorillaws.DefaultDialer.Dial(wsURL, header) + conn, _, err := websocket.Dial(context.Background(), wsURL, &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() // sessions.Store is called after StartExternalProxy, so once the session is // present the proxy attempt (and its failure) has already completed. @@ -407,12 +402,9 @@ func TestHandleWebSocket_DisplacesExistingSession(t *testing.T) { wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/ws" header := http.Header{"Authorization": []string{"Bearer " + token}} - conn, dialResp, err := gorillaws.DefaultDialer.Dial(wsURL, header) + conn, _, err := websocket.Dial(context.Background(), wsURL, &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() // Wait for the new session to be stored, which means HandleWebSocket has run // past the displacement block and stored the replacement session. diff --git a/backend/internal/orthrus/server_external_proxy_test.go b/backend/internal/orthrus/server_external_proxy_test.go index c10dc2854..bd05dd074 100644 --- a/backend/internal/orthrus/server_external_proxy_test.go +++ b/backend/internal/orthrus/server_external_proxy_test.go @@ -1,6 +1,7 @@ package orthrus import ( + "context" "fmt" "net" "net/http" @@ -9,8 +10,8 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - gorillaws "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/crypto/bcrypt" @@ -52,12 +53,9 @@ func TestOrthrusServer_HandleWebSocket_StartsExternalProxy(t *testing.T) { wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/ws" header := http.Header{"Authorization": []string{"Bearer " + token}} - conn, dialResp, err := gorillaws.DefaultDialer.Dial(wsURL, header) + conn, _, err := websocket.Dial(context.Background(), wsURL, &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() assert.Eventually(t, func() bool { status, ok := srv.GetExternalProxyStatus("ext-srv01-uuid") diff --git a/backend/internal/orthrus/server_proxy_test.go b/backend/internal/orthrus/server_proxy_test.go index bfca9a5ad..1fbdb6925 100644 --- a/backend/internal/orthrus/server_proxy_test.go +++ b/backend/internal/orthrus/server_proxy_test.go @@ -1,6 +1,7 @@ package orthrus import ( + "context" "net" "net/http" "net/http/httptest" @@ -8,8 +9,8 @@ import ( "testing" "time" + "github.com/coder/websocket" "github.com/gin-gonic/gin" - gorillaws "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/crypto/bcrypt" @@ -47,12 +48,9 @@ func TestOrthrusServer_HandleWebSocket_StartsDockerProxy(t *testing.T) { wsURL := "ws" + strings.TrimPrefix(ts.URL, "http") + "/ws" header := http.Header{"Authorization": []string{"Bearer " + token}} - conn, dialResp, err := gorillaws.DefaultDialer.Dial(wsURL, header) + conn, _, err := websocket.Dial(context.Background(), wsURL, &websocket.DialOptions{HTTPHeader: header}) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } - defer func() { _ = conn.Close() }() + defer func() { _ = conn.CloseNow() }() assert.Eventually(t, func() bool { addr, ok := srv.GetProxyAddr("proxy-s1-uuid") diff --git a/backend/internal/orthrus/session.go b/backend/internal/orthrus/session.go index f4bb1fd55..fb13da242 100644 --- a/backend/internal/orthrus/session.go +++ b/backend/internal/orthrus/session.go @@ -13,7 +13,7 @@ import ( "time" "github.com/Wikid82/charon/backend/internal/logger" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/hashicorp/yamux" "golang.org/x/time/rate" ) @@ -36,73 +36,11 @@ const ( // Must match the constant in the Orthrus agent (agent/leash/leash.go). const streamTypeDocker = byte(0x01) -// wsNetConn adapts a gorilla WebSocket connection to the net.Conn interface -// so that yamux can use it as a byte-stream transport. -// gorilla websocket supports one concurrent reader and one concurrent writer; -// rMu serialises reads and wMu serialises writes accordingly. -type wsNetConn struct { - conn *websocket.Conn - rMu sync.Mutex - wMu sync.Mutex - reader io.Reader -} - -func newWSNetConn(conn *websocket.Conn) net.Conn { - return &wsNetConn{conn: conn} -} - -func (c *wsNetConn) Read(b []byte) (int, error) { - c.rMu.Lock() - defer c.rMu.Unlock() - - for { - if c.reader != nil { - n, err := c.reader.Read(b) - if err == io.EOF { - c.reader = nil - continue - } - return n, err - } - - _, msg, err := c.conn.NextReader() - if err != nil { - return 0, err - } - c.reader = msg - } -} - -func (c *wsNetConn) Write(b []byte) (int, error) { - c.wMu.Lock() - defer c.wMu.Unlock() - - if err := c.conn.WriteMessage(websocket.BinaryMessage, b); err != nil { - return 0, err - } - return len(b), nil -} - -func (c *wsNetConn) Close() error { - return c.conn.Close() -} - -func (c *wsNetConn) LocalAddr() net.Addr { return c.conn.NetConn().LocalAddr() } -func (c *wsNetConn) RemoteAddr() net.Addr { return c.conn.NetConn().RemoteAddr() } - -func (c *wsNetConn) SetDeadline(t time.Time) error { - if err := c.conn.SetReadDeadline(t); err != nil { - return err - } - return c.conn.SetWriteDeadline(t) -} - -func (c *wsNetConn) SetReadDeadline(t time.Time) error { - return c.conn.SetReadDeadline(t) -} - -func (c *wsNetConn) SetWriteDeadline(t time.Time) error { - return c.conn.SetWriteDeadline(t) +// newWSNetConn adapts a WebSocket connection to net.Conn so that yamux can use +// it as a byte-stream transport. Each Write sends one binary message. ctx bounds +// the adapter's reads and writes: cancelling it closes the WebSocket. +func newWSNetConn(ctx context.Context, conn *websocket.Conn) net.Conn { + return websocket.NetConn(ctx, conn, websocket.MessageBinary) } // ExternalProxyStatus describes the runtime state of the external Docker proxy @@ -125,7 +63,6 @@ type ExternalProxyStatus struct { type AgentSession struct { agentUUID string agentName string - conn *websocket.Conn session *yamux.Session cancel context.CancelFunc proxyPort int // ephemeral port; 0 until StartDockerProxy succeeds @@ -149,17 +86,18 @@ type AgentSession struct { } // NewAgentSession wraps the WebSocket connection in a Yamux server session. +// The session owns conn: Close tears the WebSocket down. func NewAgentSession(agentUUID, agentName string, writeEnabled bool, auditLogger AuditLogger, conn *websocket.Conn) (*AgentSession, error) { cfg := yamux.DefaultConfig() cfg.LogOutput = io.Discard - session, err := yamux.Server(newWSNetConn(conn), cfg) + ctx, cancel := context.WithCancel(context.Background()) + session, err := yamux.Server(newWSNetConn(ctx, conn), cfg) if err != nil { + cancel() return nil, fmt.Errorf("orthrus: create yamux session: %w", err) } - _, cancel := context.WithCancel(context.Background()) - var writeLimiter *rate.Limiter if writeEnabled { writeLimiter = rate.NewLimiter(rate.Every(writeLimiterEvery), writeLimiterBurst) @@ -168,7 +106,6 @@ func NewAgentSession(agentUUID, agentName string, writeEnabled bool, auditLogger return &AgentSession{ agentUUID: agentUUID, agentName: agentName, - conn: conn, session: session, writeLimiter: writeLimiter, auditLogger: auditLogger, diff --git a/backend/internal/orthrus/session_coverage_test.go b/backend/internal/orthrus/session_coverage_test.go index d267d1b7a..6c432f043 100644 --- a/backend/internal/orthrus/session_coverage_test.go +++ b/backend/internal/orthrus/session_coverage_test.go @@ -1,71 +1,65 @@ package orthrus import ( + "context" "errors" "net" "testing" "time" + "github.com/coder/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -func TestWSNetConn_Write_Success(t *testing.T) { - serverConn, done := testWSPair(t) +// TestNewWSNetConn_WriteSendsOneBinaryMessage checks that each Write on the +// adapter reaches the peer as one binary WebSocket message. +func TestNewWSNetConn_WriteSendsOneBinaryMessage(t *testing.T) { + serverConn, clientConn, done := testWSPairBoth(t) defer done() - nc := newWSNetConn(serverConn) + nc := newWSNetConn(context.Background(), serverConn) n, err := nc.Write([]byte("hello")) - assert.NoError(t, err) + require.NoError(t, err) assert.Equal(t, 5, n) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + typ, data, err := clientConn.Read(ctx) + require.NoError(t, err) + assert.Equal(t, websocket.MessageBinary, typ) + assert.Equal(t, []byte("hello"), data) } -func TestWSNetConn_LocalAddr_NotNil(t *testing.T) { +func TestNewWSNetConn_LocalAddr_NotNil(t *testing.T) { serverConn, done := testWSPair(t) defer done() - nc := newWSNetConn(serverConn) + nc := newWSNetConn(context.Background(), serverConn) addr := nc.LocalAddr() assert.NotNil(t, addr) assert.Implements(t, (*net.Addr)(nil), addr) } -func TestWSNetConn_RemoteAddr_NotNil(t *testing.T) { +func TestNewWSNetConn_RemoteAddr_NotNil(t *testing.T) { serverConn, done := testWSPair(t) defer done() - nc := newWSNetConn(serverConn) + nc := newWSNetConn(context.Background(), serverConn) addr := nc.RemoteAddr() assert.NotNil(t, addr) assert.Implements(t, (*net.Addr)(nil), addr) } -func TestWSNetConn_SetDeadline_NoError(t *testing.T) { +func TestNewWSNetConn_Deadlines_NoError(t *testing.T) { serverConn, done := testWSPair(t) defer done() - nc := newWSNetConn(serverConn) + nc := newWSNetConn(context.Background(), serverConn) // Zero time clears the deadline — should succeed on a live connection. - err := nc.SetDeadline(time.Time{}) - assert.NoError(t, err) -} - -func TestWSNetConn_SetReadDeadline_NoError(t *testing.T) { - serverConn, done := testWSPair(t) - defer done() - - nc := newWSNetConn(serverConn) - err := nc.SetReadDeadline(time.Time{}) - assert.NoError(t, err) -} - -func TestWSNetConn_SetWriteDeadline_NoError(t *testing.T) { - serverConn, done := testWSPair(t) - defer done() - - nc := newWSNetConn(serverConn) - err := nc.SetWriteDeadline(time.Time{}) - assert.NoError(t, err) + assert.NoError(t, nc.SetDeadline(time.Time{})) + assert.NoError(t, nc.SetReadDeadline(time.Time{})) + assert.NoError(t, nc.SetWriteDeadline(time.Time{})) } func TestAgentSession_GetProxyAddr_WithPort(t *testing.T) { @@ -84,21 +78,6 @@ func TestAgentSession_GetProxyAddr_WithPort(t *testing.T) { assert.Equal(t, "127.0.0.1:8080", addr) } -// TestWSNetConn_SetDeadline_ClosedConn_ReturnsError covers the error-return path -// inside SetDeadline when the underlying SetReadDeadline call fails on a closed conn. -func TestWSNetConn_SetDeadline_ClosedConn_ReturnsError(t *testing.T) { - serverConn, done := testWSPair(t) - defer done() - - nc := newWSNetConn(serverConn) - // Closing the underlying WebSocket connection invalidates the TCP socket; - // the subsequent SetReadDeadline call inside SetDeadline will return an error. - _ = serverConn.Close() - - err := nc.SetDeadline(time.Now().Add(time.Second)) - assert.Error(t, err) -} - // TestGetExternalProxyStatus_ErrorFieldPopulated covers session.go:336-338 — // the errStr assignment when extErr is non-nil. func TestGetExternalProxyStatus_ErrorFieldPopulated(t *testing.T) { diff --git a/backend/internal/orthrus/session_external_proxy_test.go b/backend/internal/orthrus/session_external_proxy_test.go index febb02fba..d19fb56a9 100644 --- a/backend/internal/orthrus/session_external_proxy_test.go +++ b/backend/internal/orthrus/session_external_proxy_test.go @@ -2,6 +2,7 @@ package orthrus import ( "bufio" + "context" "fmt" "io" "net" @@ -35,7 +36,7 @@ func sessionWithLoopback(t *testing.T) (sess *AgentSession, clientYamux *yamux.S cfg := yamux.DefaultConfig() cfg.LogOutput = io.Discard - cy, err := yamux.Client(newWSNetConn(clientConn), cfg) + cy, err := yamux.Client(newWSNetConn(context.Background(), clientConn), cfg) require.NoError(t, err) s, err := NewAgentSession("ext-test-uuid", "ext-test-agent", false, nil, serverConn) diff --git a/backend/internal/orthrus/session_proxy_test.go b/backend/internal/orthrus/session_proxy_test.go index ce30972bf..1d3db2bf2 100644 --- a/backend/internal/orthrus/session_proxy_test.go +++ b/backend/internal/orthrus/session_proxy_test.go @@ -1,6 +1,7 @@ package orthrus import ( + "context" "io" "net" "net/http" @@ -9,7 +10,7 @@ import ( "testing" "time" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/hashicorp/yamux" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -19,11 +20,10 @@ import ( func testWSPairBoth(t *testing.T) (serverConn, clientConn *websocket.Conn, done func()) { t.Helper() - upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} ready := make(chan *websocket.Conn, 1) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true}) if err != nil { return } @@ -31,15 +31,12 @@ func testWSPairBoth(t *testing.T) (serverConn, clientConn *websocket.Conn, done })) url := "ws" + strings.TrimPrefix(srv.URL, "http") - client, dialResp, err := websocket.DefaultDialer.Dial(url, nil) + client, _, err := websocket.Dial(context.Background(), url, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } srvConn := <-ready return srvConn, client, func() { - _ = client.Close() + _ = client.CloseNow() srv.Close() } } @@ -91,7 +88,7 @@ func TestStartDockerProxy_AcceptsAndForwards(t *testing.T) { // Build a yamux client from the client websocket conn. cfg := yamux.DefaultConfig() cfg.LogOutput = io.Discard - clientYamux, err := yamux.Client(newWSNetConn(clientConn), cfg) + clientYamux, err := yamux.Client(newWSNetConn(context.Background(), clientConn), cfg) require.NoError(t, err) defer func() { _ = clientYamux.Close() }() diff --git a/backend/internal/orthrus/session_test.go b/backend/internal/orthrus/session_test.go index 2829f37f5..e76f159a0 100644 --- a/backend/internal/orthrus/session_test.go +++ b/backend/internal/orthrus/session_test.go @@ -1,6 +1,7 @@ package orthrus import ( + "context" "fmt" "io" "net" @@ -11,7 +12,7 @@ import ( "testing" "time" - "github.com/gorilla/websocket" + "github.com/coder/websocket" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) @@ -21,11 +22,10 @@ import ( func testWSPair(t *testing.T) (serverConn *websocket.Conn, done func()) { t.Helper() - upgrader := websocket.Upgrader{CheckOrigin: func(*http.Request) bool { return true }} ready := make(chan *websocket.Conn, 1) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - conn, err := upgrader.Upgrade(w, r, nil) + conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{InsecureSkipVerify: true}) if err != nil { return } @@ -33,15 +33,12 @@ func testWSPair(t *testing.T) (serverConn *websocket.Conn, done func()) { })) url := "ws" + strings.TrimPrefix(srv.URL, "http") - clientConn, dialResp, err := websocket.DefaultDialer.Dial(url, nil) + clientConn, _, err := websocket.Dial(context.Background(), url, nil) //nolint:bodyclose // coder/websocket manages the handshake response body require.NoError(t, err) - if dialResp != nil { - _ = dialResp.Body.Close() - } serverConn = <-ready return serverConn, func() { - _ = clientConn.Close() + _ = clientConn.CloseNow() srv.Close() } } diff --git a/playwright.config.js b/playwright.config.js index 957edf456..9f49957fa 100644 --- a/playwright.config.js +++ b/playwright.config.js @@ -70,6 +70,21 @@ const BROWSER_LAUNCH_OPTIONS = { webkit: {}, }; +/** + * Engine for the `security-tests` project. CI keeps Chromium for it, exactly as before. A local run uses Firefox, + * the single-browser middle ground between Chromium and WebKit, because CI already covers the others. + * PLAYWRIGHT_SECURITY_BROWSER overrides either default (chromium, firefox or webkit). + */ +const SECURITY_DEVICES = { + chromium: devices['Desktop Chrome'], + firefox: devices['Desktop Firefox'], + webkit: devices['Desktop Safari'], +}; +const securityBrowser = process.env.PLAYWRIGHT_SECURITY_BROWSER ?? (process.env.CI ? 'chromium' : 'firefox'); +if (!SECURITY_DEVICES[securityBrowser]) { + throw new Error(`PLAYWRIGHT_SECURITY_BROWSER must be chromium, firefox or webkit, got "${securityBrowser}"`); +} + const coverageReporterConfig = enableCoverage ? defineCoverageReporterConfig({ sourceRoot: __dirname, exclude: [ @@ -276,7 +291,7 @@ export default defineConfig({ dependencies: ['setup'], }, - // Security Tests - Run WITH security enabled (SEQUENTIAL, Chromium only) + // Security Tests - Run WITH security enabled (SEQUENTIAL, one engine: Chromium in CI, Firefox locally) { name: 'security-tests', testDir: './tests', @@ -289,7 +304,8 @@ export default defineConfig({ fullyParallel: false, workers: 1, use: { - ...devices['Desktop Chrome'], + ...SECURITY_DEVICES[securityBrowser], + launchOptions: BROWSER_LAUNCH_OPTIONS[securityBrowser], headless: true, storageState: STORAGE_STATE, }, diff --git a/renovate.json b/renovate.json index 50318db37..af97399f9 100644 --- a/renovate.json +++ b/renovate.json @@ -29,7 +29,6 @@ ], "ignoreDeps": [ - "github.com/gorilla/websocket", "github.com/golang-jwt/jwt/v5", "@types/eslint-plugin-jsx-a11y", "google.golang.org/grpc" diff --git a/tests/security-enforcement/acl-creation.spec.ts b/tests/security-enforcement/acl-creation.spec.ts index 3ac0b5b6d..27aa10410 100644 --- a/tests/security-enforcement/acl-creation.spec.ts +++ b/tests/security-enforcement/acl-creation.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from '@playwright/test'; +import { deleteAccessListViaAPI } from '../utils/api-helpers'; const TEST_EMAIL = process.env.E2E_TEST_EMAIL || 'e2e-test@example.com'; const TEST_PASSWORD = process.env.E2E_TEST_PASSWORD || 'TestPassword123!'; @@ -23,61 +24,69 @@ test.describe('ACL Creation Baseline', () => { const unique = Date.now(); const aclName = `ACL Baseline ${unique}`; const profileName = `Headers Baseline ${unique}`; + let aclUuid: string | undefined; - await test.step('Create ACL baseline entry', async () => { - const aclResponse = await request.post('/api/v1/access-lists', { - headers: { - Authorization: `Bearer ${token}`, - }, - data: { - name: aclName, - type: 'whitelist', - enabled: true, - ip_rules: JSON.stringify([ - { - cidr: '127.0.0.1/32', - description: 'Local test runner', - }, - ]), - }, - }); - - expect(aclResponse.ok()).toBeTruthy(); - }); + try { + await test.step('Create ACL baseline entry', async () => { + const aclResponse = await request.post('/api/v1/access-lists', { + headers: { + Authorization: `Bearer ${token}`, + }, + data: { + name: aclName, + type: 'whitelist', + enabled: true, + ip_rules: JSON.stringify([ + { + cidr: '127.0.0.1/32', + description: 'Local test runner', + }, + ]), + }, + }); - await test.step('Create security headers profile baseline entry', async () => { - const profileResponse = await request.post('/api/v1/security/headers/profiles', { - headers: { - Authorization: `Bearer ${token}`, - }, - data: { - name: profileName, - }, + expect(aclResponse.ok()).toBeTruthy(); + aclUuid = (await aclResponse.json()).uuid; }); - expect(profileResponse.status()).toBe(201); - }); + await test.step('Create security headers profile baseline entry', async () => { + const profileResponse = await request.post('/api/v1/security/headers/profiles', { + headers: { + Authorization: `Bearer ${token}`, + }, + data: { + name: profileName, + }, + }); - await test.step('Verify baseline entries are queryable', async () => { - const aclListResponse = await request.get('/api/v1/access-lists', { - headers: { - Authorization: `Bearer ${token}`, - }, + expect(profileResponse.status()).toBe(201); }); - expect(aclListResponse.ok()).toBeTruthy(); - const aclList = await aclListResponse.json(); - expect(Array.isArray(aclList)).toBeTruthy(); - expect(aclList.some((item: { name?: string }) => item.name === aclName)).toBeTruthy(); - const profileListResponse = await request.get('/api/v1/security/headers/profiles', { - headers: { - Authorization: `Bearer ${token}`, - }, + await test.step('Verify baseline entries are queryable', async () => { + const aclListResponse = await request.get('/api/v1/access-lists', { + headers: { + Authorization: `Bearer ${token}`, + }, + }); + expect(aclListResponse.ok()).toBeTruthy(); + const aclList = await aclListResponse.json(); + expect(Array.isArray(aclList)).toBeTruthy(); + expect(aclList.some((item: { name?: string }) => item.name === aclName)).toBeTruthy(); + + const profileListResponse = await request.get('/api/v1/security/headers/profiles', { + headers: { + Authorization: `Bearer ${token}`, + }, + }); + expect(profileListResponse.ok()).toBeTruthy(); + const profilePayload = await profileListResponse.json(); + const profiles = Array.isArray(profilePayload?.profiles) ? profilePayload.profiles : []; + expect(profiles.some((item: { name?: string }) => item.name === profileName)).toBeTruthy(); }); - expect(profileListResponse.ok()).toBeTruthy(); - const profilePayload = await profileListResponse.json(); - const profiles = Array.isArray(profilePayload?.profiles) ? profilePayload.profiles : []; - expect(profiles.some((item: { name?: string }) => item.name === profileName)).toBeTruthy(); - }); + } finally { + if (aclUuid) { + await deleteAccessListViaAPI(request, aclUuid, token); + } + } }); }); diff --git a/tests/security-enforcement/acl-dropdown-regression.spec.ts b/tests/security-enforcement/acl-dropdown-regression.spec.ts index 18a837b56..51649147f 100644 --- a/tests/security-enforcement/acl-dropdown-regression.spec.ts +++ b/tests/security-enforcement/acl-dropdown-regression.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from '@playwright/test'; +import { deleteAccessList, deleteProxyHostByDomain } from '../utils/proxy-refusal-helpers'; type SelectionPair = { aclLabel: string; @@ -109,10 +110,41 @@ async function selectNonDefaultPair( } test.describe('ProxyHostForm ACL and Security Headers Dropdown Regression', () => { + // Set by the test as it creates fixtures. An enabled ACL or a proxy host attached to one is enforced by later specs. + let createdProxyDomain: string | undefined; + let createdAclUuid: string | undefined; + + test.afterEach(async ({ page }) => { + // Proxy host first: it references the ACL. + if (createdProxyDomain) { + await deleteProxyHostByDomain(page.request, createdProxyDomain); + createdProxyDomain = undefined; + } + if (createdAclUuid) { + await deleteAccessList(page.request, createdAclUuid); + createdAclUuid = undefined; + } + }); + test('should keep ACL and Security Headers behavior equivalent across create/edit flows', async ({ page }) => { const suffix = Date.now(); const proxyName = `Dropdown Regression ${suffix}`; const proxyDomain = `dropdown-${suffix}.test.local`; + createdProxyDomain = proxyDomain; + + await test.step('Seed an access list so the dropdown has a non-default option', async () => { + // 203.0.113.0/24 is the RFC 5737 documentation range, so the seed admits no real client. + const response = await page.request.post('/api/v1/access-lists', { + data: { + name: `Dropdown Regression ACL ${suffix}`, + type: 'whitelist', + enabled: true, + ip_rules: JSON.stringify([{ cidr: '203.0.113.0/24', description: 'E2E seed rule' }]), + }, + }); + expect(response.ok()).toBeTruthy(); + createdAclUuid = (await response.json()).uuid; + }); await test.step('Navigate to Proxy Hosts', async () => { await page.goto('/proxy-hosts'); diff --git a/tests/security-enforcement/acl-enforcement.spec.ts b/tests/security-enforcement/acl-enforcement.spec.ts index 1beb79d41..ac3061c1f 100644 --- a/tests/security-enforcement/acl-enforcement.spec.ts +++ b/tests/security-enforcement/acl-enforcement.spec.ts @@ -208,34 +208,33 @@ test.describe('ACL Enforcement', () => { const createdList = await createResponse.json(); expect(createdList.uuid).toBeDefined(); - // Verify the list was created with correct structure - expect(createdList.name).toBe('Test Enforcement ACL'); - - // Test a listed, non-admin address using POST: the tester reports the proxy's refusal - const testResponse = await requestContext.post( - `/api/v1/access-lists/${createdList.uuid}/test`, - { data: { ip_address: '203.0.113.77' } } - ); - expect(testResponse.ok()).toBe(true); - const testResult = await testResponse.json(); - expect(testResult.allowed).toBe(false); - expect(testResult.reason).toBe('Blocked by blacklist rule'); - - // Test a listed address inside the admin whitelist: the proxy exempts it, so the tester must admit it - const adminResponse = await requestContext.post( - `/api/v1/access-lists/${createdList.uuid}/test`, - { data: { ip_address: '10.255.255.255' } } - ); - expect(adminResponse.ok()).toBe(true); - const adminResult = await adminResponse.json(); - expect(adminResult.allowed).toBe(true); - expect(adminResult.reason).toBe('Address is in the admin whitelist'); - - // Cleanup: Delete the test ACL - const deleteResponse = await requestContext.delete( - `/api/v1/access-lists/${createdList.uuid}` - ); - expect(deleteResponse.ok()).toBe(true); + try { + // Verify the list was created with correct structure + expect(createdList.name).toBe('Test Enforcement ACL'); + + // Test a listed, non-admin address using POST: the tester reports the proxy's refusal + const testResponse = await requestContext.post( + `/api/v1/access-lists/${createdList.uuid}/test`, + { data: { ip_address: '203.0.113.77' } } + ); + expect(testResponse.ok()).toBe(true); + const testResult = await testResponse.json(); + expect(testResult.allowed).toBe(false); + expect(testResult.reason).toBe('Blocked by blacklist rule'); + + // Test a listed address inside the admin whitelist: the proxy exempts it, so the tester must admit it + const adminResponse = await requestContext.post( + `/api/v1/access-lists/${createdList.uuid}/test`, + { data: { ip_address: '10.255.255.255' } } + ); + expect(adminResponse.ok()).toBe(true); + const adminResult = await adminResponse.json(); + expect(adminResult.allowed).toBe(true); + expect(adminResult.reason).toBe('Address is in the admin whitelist'); + } finally { + // Removed even when an assertion above fails, so the enabled blacklist cannot outlive this test + await deleteAccessList(requestContext, createdList.uuid); + } } else { // May fail if ACL already exists or other issue const errorBody = await createResponse.text(); diff --git a/tests/security-enforcement/acl-waf-layering.spec.ts b/tests/security-enforcement/acl-waf-layering.spec.ts index f2398c5df..7a4512a15 100644 --- a/tests/security-enforcement/acl-waf-layering.spec.ts +++ b/tests/security-enforcement/acl-waf-layering.spec.ts @@ -2,6 +2,7 @@ import { test, expect, loginUser, logoutUser, TEST_PASSWORD } from '../fixtures/ import { dismissNewDomainPromptIfPresent, retryAction, waitForLoadingComplete } from '../utils/wait-helpers'; import { caddyProxyOrigin, createUserViaApi, getAuthTokenFromPage } from '../utils/api-helpers'; +import { deleteAccessList, deleteProxyHostByDomain } from '../utils/proxy-refusal-helpers'; /** * Integration: ACL & WAF Layering (Defense in Depth) @@ -23,6 +24,9 @@ import { caddyProxyOrigin, createUserViaApi, getAuthTokenFromPage } from '../uti */ test.describe('ACL & WAF Layering', () => { + // Set by the deny-all ACL test; afterEach removes it after the proxy host that references it. + let createdAccessListUuid: string | undefined; + // Assigned fresh per test in beforeEach below (not a single shared // const across all 4 tests). Empirically, all 4 tests racing for the // SAME static domain caused intermittent "domain already exists" 400s @@ -126,43 +130,31 @@ test.describe('ACL & WAF Layering', () => { // with the orphaned one. Plain API calls never navigate the page at // all, so they can't race a pending navigation, and they're // substantially faster than a full SPA page load + `networkidle` wait. - test.afterEach(async ({ page, adminUser }) => { + test.afterEach(async ({ request }) => { + // Runs on the fixture's request context (admin storage state), not on `page`: a test may have logged `page` in as + // the restricted user, and that session cannot delete proxy hosts, access lists or users. try { - // Cleanup proxy - const proxiesResponse = await page.request.get('/api/v1/proxy-hosts', { - headers: { Authorization: `Bearer ${adminUser.token}` }, - }); - if (proxiesResponse.ok()) { - const proxies = await proxiesResponse.json(); - const matchingProxy = Array.isArray(proxies) - ? proxies.find((proxy: { domain_names?: string; domainNames?: string }) => - proxy.domain_names === testProxy.domain || proxy.domainNames === testProxy.domain - ) - : undefined; - if (matchingProxy?.uuid) { - await page.request.delete(`/api/v1/proxy-hosts/${matchingProxy.uuid}`, { - headers: { Authorization: `Bearer ${adminUser.token}` }, - }); - } + // Proxy host first: it references the access list removed below + await deleteProxyHostByDomain(request, testProxy.domain); + + if (createdAccessListUuid) { + await deleteAccessList(request, createdAccessListUuid); + createdAccessListUuid = undefined; } - // Cleanup user - const usersResponse = await page.request.get('/api/v1/users', { - headers: { Authorization: `Bearer ${adminUser.token}` }, - }); + const usersResponse = await request.get('/api/v1/users'); if (usersResponse.ok()) { const users = await usersResponse.json(); const matchingUser = Array.isArray(users) ? users.find((user: { email?: string }) => user.email === testUser.email) : undefined; if (matchingUser?.id) { - await page.request.delete(`/api/v1/users/${matchingUser.id}`, { - headers: { Authorization: `Bearer ${adminUser.token}` }, - }); + await request.delete(`/api/v1/users/${matchingUser.id}`); } } - } catch { - // Ignore cleanup errors + } catch (error) { + // Surface cleanup failures: a silent failure here leaves enabled lists behind for later specs + console.warn('ACL & WAF Layering cleanup failed:', error); } }); @@ -455,6 +447,7 @@ test.describe('ACL & WAF Layering', () => { headers: { Authorization: `Bearer ${adminToken}` }, }); expect(accessListResponse.ok()).toBe(true); + createdAccessListUuid = (await accessListResponse.json()).uuid; await page.goto('/proxy-hosts'); diff --git a/tests/security-enforcement/zzz-security-ui/access-lists-crud.spec.ts b/tests/security-enforcement/zzz-security-ui/access-lists-crud.spec.ts index 3fc090249..6440cd133 100644 --- a/tests/security-enforcement/zzz-security-ui/access-lists-crud.spec.ts +++ b/tests/security-enforcement/zzz-security-ui/access-lists-crud.spec.ts @@ -18,7 +18,8 @@ import { waitForLoadingComplete, waitForModal, waitForDialog, waitForDebounce } import { waitForAPIHealth, getAccessListViaAPI } from '../../utils/api-helpers'; import { clickSwitch } from '../../utils/ui-helpers'; import { generateUniqueId } from '../../fixtures/test-data'; -import type { Page } from '@playwright/test'; +import type { APIRequestContext, Page } from '@playwright/test'; +import { deleteAccessList } from '../../utils/proxy-refusal-helpers'; /** * RFC 5737 TEST-NET-3 documentation range. No real client uses it, and the @@ -46,6 +47,15 @@ const DOCUMENTATION_CIDR = '203.0.113.0/24'; * as certificates.spec.ts's `createCustomCertViaAPI` for * `CertificateResponse`'s analogous mismatch. */ +async function listAccessListUuids(request: APIRequestContext): Promise { + const response = await request.get('/api/v1/access-lists'); + if (!response.ok()) { + throw new Error(`Failed to list access lists: ${response.status()} ${await response.text()}`); + } + const lists = (await response.json()) as Array<{ uuid: string }>; + return lists.map((list) => list.uuid); +} + async function seedAccessList( page: Page, overrides: { name?: string; type?: string } = {} @@ -69,6 +79,14 @@ async function seedAccessList( } test.describe('Access Lists - CRUD Operations', () => { + // Lists that existed before this spec ran. The spec creates lists through seeds and through the UI, and an enabled + // list left behind is enforced by every later spec, so afterAll removes each list created since this snapshot. + let preexistingAccessListUuids: Set | undefined; + + test.beforeAll(async ({ request }) => { + preexistingAccessListUuids = new Set(await listAccessListUuids(request)); + }); + test.beforeEach(async ({ page, adminUser }) => { await waitForAPIHealth(page.request); await loginUser(page, adminUser); @@ -78,6 +96,18 @@ test.describe('Access Lists - CRUD Operations', () => { await waitForLoadingComplete(page); }); + test.afterAll(async ({ request }) => { + // Without a snapshot there is no way to tell this spec's lists from pre-existing ones, so delete nothing. + if (!preexistingAccessListUuids) { + return; + } + for (const uuid of await listAccessListUuids(request)) { + if (!preexistingAccessListUuids.has(uuid)) { + await deleteAccessList(request, uuid); + } + } + }); + // Helper to get the Create Access List button const getCreateButton = (page: import('@playwright/test').Page) => page.getByRole('button', { name: /create.*access.*list/i }).first(); diff --git a/tests/security-enforcement/zzz-security-ui/real-time-logs.spec.ts b/tests/security-enforcement/zzz-security-ui/real-time-logs.spec.ts index 3cc79f9e7..982460042 100644 --- a/tests/security-enforcement/zzz-security-ui/real-time-logs.spec.ts +++ b/tests/security-enforcement/zzz-security-ui/real-time-logs.spec.ts @@ -13,7 +13,10 @@ * - Performance (2 tests): high volume logs, buffer limits */ +import { request as playwrightRequest, type APIRequestContext } from '@playwright/test'; import { test, expect, loginUser } from '../../fixtures/auth-fixtures'; +import { STORAGE_STATE } from '../../constants'; +import { withSecurityEnabled } from '../../utils/security-helpers'; import { waitForLoadingComplete } from '../../utils/wait-helpers'; /** @@ -41,7 +44,8 @@ const SELECTORS = { blockedOnlyCheckbox: 'input[type="checkbox"]', // Log display - logContainer: '.font-mono.text-xs', + // Scrollable log pane only. The bare '.font-mono.text-xs' also matches the GeoIP card's
on /security. + logContainer: '.overflow-y-auto.font-mono.text-xs', logEntry: '[data-testid="log-entry"]', logCount: '[data-testid="log-count"]', emptyState: 'text=No logs yet', @@ -86,30 +90,19 @@ async function waitForWebSocketConnection(page: import('@playwright/test').Page) } test.describe('Real-Time Logs Viewer', () => { - // Note: These tests require Cerberus (security module) to be enabled. - // The LiveLogViewer component is only rendered when Cerberus is active. - // Tests will be skipped if the component is not visible on the /security page. - - // Track if LiveLogViewer is available (Cerberus enabled) - let cerberusEnabled = false; - - test.beforeAll(async ({ browser }) => { - // Check once at the start if Cerberus is enabled - // Use stored auth state from playwright/.auth/user.json - const context = await browser.newContext({ - storageState: 'playwright/.auth/user.json', - }); - const page = await context.newPage(); - - // Navigate to security page - await page.goto('/security'); - await expect(page.getByRole('main')).toBeVisible({ timeout: 10000 }); - - // Check if LiveLogViewer is visible (only shown when Cerberus is enabled) - const connectionStatus = page.locator('[data-testid="connection-status"]'); - cerberusEnabled = await connectionStatus.isVisible({ timeout: 3000 }).catch(() => false); + // The LiveLogViewer is only rendered when Cerberus is enabled, and the security shard baseline leaves it + // off. Turn it on for this file and restore the previous state afterwards. + let requestContext: APIRequestContext; + let restoreSecurityState: () => Promise; + + test.beforeAll(async ({ baseURL }) => { + requestContext = await playwrightRequest.newContext({ baseURL, storageState: STORAGE_STATE }); + restoreSecurityState = await withSecurityEnabled(requestContext, { cerberus: true }); + }); - await context.close(); + test.afterAll(async () => { + await restoreSecurityState(); + await requestContext.dispose(); }); // ========================================================================= diff --git a/tests/security/emergency-operations.spec.ts b/tests/security/emergency-operations.spec.ts index 234f6ea27..d69c36b84 100644 --- a/tests/security/emergency-operations.spec.ts +++ b/tests/security/emergency-operations.spec.ts @@ -1,4 +1,5 @@ import { test, expect } from '@playwright/test'; +import { deleteAccessList, deleteProxyHostByDomain } from '../utils/proxy-refusal-helpers'; /** * Emergency & Break-Glass Operations Workflow @@ -79,57 +80,69 @@ test.describe('Emergency & Break-Glass Operations', () => { const suffix = Date.now(); const aclName = `Emergency-ACL-${suffix}`; const proxyDomain = `emergency-acl-${suffix}.test.local`; - - await test.step('Create ACL prerequisite through API for deterministic dropdown options', async () => { - const createAclResponse = await page.request.post('/api/v1/access-lists', { - data: { - name: aclName, - type: 'whitelist', - description: 'ACL prerequisite for emergency regression test', - enabled: true, - ip_rules: JSON.stringify([{ cidr: '10.0.0.0/8' }]), - }, + // Tracked so the finally block removes exactly what this test created. An enabled ACL left behind + // is enforced against every later request from a non-whitelisted admin IP. + let aclUuid: string | undefined; + + try { + await test.step('Create ACL prerequisite through API for deterministic dropdown options', async () => { + const createAclResponse = await page.request.post('/api/v1/access-lists', { + data: { + name: aclName, + type: 'whitelist', + description: 'ACL prerequisite for emergency regression test', + enabled: true, + ip_rules: JSON.stringify([{ cidr: '10.0.0.0/8' }]), + }, + }); + expect(createAclResponse.ok()).toBeTruthy(); + aclUuid = (await createAclResponse.json()).uuid; }); - expect(createAclResponse.ok()).toBeTruthy(); - }); - await test.step('Create proxy host and select created ACL in dropdown', async () => { - await page.goto('/proxy-hosts'); - await expect(page.getByRole('main')).toBeVisible({ timeout: 10000 }); + await test.step('Create proxy host and select created ACL in dropdown', async () => { + await page.goto('/proxy-hosts'); + await expect(page.getByRole('main')).toBeVisible({ timeout: 10000 }); - await openCreateProxyModal(page); - const dialog = page.getByRole('dialog'); + await openCreateProxyModal(page); + const dialog = page.getByRole('dialog'); - await dialog.locator('#proxy-name').fill(`Emergency ACL Regression ${suffix}`); - await dialog.locator('#domain-names').click(); - await page.keyboard.type(proxyDomain); - await page.keyboard.press('Tab'); - await dismissDomainDialog(page); + await dialog.locator('#proxy-name').fill(`Emergency ACL Regression ${suffix}`); + await dialog.locator('#domain-names').click(); + await page.keyboard.type(proxyDomain); + await page.keyboard.press('Tab'); + await dismissDomainDialog(page); - await dialog.locator('#forward-host').fill('127.0.0.1'); - await dialog.locator('#forward-port').fill('8080'); + await dialog.locator('#forward-host').fill('127.0.0.1'); + await dialog.locator('#forward-port').fill('8080'); - const aclTrigger = dialog.getByRole('combobox', { name: /access control list/i }); - const selectedAclLabel = await selectOptionByName( - page, - aclTrigger, - new RegExp(aclName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'i') - ); - await expect(aclTrigger).toContainText(selectedAclLabel); + const aclTrigger = dialog.getByRole('combobox', { name: /access control list/i }); + const selectedAclLabel = await selectOptionByName( + page, + aclTrigger, + new RegExp(aclName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'i') + ); + await expect(aclTrigger).toContainText(selectedAclLabel); - await saveProxyHost(page); - }); + await saveProxyHost(page); + }); - await test.step('Edit proxy host and verify ACL selection persisted', async () => { - await openEditProxyModalForDomain(page, proxyDomain); + await test.step('Edit proxy host and verify ACL selection persisted', async () => { + await openEditProxyModalForDomain(page, proxyDomain); - const dialog = page.getByRole('dialog'); - const aclTrigger = dialog.getByRole('combobox', { name: /access control list/i }); - await expect(aclTrigger).toContainText(new RegExp(aclName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'i')); + const dialog = page.getByRole('dialog'); + const aclTrigger = dialog.getByRole('combobox', { name: /access control list/i }); + await expect(aclTrigger).toContainText(new RegExp(aclName.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'), 'i')); - await dialog.getByRole('button', { name: /cancel/i }).click(); - await expect(dialog).not.toBeVisible({ timeout: 5000 }); - }); + await dialog.getByRole('button', { name: /cancel/i }).click(); + await expect(dialog).not.toBeVisible({ timeout: 5000 }); + }); + } finally { + // Proxy host first: a host that still references the ACL would block its deletion. + await deleteProxyHostByDomain(page.request, proxyDomain); + if (aclUuid) { + await deleteAccessList(page.request, aclUuid); + } + } }); // Use emergency token diff --git a/tests/security/geoip-database-card.spec.ts b/tests/security/geoip-database-card.spec.ts index e4d405a59..10b228d23 100644 --- a/tests/security/geoip-database-card.spec.ts +++ b/tests/security/geoip-database-card.spec.ts @@ -191,12 +191,17 @@ test.describe('GeoIP database card', () => { }); await test.step('Confirming removes the data-folder copy and returns to the bundled source', async () => { - await page.getByRole('dialog').getByRole('button', { name: USE_INCLUDED_NAME }).click(); + // The dialog closes only once the removal has succeeded. While it is open the page behind it is aria-hidden, so + // the card's own button is out of the accessibility tree and a count of 0 would pass before the DELETE completes. + // Waiting for the dialog to close is the settle signal; the card button is then checked against the refreshed UI. + const dialog = page.getByRole('dialog'); + await dialog.getByRole('button', { name: USE_INCLUDED_NAME }).click(); + await expect(dialog).toBeHidden(); await expect(card.getByRole('button', { name: USE_INCLUDED_NAME })).toHaveCount(0); - const status = await readStatus(requestContext); - expect(status.source).toBe('bundled'); - expect(status.data_copy_present).toBe(false); + await expect + .poll(async () => readStatus(requestContext)) + .toMatchObject({ source: 'bundled', data_copy_present: false }); }); await test.step('No GeoIP alert is raised while the bundled copy is in use', async () => { diff --git a/tests/utils/phase5-helpers.ts b/tests/utils/phase5-helpers.ts index 739b63030..31e5bd2eb 100644 --- a/tests/utils/phase5-helpers.ts +++ b/tests/utils/phase5-helpers.ts @@ -1023,7 +1023,8 @@ export const LIVE_LOG_SELECTORS = { levelSelect: 'select >> text=All Levels', sourceSelect: 'select >> text=All Sources', blockedOnlyCheckbox: 'input[type="checkbox"]', - logContainer: '.font-mono.text-xs', + // Scrollable log pane only; see real-time-logs.spec.ts for why the bare classes are ambiguous on /security. + logContainer: '.overflow-y-auto.font-mono.text-xs', logEntry: '[data-testid="log-entry"]', blockedEntry: '.bg-red-900\\/30', logCount: '[data-testid="log-count"]', diff --git a/tests/utils/proxy-refusal-helpers.ts b/tests/utils/proxy-refusal-helpers.ts index bf8ab99de..a9cb38545 100644 --- a/tests/utils/proxy-refusal-helpers.ts +++ b/tests/utils/proxy-refusal-helpers.ts @@ -144,3 +144,16 @@ export async function createProxyHost(request: APIRequestContext, options: Proxy export async function deleteProxyHost(request: APIRequestContext, uuid: string): Promise { await withRateLimitRetry(() => request.delete(`${PROXY_HOSTS_PATH}/${uuid}`)); } + +/** Deletes the proxy host whose `domain_names` equals the given domain. A no-op when none matches. */ +export async function deleteProxyHostByDomain(request: APIRequestContext, domain: string): Promise { + const listResponse = await withRateLimitRetry(() => request.get(PROXY_HOSTS_PATH)); + if (!listResponse.ok()) { + return; + } + const hosts = (await listResponse.json()) as Array<{ uuid: string; domain_names: string }>; + const match = hosts.find((host) => host.domain_names === domain); + if (match) { + await deleteProxyHost(request, match.uuid); + } +}