From 74cb7f660faad2965d2571ac7af1a8604fc1486c Mon Sep 17 00:00:00 2001 From: Steve Cliff Date: Sat, 22 Aug 2026 10:46:00 +0100 Subject: [PATCH] fix(ws): raise bounded agent message limit --- internal/agent/wsclient/client.go | 1 + internal/agent/wsclient/client_test.go | 57 +++++++++++++++++++++++ internal/api/wire.go | 7 +++ internal/server/ws/handler.go | 1 + internal/server/ws/hub_test.go | 62 ++++++++++++++++++++++++++ 5 files changed, 128 insertions(+) diff --git a/internal/agent/wsclient/client.go b/internal/agent/wsclient/client.go index 625b7f7..51b2e92 100644 --- a/internal/agent/wsclient/client.go +++ b/internal/agent/wsclient/client.go @@ -108,6 +108,7 @@ func connectOnce(ctx context.Context, cfg Config, handle Handler) error { if err != nil { return fmt.Errorf("dial: %w", err) } + conn.SetReadLimit(api.MaxWebSocketMessageBytes) // On a successful upgrade coder/websocket transfers ownership of the // response stream to conn and deliberately sets res.Body to nil. Closing // the connection below releases that stream. diff --git a/internal/agent/wsclient/client_test.go b/internal/agent/wsclient/client_test.go index 5be6220..bda07d4 100644 --- a/internal/agent/wsclient/client_test.go +++ b/internal/agent/wsclient/client_test.go @@ -2,12 +2,16 @@ package wsclient import ( "context" + "encoding/json" "net/http" "net/http/httptest" + "strings" "testing" "time" "github.com/coder/websocket" + + "gitea.dcglab.co.uk/steve/restic-manager/internal/api" ) func TestConnectOnceCleanDisconnectDoesNotPanic(t *testing.T) { @@ -44,3 +48,56 @@ func TestConnectOnceCleanDisconnectDoesNotPanic(t *testing.T) { t.Fatalf("server websocket: %v", err) } } + +func TestConnectOnceAcceptsMessageLargerThanDefaultReadLimit(t *testing.T) { + received := make(chan struct{}, 1) + serverErr := make(chan error, 1) + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + conn, err := websocket.Accept(w, r, nil) + if err != nil { + serverErr <- err + return + } + defer conn.CloseNow() //nolint:errcheck + if _, _, err := conn.Read(r.Context()); err != nil { + serverErr <- err + return + } + env := api.Envelope{ + Type: api.MsgConfigUpdate, + Payload: json.RawMessage(`{"padding":"` + strings.Repeat("x", 40*1024) + `"}`), + } + raw, _ := json.Marshal(env) + serverErr <- conn.Write(r.Context(), websocket.MessageText, raw) + <-r.Context().Done() + })) + defer srv.Close() + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- connectOnce(ctx, Config{ + ServerURL: srv.URL, + AgentToken: "test-token", + HeartbeatPeriod: time.Hour, + }, func(_ context.Context, env api.Envelope, _ Sender) error { + if env.Type == api.MsgConfigUpdate { + received <- struct{}{} + cancel() + } + return nil + }) + }() + + select { + case <-received: + case <-time.After(5 * time.Second): + t.Fatal("agent did not receive oversized server message") + } + if err := <-serverErr; err != nil { + t.Fatalf("server websocket: %v", err) + } + if err := <-done; err == nil { + t.Fatal("connectOnce returned nil") + } +} diff --git a/internal/api/wire.go b/internal/api/wire.go index 4573738..1de58ae 100644 --- a/internal/api/wire.go +++ b/internal/api/wire.go @@ -10,6 +10,13 @@ import ( // (not iota ints) makes traffic readable in logs and packet captures. type MessageType string +// MaxWebSocketMessageBytes is the protocol-wide upper bound for one agent ↔ +// server envelope. Snapshot projections and restic JSON log events can exceed +// coder/websocket's 32 KiB default on ordinary repositories, so both peers set +// this limit explicitly. It remains bounded to protect either process from an +// untrusted or malfunctioning peer allocating without limit. +const MaxWebSocketMessageBytes int64 = 8 << 20 + // Agent → server message types. const ( MsgHello MessageType = "hello" diff --git a/internal/server/ws/handler.go b/internal/server/ws/handler.go index 6c54b81..88b9ec9 100644 --- a/internal/server/ws/handler.go +++ b/internal/server/ws/handler.go @@ -77,6 +77,7 @@ func AgentHandler(deps HandlerDeps) stdhttp.Handler { slog.Warn("ws accept failed", "err", err, "host_id", host.ID) return } + conn.SetReadLimit(api.MaxWebSocketMessageBytes) c := NewConn(host.ID, conn) // Keep agents alive across NAT boxes; coder/websocket diff --git a/internal/server/ws/hub_test.go b/internal/server/ws/hub_test.go index 3753cc2..f4b9248 100644 --- a/internal/server/ws/hub_test.go +++ b/internal/server/ws/hub_test.go @@ -123,6 +123,68 @@ func TestWSHelloAndHeartbeat(t *testing.T) { t.Error("heartbeat did not update last_seen_at") } +func TestWSAcceptsSnapshotReportLargerThanDefaultReadLimit(t *testing.T) { + t.Parallel() + url, token, hostID, st, hub := setupTestHub(t) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + c, _, err := websocket.Dial(ctx, url, &websocket.DialOptions{ + HTTPHeader: stdhttp.Header{"Authorization": []string{"Bearer " + token}}, + }) + if err != nil { + t.Fatalf("dial: %v", err) + } + defer c.CloseNow() //nolint:errcheck + + hello, _ := api.Marshal(api.MsgHello, "", api.HelloPayload{ + ProtocolVersion: api.CurrentProtocolVersion, + AgentVersion: "0.1.0", + ResticVersion: "0.17.1", + Hostname: "h1", + OS: api.OSLinux, + Arch: api.ArchAmd64, + }) + helloRaw, _ := json.Marshal(hello) + if err := c.Write(ctx, websocket.MessageText, helloRaw); err != nil { + t.Fatalf("write hello: %v", err) + } + deadline := time.Now().Add(time.Second) + for !hub.Connected(hostID) && time.Now().Before(deadline) { + time.Sleep(10 * time.Millisecond) + } + + report, _ := api.Marshal(api.MsgSnapshotsRpt, "", api.SnapshotsReportPayload{ + Snapshots: []api.Snapshot{{ + ID: strings.Repeat("a", 64), + ShortID: "aaaaaaaa", + Time: time.Now().UTC(), + Hostname: "h1", + Paths: []string{"/" + strings.Repeat("long-path/", 5000)}, + }}, + }) + reportRaw, _ := json.Marshal(report) + if len(reportRaw) <= 32*1024 { + t.Fatalf("test payload is only %d bytes; must exceed old limit", len(reportRaw)) + } + if int64(len(reportRaw)) >= api.MaxWebSocketMessageBytes { + t.Fatalf("test payload %d exceeds protocol limit", len(reportRaw)) + } + if err := c.Write(ctx, websocket.MessageText, reportRaw); err != nil { + t.Fatalf("write snapshots.report: %v", err) + } + + deadline = time.Now().Add(2 * time.Second) + for time.Now().Before(deadline) { + host, err := st.GetHost(context.Background(), hostID) + if err == nil && host.SnapshotCount == 1 { + return + } + time.Sleep(10 * time.Millisecond) + } + t.Fatal("oversized snapshots.report was not projected") +} + func TestWSRejectsOldProtocol(t *testing.T) { t.Parallel() url, token, _, _, _ := setupTestHub(t)