Fix WebSocket limit for large agent messages #45
@@ -108,6 +108,7 @@ func connectOnce(ctx context.Context, cfg Config, handle Handler) error {
|
|||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("dial: %w", err)
|
return fmt.Errorf("dial: %w", err)
|
||||||
}
|
}
|
||||||
|
conn.SetReadLimit(api.MaxWebSocketMessageBytes)
|
||||||
// On a successful upgrade coder/websocket transfers ownership of the
|
// On a successful upgrade coder/websocket transfers ownership of the
|
||||||
// response stream to conn and deliberately sets res.Body to nil. Closing
|
// response stream to conn and deliberately sets res.Body to nil. Closing
|
||||||
// the connection below releases that stream.
|
// the connection below releases that stream.
|
||||||
|
|||||||
@@ -2,12 +2,16 @@ package wsclient
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
|
"encoding/json"
|
||||||
"net/http"
|
"net/http"
|
||||||
"net/http/httptest"
|
"net/http/httptest"
|
||||||
|
"strings"
|
||||||
"testing"
|
"testing"
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/coder/websocket"
|
"github.com/coder/websocket"
|
||||||
|
|
||||||
|
"gitea.dcglab.co.uk/steve/restic-manager/internal/api"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestConnectOnceCleanDisconnectDoesNotPanic(t *testing.T) {
|
func TestConnectOnceCleanDisconnectDoesNotPanic(t *testing.T) {
|
||||||
@@ -44,3 +48,56 @@ func TestConnectOnceCleanDisconnectDoesNotPanic(t *testing.T) {
|
|||||||
t.Fatalf("server websocket: %v", err)
|
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")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -10,6 +10,13 @@ import (
|
|||||||
// (not iota ints) makes traffic readable in logs and packet captures.
|
// (not iota ints) makes traffic readable in logs and packet captures.
|
||||||
type MessageType string
|
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.
|
// Agent → server message types.
|
||||||
const (
|
const (
|
||||||
MsgHello MessageType = "hello"
|
MsgHello MessageType = "hello"
|
||||||
|
|||||||
@@ -77,6 +77,7 @@ func AgentHandler(deps HandlerDeps) stdhttp.Handler {
|
|||||||
slog.Warn("ws accept failed", "err", err, "host_id", host.ID)
|
slog.Warn("ws accept failed", "err", err, "host_id", host.ID)
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
conn.SetReadLimit(api.MaxWebSocketMessageBytes)
|
||||||
|
|
||||||
c := NewConn(host.ID, conn)
|
c := NewConn(host.ID, conn)
|
||||||
// Keep agents alive across NAT boxes; coder/websocket
|
// Keep agents alive across NAT boxes; coder/websocket
|
||||||
|
|||||||
@@ -123,6 +123,68 @@ func TestWSHelloAndHeartbeat(t *testing.T) {
|
|||||||
t.Error("heartbeat did not update last_seen_at")
|
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) {
|
func TestWSRejectsOldProtocol(t *testing.T) {
|
||||||
t.Parallel()
|
t.Parallel()
|
||||||
url, token, _, _, _ := setupTestHub(t)
|
url, token, _, _, _ := setupTestHub(t)
|
||||||
|
|||||||
Reference in New Issue
Block a user