rpc: count error responses toward batch response size limit

Fixes #33814

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.encodedSize()
 				if responseBytes > h.batchResponseMaxSize {
 					err := &internalServerError{errcodeResponseTooLarge, errMsgResponseTooLarge}
 					callBuffer.respondWithError(cp.ctx, h.conn, err)
@@ -673,6 +673,12 @@ func (buf *limitedBuffer) Write(data []byte) (int, error) {
 }
 
 func formatErrorData(v any) string {
+	if raw, ok := v.(json.RawMessage); ok {
+		if len(raw) <= 1024 {
+			return string(raw)
+		}
+		return string(raw[:1024]) + "... (truncated)"
+	}
 	buf := limitedBuffer{limit: 1024}
 	err := json.NewEncoder(&buf).Encode(v)
 	switch {

rpc/json.go

@@ -107,6 +107,14 @@ func (msg *jsonrpcMessage) String() string {
 	return string(b)
 }
 
+// encodedSize returns the length of the message when JSON-encoded.
+// Used for batch response size accounting. Error.Data is already RawMessage
+// so it is not double-encoded when marshalling.
+func (msg *jsonrpcMessage) encodedSize() int {
+	b, _ := json.Marshal(msg)
+	return len(b)
+}
+
 func (msg *jsonrpcMessage) errorResponse(err error) *jsonrpcMessage {
 	resp := errorMessage(err)
 	resp.ID = msg.ID
@@ -132,15 +140,19 @@ func errorMessage(err error) *jsonrpcMessage {
 	}
 	de, ok := err.(DataError)
 	if ok {
-		msg.Error.Data = de.ErrorData()
+		if data := de.ErrorData(); data != nil {
+			if enc, err := json.Marshal(data); err == nil {
+				msg.Error.Data = enc
+			}
+		}
 	}
 	return msg
 }
 
 type jsonError struct {
-	Code    int         `json:"code"`
-	Message string      `json:"message"`
-	Data    interface{} `json:"data,omitempty"`
+	Code    int             `json:"code"`
+	Message string          `json:"message"`
+	Data    json.RawMessage `json:"data,omitempty"`
 }
 
 func (err *jsonError) Error() string {
@@ -155,6 +167,13 @@ func (err *jsonError) ErrorCode() int {
 }
 
 func (err *jsonError) ErrorData() interface{} {
+	if len(err.Data) == 0 {
+		return nil
+	}
+	var v interface{}
+	if json.Unmarshal(err.Data, &v) == nil {
+		return v
+	}
 	return err.Data
 }
 

rpc/server_test.go

@@ -53,9 +53,9 @@ func TestServerRegisterName(t *testing.T) {
 		t.Fatalf("Expected service %s to be registered", svcName)
 	}
 
-	wantCallbacks := 14
+	wantCallbacks := 15 // testService methods including ReturnLargeDataError
 	if len(svc.callbacks) != wantCallbacks {
-		t.Errorf("Expected %d callbacks for service 'service', got %d", wantCallbacks, len(svc.callbacks))
+		t.Errorf("Expected %d callbacks for service %q, got %d", wantCallbacks, svcName, len(svc.callbacks))
 	}
 }
 
@@ -173,7 +173,9 @@ func TestServerBatchResponseSizeLimit(t *testing.T) {
 
 	server := newTestServer()
 	defer server.Stop()
-	server.SetBatchLimits(100, 60)
+	// Limit counts full response size (success: jsonrpc+id+result). Each echo response ~65 bytes.
+	// Use limit so two responses exceed it (break when responseBytes > limit).
+	server.SetBatchLimits(100, 129)
 	var (
 		batch  []BatchElem
 		client = DialInProc(server)
@@ -208,6 +210,45 @@ func TestServerBatchResponseSizeLimit(t *testing.T) {
 	}
 }
 
+// TestServerBatchResponseSizeLimitCountsErrorResponses verifies that error responses
+// (including large error.data) count toward the batch response size limit, and that
+// -32003 (response too large) is returned for subsequent items once the limit is exceeded.
+func TestServerBatchResponseSizeLimitCountsErrorResponses(t *testing.T) {
+	t.Parallel()
+
+	server := newTestServer()
+	defer server.Stop()
+	// Each error response from test_returnLargeDataError(60) is ~150 bytes (error envelope + 60-char data).
+	// Set limit so that one error response fits but two exceed it.
+	server.SetBatchLimits(100, 160)
+
+	client := DialInProc(server)
+	batch := []BatchElem{
+		{Method: "test_returnLargeDataError", Args: []any{60}, Result: new(any)},
+		{Method: "test_returnLargeDataError", Args: []any{60}, Result: new(any)},
+		{Method: "test_returnLargeDataError", Args: []any{60}, Result: new(any)},
+	}
+	if err := client.BatchCall(batch); err != nil {
+		t.Fatal("batch call error:", err)
+	}
+
+	// First batch element(s) get the method error (largeDataError); rest get -32003 (response too large).
+	var gotResponseTooLarge int
+	for i := range batch {
+		e := batch[i].Error
+		if e == nil {
+			t.Fatalf("batch elem %d: expected error (method returns error), got nil", i)
+		}
+		if re, ok := e.(Error); ok && re.ErrorCode() == errcodeResponseTooLarge {
+			gotResponseTooLarge++
+		}
+	}
+	// We expect at least one -32003 (response too large) because error responses count toward the limit.
+	if gotResponseTooLarge < 1 {
+		t.Errorf("expected at least one batch elem with -32003 (response too large), got %d; error responses should count toward batch size limit", gotResponseTooLarge)
+	}
+}
+
 func TestServerWebsocketReadLimit(t *testing.T) {
 	t.Parallel()
 

rpc/testservice_test.go

@@ -128,6 +128,19 @@ func (s *testService) ReturnError() error {
 	return testError{}
 }
 
+// largeDataError is used by tests to trigger error responses with large error data.
+type largeDataError struct {
+	Data string
+}
+
+func (e largeDataError) Error() string     { return "largeDataError" }
+func (e largeDataError) ErrorCode() int    { return 555 }
+func (e largeDataError) ErrorData() interface{} { return e.Data }
+
+func (s *testService) ReturnLargeDataError(n int) error {
+	return largeDataError{Data: strings.Repeat("x", n)}
+}
+
 func (s *testService) MarshalError() *MarshalErrObj {
 	return &MarshalErrObj{}
 }