Skip to content

Commit

Permalink
eth, rpc: add configurable option for wsMessageSizeLimit (ethereum#27801
Browse files Browse the repository at this point in the history
)

This change adds a configurable limit to websocket message. 
---------

Co-authored-by: Martin Holst Swende <martin@swende.se>
  • Loading branch information
2 people authored and hyunchel committed Oct 6, 2023
1 parent 5e18739 commit 315607d
Show file tree
Hide file tree
Showing 5 changed files with 85 additions and 8 deletions.
11 changes: 10 additions & 1 deletion rpc/client_opt.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@ type clientConfig struct {
httpAuth HTTPAuth

// WebSocket options
wsDialer *websocket.Dialer
wsDialer *websocket.Dialer
wsMessageSizeLimit *int64 // wsMessageSizeLimit nil = default, 0 = no limit

// RPC handler options
idgen func() ID
Expand Down Expand Up @@ -66,6 +67,14 @@ func WithWebsocketDialer(dialer websocket.Dialer) ClientOption {
})
}

// WithWebsocketMessageSizeLimit configures the websocket message size limit used by the RPC
// client. Passing a limit of 0 means no limit.
func WithWebsocketMessageSizeLimit(messageSizeLimit int64) ClientOption {
return optionFunc(func(cfg *clientConfig) {
cfg.wsMessageSizeLimit = &messageSizeLimit
})
}

// WithHeader configures HTTP headers set by the RPC client. Headers set using this option
// will be used for both HTTP and WebSocket connections.
func WithHeader(key, value string) ClientOption {
Expand Down
2 changes: 1 addition & 1 deletion rpc/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ func TestServerRegisterName(t *testing.T) {
t.Fatalf("Expected service calc to be registered")
}

wantCallbacks := 13
wantCallbacks := 14
if len(svc.callbacks) != wantCallbacks {
t.Errorf("Expected %d callbacks for service 'service', got %d", wantCallbacks, len(svc.callbacks))
}
Expand Down
4 changes: 4 additions & 0 deletions rpc/testservice_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,10 @@ func (s *testService) EchoWithCtx(ctx context.Context, str string, i int, args *
return echoResult{str, i, args}
}

func (s *testService) Repeat(msg string, i int) string {
return strings.Repeat(msg, i)
}

func (s *testService) PeerInfo(ctx context.Context) PeerInfo {
return PeerInfoFromContext(ctx)
}
Expand Down
14 changes: 9 additions & 5 deletions rpc/websocket.go
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@ const (
wsPingInterval = 30 * time.Second
wsPingWriteTimeout = 5 * time.Second
wsPongTimeout = 30 * time.Second
wsMessageSizeLimit = 32 * 1024 * 1024
wsDefaultReadLimit = 32 * 1024 * 1024
)

var wsBufferPool = new(sync.Pool)
Expand All @@ -60,7 +60,7 @@ func (s *Server) WebsocketHandler(allowedOrigins []string) http.Handler {
log.Debug("WebSocket upgrade failed", "err", err)
return
}
codec := newWebsocketCodec(conn, r.Host, r.Header)
codec := newWebsocketCodec(conn, r.Host, r.Header, wsDefaultReadLimit)
s.ServeCodec(codec, 0)
})
}
Expand Down Expand Up @@ -251,7 +251,11 @@ func newClientTransportWS(endpoint string, cfg *clientConfig) (reconnectFunc, er
}
return nil, hErr
}
return newWebsocketCodec(conn, dialURL, header), nil
messageSizeLimit := int64(wsDefaultReadLimit)
if cfg.wsMessageSizeLimit != nil && *cfg.wsMessageSizeLimit >= 0 {
messageSizeLimit = *cfg.wsMessageSizeLimit
}
return newWebsocketCodec(conn, dialURL, header, messageSizeLimit), nil
}
return connect, nil
}
Expand Down Expand Up @@ -283,8 +287,8 @@ type websocketCodec struct {
pongReceived chan struct{}
}

func newWebsocketCodec(conn *websocket.Conn, host string, req http.Header) ServerCodec {
conn.SetReadLimit(wsMessageSizeLimit)
func newWebsocketCodec(conn *websocket.Conn, host string, req http.Header, readLimit int64) ServerCodec {
conn.SetReadLimit(readLimit)
encode := func(v interface{}, isErrorResponse bool) error {
return conn.WriteJSON(v)
}
Expand Down
62 changes: 61 additions & 1 deletion rpc/websocket_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -113,6 +113,66 @@ func TestWebsocketLargeCall(t *testing.T) {
}
}

// This test checks whether the wsMessageSizeLimit option is obeyed.
func TestWebsocketLargeRead(t *testing.T) {
t.Parallel()

var (
srv = newTestServer()
httpsrv = httptest.NewServer(srv.WebsocketHandler([]string{"*"}))
wsURL = "ws:" + strings.TrimPrefix(httpsrv.URL, "http:")
)
defer srv.Stop()
defer httpsrv.Close()

testLimit := func(limit *int64) {
opts := []ClientOption{}
expLimit := int64(wsDefaultReadLimit)
if limit != nil && *limit >= 0 {
opts = append(opts, WithWebsocketMessageSizeLimit(*limit))
if *limit > 0 {
expLimit = *limit // 0 means infinite
}
}
client, err := DialOptions(context.Background(), wsURL, opts...)
if err != nil {
t.Fatalf("can't dial: %v", err)
}
defer client.Close()
// Remove some bytes for json encoding overhead.
underLimit := int(expLimit - 128)
overLimit := expLimit + 1
if expLimit == wsDefaultReadLimit {
// No point trying the full 32MB in tests. Just sanity-check that
// it's not obviously limited.
underLimit = 1024
overLimit = -1
}
var res string
// Check under limit
if err = client.Call(&res, "test_repeat", "A", underLimit); err != nil {
t.Fatalf("unexpected error with limit %d: %v", expLimit, err)
}
if len(res) != underLimit || strings.Count(res, "A") != underLimit {
t.Fatal("incorrect data")
}
// Check over limit
if overLimit > 0 {
err = client.Call(&res, "test_repeat", "A", expLimit+1)
if err == nil || err != websocket.ErrReadLimit {
t.Fatalf("wrong error with limit %d: %v expecting %v", expLimit, err, websocket.ErrReadLimit)
}
}
}
ptr := func(v int64) *int64 { return &v }

testLimit(ptr(-1)) // Should be ignored (use default)
testLimit(ptr(0)) // Should be ignored (use default)
testLimit(nil) // Should be ignored (use default)
testLimit(ptr(200))
testLimit(ptr(wsDefaultReadLimit * 2))
}

func TestWebsocketPeerInfo(t *testing.T) {
var (
s = newTestServer()
Expand Down Expand Up @@ -206,7 +266,7 @@ func TestClientWebsocketLargeMessage(t *testing.T) {
defer srv.Stop()
defer httpsrv.Close()

respLength := wsMessageSizeLimit - 50
respLength := wsDefaultReadLimit - 50
srv.RegisterName("test", largeRespService{respLength})

c, err := DialWebsocket(context.Background(), wsURL, "")
Expand Down

0 comments on commit 315607d

Please sign in to comment.