Fix WebSocket limit for large agent messages #45

Merged
steve merged 1 commits from fix-issue-44-websocket-read-limit into main 2026-08-22 10:48:32 +01:00
5 changed files with 128 additions and 0 deletions
+1
View File
@@ -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.
+57
View File
@@ -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")
}
}
+7
View File
@@ -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"
+1
View File
@@ -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
+62
View File
@@ -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)