From d1a6871e01033d3ad412143c37f1a7a9f28b850e Mon Sep 17 00:00:00 2001
From: THE MACHINE <themachine@themachine.ai>
Date: Mon, 23 Mar 2026 08:09:07 +0800
Subject: [PATCH] rpc: count error response size in batch response limit

Fix Bug #33814: handleBatch only accounted for len(resp.Result) when
accumulating responseBytes for maxResponseSize. Error responses have
resp.Result == nil, so large error payloads (via rpc.DataError) were
not counted toward the batch response size limit.

Changes:
- rpc/json.go: add size() method to jsonrpcMessage that returns
  len(Result) for success responses, or the serialized size of the
  full error response (including Error.Data) for error responses
- rpc/handler.go: use resp.size() instead of len(resp.Result) in
  handleBatch to count all response types toward maxResponseSize
- rpc/server_test.go: add TestServerBatchResponseSizeLimitError to
  verify error responses with large data are correctly counted
- rpc/testservice_test.go: add largeDataError type and
  ReturnLargeError method to generate error responses with data
  exceeding the batch response size limit
---
 rpc/handler.go          |  2 +-
 rpc/json.go             | 19 ++++++++++++++++
 rpc/server_test.go      | 49 +++++++++++++++++++++++++++++++++++++++++
 rpc/testservice_test.go | 17 ++++++++++++++
 4 files changed, 86 insertions(+), 1 deletion(-)

diff --git a/rpc/handler.go b/rpc/handler.go
index c0af162f13..2f66df0924 100644
--- a/rpc/handler.go
+++ b/rpc/handler.go
@@ -237,7 +237,7 @@ func (h *handler) handleBatch(msgs []*jsonrpcMessage) {
 			resp := h.handleCallMsg(cp, msg)
 			callBuffer.pushResponse(resp)
 			if resp != nil && h.batchResponseMaxSize != 0 {
-				responseBytes += len(resp.Result)
+				responseBytes += resp.size()
 				if responseBytes > h.batchResponseMaxSize {
 					err := &internalServerError{errcodeResponseTooLarge, errMsgResponseTooLarge}
 					callBuffer.respondWithError(cp.ctx, h.conn, err)
diff --git a/rpc/json.go b/rpc/json.go
index fcd801fc95..4ba9ff8dae 100644
--- a/rpc/json.go
+++ b/rpc/json.go
@@ -107,6 +107,25 @@ func (msg *jsonrpcMessage) String() string {
 	return string(b)
 }
 
+// size returns the approximate wire-size of the response.
+// For successful responses it returns len(Result), which is already a json.RawMessage.
+// For error responses it returns the serialized size of the entire error response,
+// since Error.Data may contain large payloads that would be serialized to the wire.
+func (msg *jsonrpcMessage) size() int {
+	if msg == nil {
+		return 0
+	}
+	if msg.Result != nil {
+		return len(msg.Result)
+	}
+	if msg.Error != nil {
+		// Re-serialize to capture full error response size including Error.Data
+		b, _ := json.Marshal(msg)
+		return len(b)
+	}
+	return 0
+}
+
 func (msg *jsonrpcMessage) errorResponse(err error) *jsonrpcMessage {
 	resp := errorMessage(err)
 	resp.ID = msg.ID
diff --git a/rpc/server_test.go b/rpc/server_test.go
index 8334d4e80d..5e1ceb3319 100644
--- a/rpc/server_test.go
+++ b/rpc/server_test.go
@@ -208,6 +208,55 @@ func TestServerBatchResponseSizeLimit(t *testing.T) {
 	}
 }
 
+func TestServerBatchResponseSizeLimitError(t *testing.T) {
+	t.Parallel()
+
+	server := newTestServer()
+	defer server.Stop()
+	// Set itemLimit high enough that we hit maxResponseSize from error responses alone.
+	// Each error response (~54 bytes serialized) should count toward the limit.
+	server.SetBatchLimits(100, 80)
+	var (
+		batch  []BatchElem
+		client = DialInProc(server)
+	)
+	// With maxResponseSize=80 and each error being ~54 bytes,
+	// the first error fits (0→54), second error would exceed (54→108>80),
+	// so the second call should get -32003.
+	for i := 0; i < 3; i++ {
+		batch = append(batch, BatchElem{
+			Method: "test_returnLargeError",
+			Args:   nil,
+			Result: new(any),
+		})
+	}
+	if err := client.BatchCall(batch); err != nil {
+		t.Fatal("error sending batch:", err)
+	}
+	for i := range batch {
+		if i < 1 {
+			// First call should return the actual error, not "too large"
+			re, ok := batch[i].Error.(Error)
+			if !ok {
+				t.Fatalf("batch elem %d has wrong error type: %T", i, batch[i].Error)
+			}
+			if re.ErrorCode() != 444 {
+				t.Errorf("batch elem %d wrong error code, have %d want 444", i, re.ErrorCode())
+			}
+			continue
+		}
+		// After the first, we expect "response too large" (-32003)
+		re, ok := batch[i].Error.(Error)
+		if !ok {
+			t.Fatalf("batch elem %d has wrong error type: %T", i, batch[i].Error)
+		}
+		wantedCode := errcodeResponseTooLarge
+		if re.ErrorCode() != wantedCode {
+			t.Errorf("batch elem %d wrong error code, have %d want %d", i, re.ErrorCode(), wantedCode)
+		}
+	}
+}
+
 func TestServerWebsocketReadLimit(t *testing.T) {
 	t.Parallel()
 
diff --git a/rpc/testservice_test.go b/rpc/testservice_test.go
index 69199e21b7..88d3d56719 100644
--- a/rpc/testservice_test.go
+++ b/rpc/testservice_test.go
@@ -70,6 +70,19 @@ func (testError) Error() string          { return "testError" }
 func (testError) ErrorCode() int         { return 444 }
 func (testError) ErrorData() interface{} { return "testError data" }
 
+type largeDataError struct{}
+
+func (largeDataError) Error() string { return "largeDataError" }
+func (largeDataError) ErrorCode() int {
+	return 444
+}
+func (largeDataError) ErrorData() interface{} {
+	// Returns a 100-byte string to exceed the 60-byte batch response limit
+	// when serialized as JSON: "largeDataError data padding..." (~33 bytes)
+	// plus JSON overhead, the error response itself is ~50 bytes
+	return "largeDataError data padding padding padding"
+}
+
 type MarshalErrObj struct{}
 
 func (o *MarshalErrObj) MarshalText() ([]byte, error) {
@@ -128,6 +141,10 @@ func (s *testService) ReturnError() error {
 	return testError{}
 }
 
+func (s *testService) ReturnLargeError() error {
+	return largeDataError{}
+}
+
 func (s *testService) MarshalError() *MarshalErrObj {
 	return &MarshalErrObj{}
 }