concat batch responses manually

Author: NazariiDenha

jsonrpc/server.go

@@ -535,9 +535,8 @@ func (s *Server) handleBatchRequest(ctx context.Context, batchReq []json.RawMess
 		return nil, finalHeaders, nil
 	}
 
-	result, err := json.Marshal(responses)
-
-	return result, finalHeaders, err // todo: fix batch request aggregate header
+	// todo: fix batch request aggregate header
+	return concatBatchResponses(responses), finalHeaders, nil
 }
 
 func isBatch(reader *bufio.Reader) bool {
@@ -568,6 +567,27 @@ func isNilOrEmpty(i any) (bool, error) {
 	}
 }
 
+// concatBatchResponses builds the JSON array from elements that already valid
+// JSON, so it joins bytes instead of re-encoding. json.Marshal would run every
+// byte through compact() again, which is costly
+func concatBatchResponses(responses []json.RawMessage) []byte {
+	size := len(responses) + 1
+	for _, response := range responses {
+		size += len(response)
+	}
+
+	result := make([]byte, 0, size)
+	result = append(result, '[')
+	for i, response := range responses {
+		if i > 0 {
+			result = append(result, ',')
+		}
+		result = append(result, response...)
+	}
+
+	return append(result, ']')
+}
+
 // TODO: add recover() to catch panics from handlers/validators and return a JSON-RPC internal error
 // instead of crashing the HTTP connection
 func (s *Server) handleRequest(ctx context.Context, req *Request) (*response, http.Header, error) {

jsonrpc/server_test.go

@@ -1051,3 +1051,37 @@ func TestBatchResponseSizeLimit(t *testing.T) {
 		}
 	})
 }
+
+func TestBatchArrayIsByteIdenticalToJSONMarshal(t *testing.T) {
+	server := jsonrpc.NewServer(1, log.NewNopZapLogger())
+	require.NoError(t, server.RegisterMethods(jsonrpc.Method{
+		Name:    "echo",
+		Params:  []jsonrpc.Parameter{{Name: "data"}},
+		Handler: func(data string) (string, *jsonrpc.Error) { return data, nil },
+	}))
+
+	payloads := []string{
+		"<script>a && b</script>",
+		"line\u2028sep\u2029par",
+		`quote " backslash \`,
+		"",
+	}
+
+	elements := make([]string, len(payloads))
+	for i, payload := range payloads {
+		params, err := json.Marshal([]string{payload})
+		require.NoError(t, err)
+		elements[i] = fmt.Sprintf(`{"jsonrpc":"2.0","id":%d,"method":"echo","params":%s}`, i+1, params)
+	}
+
+	body, _, err := server.HandleReader(t.Context(), strings.NewReader("["+strings.Join(elements, ",")+"]"))
+	require.NoError(t, err)
+
+	var got []json.RawMessage
+	require.NoError(t, json.Unmarshal(body, &got))
+	require.Len(t, got, len(payloads))
+
+	want, err := json.Marshal(got)
+	require.NoError(t, err)
+	assert.Equal(t, string(want), string(body))
+}