Skip to content

Commit 27b3137

Browse files
authored
rpc: add separate size limit for websocket (#22385)
This makes the WebSocket message size limit independent of the limit used for HTTP requests. The new limit for WebSocket messages is 15MB.
1 parent dc109cc commit 27b3137

File tree

4 files changed

+64
-1
lines changed

4 files changed

+64
-1
lines changed

rpc/http_test.go

+25
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,28 @@ func confirmHTTPRequestYieldsStatusCode(t *testing.T, method, contentType, body
9898
func TestHTTPResponseWithEmptyGet(t *testing.T) {
9999
confirmHTTPRequestYieldsStatusCode(t, http.MethodGet, "", "", http.StatusOK)
100100
}
101+
102+
// This checks that maxRequestContentLength is not applied to the response of a request.
103+
func TestHTTPRespBodyUnlimited(t *testing.T) {
104+
const respLength = maxRequestContentLength * 3
105+
106+
s := NewServer()
107+
defer s.Stop()
108+
s.RegisterName("test", largeRespService{respLength})
109+
ts := httptest.NewServer(s)
110+
defer ts.Close()
111+
112+
c, err := DialHTTP(ts.URL)
113+
if err != nil {
114+
t.Fatal(err)
115+
}
116+
defer c.Close()
117+
118+
var r string
119+
if err := c.Call(&r, "test_largeResp"); err != nil {
120+
t.Fatal(err)
121+
}
122+
if len(r) != respLength {
123+
t.Fatalf("response has wrong length %d, want %d", len(r), respLength)
124+
}
125+
}

rpc/testservice_test.go

+10
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ import (
2020
"context"
2121
"encoding/binary"
2222
"errors"
23+
"strings"
2324
"sync"
2425
"time"
2526
)
@@ -194,3 +195,12 @@ func (s *notificationTestService) HangSubscription(ctx context.Context, val int)
194195
}()
195196
return subscription, nil
196197
}
198+
199+
// largeRespService generates arbitrary-size JSON responses.
200+
type largeRespService struct {
201+
length int
202+
}
203+
204+
func (x largeRespService) LargeResp() string {
205+
return strings.Repeat("x", x.length)
206+
}

rpc/websocket.go

+2-1
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ const (
3737
wsWriteBuffer = 1024
3838
wsPingInterval = 60 * time.Second
3939
wsPingWriteTimeout = 5 * time.Second
40+
wsMessageSizeLimit = 15 * 1024 * 1024
4041
)
4142

4243
var wsBufferPool = new(sync.Pool)
@@ -239,7 +240,7 @@ type websocketCodec struct {
239240
}
240241

241242
func newWebsocketCodec(conn *websocket.Conn) ServerCodec {
242-
conn.SetReadLimit(maxRequestContentLength)
243+
conn.SetReadLimit(wsMessageSizeLimit)
243244
wc := &websocketCodec{
244245
jsonCodec: NewFuncCodec(conn, conn.WriteJSON, conn.ReadJSON).(*jsonCodec),
245246
conn: conn,

rpc/websocket_test.go

+27
Original file line numberDiff line numberDiff line change
@@ -157,6 +157,33 @@ func TestClientWebsocketPing(t *testing.T) {
157157
}
158158
}
159159

160+
// This checks that the websocket transport can deal with large messages.
161+
func TestClientWebsocketLargeMessage(t *testing.T) {
162+
var (
163+
srv = NewServer()
164+
httpsrv = httptest.NewServer(srv.WebsocketHandler(nil))
165+
wsURL = "ws:" + strings.TrimPrefix(httpsrv.URL, "http:")
166+
)
167+
defer srv.Stop()
168+
defer httpsrv.Close()
169+
170+
respLength := wsMessageSizeLimit - 50
171+
srv.RegisterName("test", largeRespService{respLength})
172+
173+
c, err := DialWebsocket(context.Background(), wsURL, "")
174+
if err != nil {
175+
t.Fatal(err)
176+
}
177+
178+
var r string
179+
if err := c.Call(&r, "test_largeResp"); err != nil {
180+
t.Fatal("call failed:", err)
181+
}
182+
if len(r) != respLength {
183+
t.Fatalf("response has wrong length %d, want %d", len(r), respLength)
184+
}
185+
}
186+
160187
// wsPingTestServer runs a WebSocket server which accepts a single subscription request.
161188
// When a value arrives on sendPing, the server sends a ping frame, waits for a matching
162189
// pong and finally delivers a single subscription result.

0 commit comments

Comments
 (0)