Add batch request limit and response size limit in a RPC batch request (#2357)

* Add batch request limit and response size limit in a RPC batch request

* Move appending result after check of total response size in batch request

Author: Kourin1996

rpc/errors.go

@@ -54,6 +54,7 @@ var (
 	_ Error = new(invalidRequestError)
 	_ Error = new(invalidMessageError)
 	_ Error = new(invalidParamsError)
+	_ Error = new(responseTooLargeError)
 )
 
 const defaultErrorCode = -32000
@@ -101,3 +102,9 @@ type invalidParamsError struct{ message string }
 func (e *invalidParamsError) ErrorCode() int { return -32602 }
 
 func (e *invalidParamsError) Error() string { return e.message }
+
+type responseTooLargeError struct{}
+
+func (e *responseTooLargeError) ErrorCode() int { return -32003 }
+
+func (e *responseTooLargeError) Error() string { return "response too large" }

rpc/handler.go

@@ -28,6 +28,11 @@ import (
 	"github.com/celo-org/celo-blockchain/log"
 )
 
+const (
+	BatchRequestLimit    = 1000
+	BatchResponseMaxSize = 25 * 1000 * 1000
+)
+
 // handler handles JSON-RPC messages. There is one handler per connection. Note that
 // handler is not safe for concurrent use. Message handling never blocks indefinitely
 // because RPCs are processed on background goroutines launched by handler.
@@ -100,6 +105,13 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 		})
 		return
 	}
+	// Returns an error if number of requests exceeds an allowed maximum
+	if len(msgs) > BatchRequestLimit {
+		h.startCallProc(func(cp *callProc) {
+			h.conn.writeJSON(cp.ctx, errorMessage(&invalidRequestError{"batch too large"}))
+		})
+		return
+	}
 
 	// Handle non-call messages first:
 	calls := make([]*jsonrpcMessage, 0, len(msgs))
@@ -112,10 +124,20 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 		return
 	}
 	// Process calls on a goroutine because they may block indefinitely:
+	responseBytes := 0
 	h.startCallProc(func(cp *callProc) {
 		answers := make([]*jsonrpcMessage, 0, len(msgs))
-		for _, msg := range calls {
+		for idx, msg := range calls {
 			if answer := h.handleCallMsg(cp, msg); answer != nil {
+				// Once total size of responses exceeds an allowed maximum,
+				// generate error messages for all remaining calls and stop further processing
+				if responseBytes += len(answer.Result); responseBytes > BatchResponseMaxSize {
+					for i := idx; i < len(calls); i++ {
+						errMsg := msg.errorResponse(&responseTooLargeError{})
+						answers = append(answers, errMsg)
+					}
+					break
+				}
 				answers = append(answers, answer)
 			}
 		}

rpc/server_test.go

@@ -19,13 +19,16 @@ package rpc
 import (
 	"bufio"
 	"bytes"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"net"
 	"path/filepath"
 	"strings"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/require"
 )
 
 func TestServerRegisterName(t *testing.T) {
@@ -150,3 +153,166 @@ func TestServerShortLivedConn(t *testing.T) {
 		}
 	}
 }
+
+// TestBatchRequestLimit verifies that the server returns "batch too large" error when the number
+// of JSON-RPC calls in a batch exceeds the defined limit
+func TestBatchRequestLimit(t *testing.T) {
+	server := newTestServer()
+	defer server.Stop()
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal("can't listen:", err)
+	}
+	defer listener.Close()
+	go server.ServeListener(listener)
+
+	var (
+		request  = `{"jsonrpc":"2.0","id":1,"method":"rpc_modules"}`
+		wantResp = `{"jsonrpc":"2.0","id":null,"error":{"code":-32600,"message":"batch too large"}}` + "\n"
+		deadline = time.Now().Add(10 * time.Second)
+	)
+
+	// Create a batch request containing (BatchRequestLimit + 1) calls
+	var reqBuf bytes.Buffer
+	reqBuf.WriteString("[")
+	for i := 0; i < BatchRequestLimit+1; i++ {
+		if i > 0 {
+			reqBuf.WriteString(",")
+		}
+		reqBuf.WriteString(request)
+	}
+	reqBuf.WriteString("]\n")
+
+	// Write the request to the server and then close the write side of the connection
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		t.Fatal("can't dial:", err)
+	}
+	defer conn.Close()
+	conn.SetDeadline(deadline)
+	conn.Write(reqBuf.Bytes())
+	conn.(*net.TCPConn).CloseWrite()
+
+	// Verify that the server returns the "batch too large" error
+	buf := make([]byte, 100)
+	n, err := conn.Read(buf)
+	if err != nil {
+		t.Fatal("read error:", err)
+	}
+	if !bytes.Equal(buf[:n], []byte(wantResp)) {
+		t.Fatalf("wrong response: expected=%s, got=%s", wantResp, buf[:n])
+	}
+
+	// Ensure that the connection is closed and no additional data is returned (EOF expected)
+	n, err = conn.Read(make([]byte, 1))
+	require.Zero(t, n)
+	require.ErrorIs(t, io.EOF, err)
+}
+
+// TestBatchResponseMaxSize verifies that the server returns successful responses
+// until the total response size exceeds the configured maximum. Once the threshold is exceeded,
+// the server should respond with a specific "response too large" error for the remaining requests
+// in the batch and then close the connection.
+func TestBatchResponseMaxSize(t *testing.T) {
+	server := newTestServer()
+	defer server.Stop()
+
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal("can't listen:", err)
+	}
+	defer listener.Close()
+	go server.ServeListener(listener)
+
+	var (
+		strSize  = 25 * 1000
+		strValue = strings.Repeat("A", strSize)
+
+		successfulResultData    = fmt.Sprintf(`{"String":"%s","Int":1,"Args":null}`, strValue)
+		successfulResponse      = fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"result":%s}`, successfulResultData)
+		successfulResultDataLen = len(successfulResultData)
+		successfulResponseLen   = len(successfulResponse)
+		tooLargeErrorResponse   = `{"jsonrpc":"2.0","id":1,"error":{"code":-32003,"message":"response too large"}}`
+
+		deadline = time.Now().Add(10 * time.Second)
+	)
+
+	// create a batch request
+	var reqBuf bytes.Buffer
+	reqBuf.WriteString("[")
+	for i := 0; i < BatchRequestLimit; i++ {
+		if i > 0 {
+			reqBuf.WriteString(",")
+		}
+		reqBuf.WriteString(fmt.Sprintf(`{"jsonrpc":"2.0","id":1,"method":"test_echo","params":["%s",1]}`, strValue))
+	}
+	reqBuf.WriteString("]\n")
+
+	// send request
+	conn, err := net.Dial("tcp", listener.Addr().String())
+	if err != nil {
+		t.Fatal("can't dial:", err)
+	}
+	defer conn.Close()
+	conn.SetDeadline(deadline)
+	conn.Write(reqBuf.Bytes())
+	conn.(*net.TCPConn).CloseWrite()
+
+	buf := make([]byte, successfulResponseLen)
+
+	// mustConsume reads from the connection until the entire provided buffer is filled
+	var mustConsume func(t *testing.T, buf []byte) int
+	mustConsume = func(t *testing.T, buf []byte) int {
+		t.Helper()
+		n, err := conn.Read(buf)
+		if err != nil {
+			t.Fatal("read error:", err)
+		}
+		if n < len(buf) {
+			mustConsume(t, buf[n:])
+		}
+		return n
+	}
+	// mustConsumeAndCompare consumes from the connection and compares the read data with the given expected data
+	mustConsumeAndCompare := func(t *testing.T, expected []byte) {
+		t.Helper()
+		mustConsume(t, buf[:len(expected)])
+		if !bytes.Equal(buf[:len(expected)], expected) {
+			t.Fatalf("wrong response: expected=%s, actual=%s", expected, buf[:len(expected)])
+		}
+	}
+
+	var (
+		totalResultSize = 0
+		errorBeginsAt   = 0
+	)
+
+	mustConsumeAndCompare(t, []byte("["))
+	// Read through each response until the cumulative size limit is exceeded
+	for i := 0; i < BatchRequestLimit; i++ {
+		if totalResultSize += successfulResultDataLen; totalResultSize > BatchResponseMaxSize {
+			// Record the first index where the error should begin
+			errorBeginsAt = i
+			break
+		}
+
+		if i > 0 {
+			mustConsumeAndCompare(t, []byte(","))
+		}
+		mustConsumeAndCompare(t, []byte(successfulResponse))
+	}
+
+	// From the point where the total size exceeded the limit,
+	// check whether all responses for the remaining calls is "response too large"
+	for i := errorBeginsAt; i < BatchRequestLimit; i++ {
+		mustConsumeAndCompare(t, []byte(","))
+		mustConsumeAndCompare(t, []byte(tooLargeErrorResponse))
+	}
+	mustConsumeAndCompare(t, []byte("]\n"))
+
+	// Ensure that the connection is closed and no additional data is returned (EOF expected)
+	n, err := conn.Read(make([]byte, 1))
+	require.Zero(t, n)
+	require.ErrorIs(t, io.EOF, err)
+}