From 667a4088d22ada3878a500deab16dd665824a1d8 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Tue, 14 Feb 2023 10:40:21 +0800
Subject: [PATCH] limit the size of the response packet to 10MB

---
 rpc/errors.go  |  4 +++-
 rpc/handler.go | 23 +++++++++++++++++++++++
 2 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/rpc/errors.go b/rpc/errors.go
index 7188332d551eb..36d2c82a73c74 100644
--- a/rpc/errors.go
+++ b/rpc/errors.go
@@ -61,12 +61,14 @@ const (
 	errcodeDefault                  = -32000
 	errcodeNotificationsUnsupported = -32001
 	errcodeTimeout                  = -32002
+	errcodeResponseTooLarge         = -32003
 	errcodePanic                    = -32603
 	errcodeMarshalError             = -32603
 )
 
 const (
-	errMsgTimeout = "request timed out"
+	errMsgTimeout          = "request timed out"
+	errMsgResponseTooLarge = "response too large"
 )
 
 type methodNotFoundError struct{ method string }
diff --git a/rpc/handler.go b/rpc/handler.go
index 346d8152b8428..45fbb1ed48f72 100644
--- a/rpc/handler.go
+++ b/rpc/handler.go
@@ -149,6 +149,21 @@ func (b *batchCallBuffer) timeout(ctx context.Context, conn jsonWriter) {
 	b.doWrite(ctx, conn, true)
 }
 
+// responseTooLarge sends the responses added so far. For the remaining unanswered call
+// messages, it sends a response too large error response.
+func (b *batchCallBuffer) responseTooLarge(ctx context.Context, conn jsonWriter) {
+	b.mutex.Lock()
+	defer b.mutex.Unlock()
+
+	for _, msg := range b.calls {
+		if !msg.isNotification() {
+			resp := msg.errorResponse(&internalServerError{errcodeResponseTooLarge, errMsgResponseTooLarge})
+			b.resp = append(b.resp, resp)
+		}
+	}
+	b.doWrite(ctx, conn, true)
+}
+
 // doWrite actually writes the response.
 // This assumes b.mutex is held.
 func (b *batchCallBuffer) doWrite(ctx context.Context, conn jsonWriter, isErrorResponse bool) {
@@ -211,6 +226,8 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 			})
 		}
 
+		resBytes := 0
+		maxBytes := 10 * 1000 * 1000
 		for {
 			// No need to handle rest of calls if timed out.
 			if cp.ctx.Err() != nil {
@@ -222,6 +239,12 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 			}
 			resp := h.handleCallMsg(cp, msg)
 			callBuffer.pushResponse(resp)
+			if resp != nil {
+				if resBytes += len(resp.Result); resBytes > maxBytes {
+					callBuffer.responseTooLarge(cp.ctx, h.conn)
+					break
+				}
+			}
 		}
 		if timer != nil {
 			timer.Stop()