From 27b31371d46bc932853cce36078b28a53088b2b2 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Fri, 26 Feb 2021 13:40:35 +0100
Subject: [PATCH] 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.
---
 rpc/http_test.go        | 25 +++++++++++++++++++++++++
 rpc/testservice_test.go | 10 ++++++++++
 rpc/websocket.go        |  3 ++-
 rpc/websocket_test.go   | 27 +++++++++++++++++++++++++++
 4 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/rpc/http_test.go b/rpc/http_test.go
index fc939ae48f..b75af67c52 100644
--- a/rpc/http_test.go
+++ b/rpc/http_test.go
@@ -98,3 +98,28 @@ func confirmHTTPRequestYieldsStatusCode(t *testing.T, method, contentType, body
 func TestHTTPResponseWithEmptyGet(t *testing.T) {
 	confirmHTTPRequestYieldsStatusCode(t, http.MethodGet, "", "", http.StatusOK)
 }
+
+// This checks that maxRequestContentLength is not applied to the response of a request.
+func TestHTTPRespBodyUnlimited(t *testing.T) {
+	const respLength = maxRequestContentLength * 3
+
+	s := NewServer()
+	defer s.Stop()
+	s.RegisterName("test", largeRespService{respLength})
+	ts := httptest.NewServer(s)
+	defer ts.Close()
+
+	c, err := DialHTTP(ts.URL)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+
+	var r string
+	if err := c.Call(&r, "test_largeResp"); err != nil {
+		t.Fatal(err)
+	}
+	if len(r) != respLength {
+		t.Fatalf("response has wrong length %d, want %d", len(r), respLength)
+	}
+}
diff --git a/rpc/testservice_test.go b/rpc/testservice_test.go
index 6f948a1bac..62afc1df44 100644
--- a/rpc/testservice_test.go
+++ b/rpc/testservice_test.go
@@ -20,6 +20,7 @@ import (
 	"context"
 	"encoding/binary"
 	"errors"
+	"strings"
 	"sync"
 	"time"
 )
@@ -194,3 +195,12 @@ func (s *notificationTestService) HangSubscription(ctx context.Context, val int)
 	}()
 	return subscription, nil
 }
+
+// largeRespService generates arbitrary-size JSON responses.
+type largeRespService struct {
+	length int
+}
+
+func (x largeRespService) LargeResp() string {
+	return strings.Repeat("x", x.length)
+}
diff --git a/rpc/websocket.go b/rpc/websocket.go
index cd60eeb613..ab55ae69c1 100644
--- a/rpc/websocket.go
+++ b/rpc/websocket.go
@@ -37,6 +37,7 @@ const (
 	wsWriteBuffer      = 1024
 	wsPingInterval     = 60 * time.Second
 	wsPingWriteTimeout = 5 * time.Second
+	wsMessageSizeLimit = 15 * 1024 * 1024
 )
 
 var wsBufferPool = new(sync.Pool)
@@ -239,7 +240,7 @@ type websocketCodec struct {
 }
 
 func newWebsocketCodec(conn *websocket.Conn) ServerCodec {
-	conn.SetReadLimit(maxRequestContentLength)
+	conn.SetReadLimit(wsMessageSizeLimit)
 	wc := &websocketCodec{
 		jsonCodec: NewFuncCodec(conn, conn.WriteJSON, conn.ReadJSON).(*jsonCodec),
 		conn:      conn,
diff --git a/rpc/websocket_test.go b/rpc/websocket_test.go
index f54fc3cd54..37ed19476f 100644
--- a/rpc/websocket_test.go
+++ b/rpc/websocket_test.go
@@ -157,6 +157,33 @@ func TestClientWebsocketPing(t *testing.T) {
 	}
 }
 
+// This checks that the websocket transport can deal with large messages.
+func TestClientWebsocketLargeMessage(t *testing.T) {
+	var (
+		srv     = NewServer()
+		httpsrv = httptest.NewServer(srv.WebsocketHandler(nil))
+		wsURL   = "ws:" + strings.TrimPrefix(httpsrv.URL, "http:")
+	)
+	defer srv.Stop()
+	defer httpsrv.Close()
+
+	respLength := wsMessageSizeLimit - 50
+	srv.RegisterName("test", largeRespService{respLength})
+
+	c, err := DialWebsocket(context.Background(), wsURL, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var r string
+	if err := c.Call(&r, "test_largeResp"); err != nil {
+		t.Fatal("call failed:", err)
+	}
+	if len(r) != respLength {
+		t.Fatalf("response has wrong length %d, want %d", len(r), respLength)
+	}
+}
+
 // wsPingTestServer runs a WebSocket server which accepts a single subscription request.
 // When a value arrives on sendPing, the server sends a ping frame, waits for a matching
 // pong and finally delivers a single subscription result.