don't continue rest jsonrpc request of batch if overflow

Author: 317787106

framework/src/main/java/org/tron/core/services/jsonrpc/JsonRpcServlet.java

@@ -9,6 +9,7 @@
 import com.googlecode.jsonrpc4j.JsonRpcInterceptor;
 import com.googlecode.jsonrpc4j.JsonRpcServer;
 import com.googlecode.jsonrpc4j.ProxyUtil;
+import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
@@ -116,29 +117,94 @@ protected void doPost(HttpServletRequest req, HttpServletResponse resp) throws I
       return;
     }
 
+    int maxResponseSize = parameter.getJsonRpcMaxResponseSize();
+    if (isBatch) {
+      handleBatch(resp, rootNode, maxResponseSize);
+    } else {
+      handleSingle(req, resp, rootNode, body, maxResponseSize);
+    }
+  }
+
+  private void handleSingle(HttpServletRequest req, HttpServletResponse resp,
+      JsonNode rootNode, byte[] body, int maxResponseSize) throws IOException {
     CachedBodyRequestWrapper cachedReq = new CachedBodyRequestWrapper(req, body);
     BufferedResponseWrapper bufferedResp = new BufferedResponseWrapper(
-        resp, parameter.getJsonRpcMaxResponseSize());
+        resp, maxResponseSize);
 
     try {
       rpcServer.handle(cachedReq, bufferedResp);
     } catch (RuntimeException e) {
       logger.error("RPC execution failed", e);
-      JsonNode idNode = isBatch ? null : rootNode.get("id");
-      writeJsonRpcError(resp, JsonRpcError.INTERNAL_ERROR, "Internal error", idNode, isBatch);
+      writeJsonRpcError(resp, JsonRpcError.INTERNAL_ERROR, "Internal error",
+          rootNode.get("id"), false);
       return;
     }
 
     if (bufferedResp.isOverflow()) {
-      JsonNode idNode = isBatch ? null : rootNode.get("id");
       writeJsonRpcError(resp, JsonRpcError.RESPONSE_TOO_LARGE,
-          "Response exceeds the limit of " + parameter.getJsonRpcMaxResponseSize() + " bytes",
-          idNode, isBatch);
+          "Response exceeds the limit of " + maxResponseSize + " bytes",
+          rootNode.get("id"), false);
       return;
     }
     bufferedResp.commitToResponse();
   }
 
+  private void handleBatch(HttpServletResponse resp, JsonNode rootNode, int maxResponseSize)
+      throws IOException {
+
+    ArrayNode batchResult = MAPPER.createArrayNode();
+    int accumulatedSize = 2; // "[]"
+
+    for (int i = 0; i < rootNode.size(); i++) {
+      byte[] subBody;
+      try {
+        subBody = MAPPER.writeValueAsBytes(rootNode.get(i));
+      } catch (JsonProcessingException e) {
+        writeJsonRpcError(resp, JsonRpcError.INTERNAL_ERROR, "Internal error", null, true);
+        return;
+      }
+
+      ByteArrayOutputStream subOutput = new ByteArrayOutputStream();
+      try {
+        rpcServer.handleRequest(new ByteArrayInputStream(subBody), subOutput);
+      } catch (RuntimeException e) {
+        logger.error("RPC execution failed for batch sub-request {}", i, e);
+        writeJsonRpcError(resp, JsonRpcError.INTERNAL_ERROR, "Internal error", null, true);
+        return;
+      }
+
+      byte[] responseBytes = subOutput.toByteArray();
+      if (responseBytes.length == 0) {
+        continue; // notification — no response
+      }
+
+      // comma separator between array elements
+      int addition = responseBytes.length + (!batchResult.isEmpty() ? 1 : 0);
+      if (maxResponseSize > 0 && accumulatedSize + addition > maxResponseSize) {
+        writeJsonRpcError(resp, JsonRpcError.RESPONSE_TOO_LARGE,
+            "Response exceeds the limit of " + maxResponseSize + " bytes", null, true);
+        return;
+      }
+      accumulatedSize += addition;
+
+      JsonNode responseNode;
+      try {
+        responseNode = MAPPER.readTree(responseBytes);
+      } catch (IOException e) {
+        writeJsonRpcError(resp, JsonRpcError.INTERNAL_ERROR, "Internal error", null, true);
+        return;
+      }
+      batchResult.add(responseNode);
+    }
+
+    byte[] finalBytes = MAPPER.writeValueAsBytes(batchResult);
+    resp.setContentType("application/json; charset=utf-8");
+    resp.setStatus(HttpServletResponse.SC_OK);
+    resp.setContentLength(finalBytes.length);
+    resp.getOutputStream().write(finalBytes);
+    resp.getOutputStream().flush();
+  }
+
   private byte[] readBody(InputStream in) throws IOException {
     ByteArrayOutputStream buffer = new ByteArrayOutputStream();
     byte[] tmp = new byte[4096];

framework/src/test/java/org/tron/core/services/jsonrpc/JsonRpcServletTest.java

@@ -13,6 +13,8 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.googlecode.jsonrpc4j.JsonRpcServer;
 import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
 import java.lang.reflect.Field;
 import java.nio.charset.StandardCharsets;
 import javax.servlet.http.HttpServletRequest;
@@ -87,16 +89,20 @@ public void batchExceedsLimit_returnsExceedLimitAsArray() throws Exception {
   @Test
   public void batchWithinLimit_proceedsToRpcServer() throws Exception {
     CommonParameter.getInstance().jsonRpcMaxBatchSize = 5;
-    byte[] rpcResp = "[{\"result\":\"ok\"}]".getBytes(StandardCharsets.UTF_8);
+    byte[] singleResp = "{\"jsonrpc\":\"2.0\",\"result\":\"ok\",\"id\":1}"
+        .getBytes(StandardCharsets.UTF_8);
     doAnswer(inv -> {
-      HttpServletResponse r = inv.getArgument(1);
-      r.getOutputStream().write(rpcResp);
-      return null;
-    }).when(mockRpcServer).handle(any(HttpServletRequest.class), any(HttpServletResponse.class));
+      OutputStream out = inv.getArgument(1);
+      out.write(singleResp);
+      return 0;
+    }).when(mockRpcServer).handleRequest(any(InputStream.class), any(OutputStream.class));
 
     MockHttpServletResponse resp = doPost("[{\"id\":1},{\"id\":2}]");
     assertEquals(200, resp.getStatus());
-    assertArrayEquals(rpcResp, resp.getContentAsByteArray());
+    JsonNode body = MAPPER.readTree(resp.getContentAsByteArray());
+    assertTrue("batch response must be a JSON array", body.isArray());
+    assertEquals("each sub-request must produce a response", 2, body.size());
+    assertEquals("ok", body.get(0).get("result").asText());
   }
 
   @Test
@@ -113,12 +119,9 @@ public void emptyBatch_returnsInvalidRequest() throws Exception {
   @Test
   public void batchLimitDisabled_largeBatchAllowed() throws Exception {
     CommonParameter.getInstance().jsonRpcMaxBatchSize = 0;
-    byte[] rpcResp = "[]".getBytes(StandardCharsets.UTF_8);
-    doAnswer(inv -> {
-      HttpServletResponse r = inv.getArgument(1);
-      r.getOutputStream().write(rpcResp);
-      return null;
-    }).when(mockRpcServer).handle(any(HttpServletRequest.class), any(HttpServletResponse.class));
+    // write nothing — simulates notifications (no response expected)
+    doAnswer(inv -> 0).when(mockRpcServer)
+        .handleRequest(any(InputStream.class), any(OutputStream.class));
 
     StringBuilder sb = new StringBuilder("[");
     for (int i = 0; i < 500; i++) {
@@ -130,7 +133,9 @@ public void batchLimitDisabled_largeBatchAllowed() throws Exception {
     sb.append("]");
     MockHttpServletResponse resp = doPost(sb.toString());
     assertEquals(200, resp.getStatus());
-    assertArrayEquals(rpcResp, resp.getContentAsByteArray());
+    JsonNode body = MAPPER.readTree(resp.getContentAsByteArray());
+    assertTrue("response must be a JSON array", body.isArray());
+    assertEquals("all notifications produce no response entries", 0, body.size());
   }
 
   // --- rpcServer.handle exceptions ---
@@ -149,7 +154,7 @@ public void rpcServerThrowsRuntimeException_returnsInternalError() throws Except
   @Test
   public void batchRpcServerThrows_internalErrorIsArray() throws Exception {
     doThrow(new RuntimeException("boom")).when(mockRpcServer)
-        .handle((HttpServletRequest) any(), any());
+        .handleRequest(any(InputStream.class), any(OutputStream.class));
     MockHttpServletResponse resp = doPost("[{\"method\":\"eth_blockNumber\"}]");
     assertEquals(200, resp.getStatus());
     JsonNode body = MAPPER.readTree(resp.getContentAsString());
@@ -181,10 +186,10 @@ public void batchResponseTooLarge_returnsErrorArray() throws Exception {
     int limit = 50;
     CommonParameter.getInstance().jsonRpcMaxResponseSize = limit;
     doAnswer(inv -> {
-      HttpServletResponse r = inv.getArgument(1);
-      r.getOutputStream().write(new byte[limit + 1]);
-      return null;
-    }).when(mockRpcServer).handle(any(HttpServletRequest.class), any(HttpServletResponse.class));
+      OutputStream out = inv.getArgument(1);
+      out.write(new byte[limit + 1]);
+      return 0;
+    }).when(mockRpcServer).handleRequest(any(InputStream.class), any(OutputStream.class));
 
     MockHttpServletResponse resp = doPost("[{\"method\":\"eth_getLogs\"}]");
     assertEquals(200, resp.getStatus());
@@ -193,6 +198,30 @@ public void batchResponseTooLarge_returnsErrorArray() throws Exception {
     assertEquals(-32003, body.get(0).get("error").get("code").asInt());
   }
 
+  @Test
+  public void batchShortCircuitsOnOverflow() throws Exception {
+    int limit = 50;
+    CommonParameter.getInstance().jsonRpcMaxResponseSize = limit;
+    int[] callCount = {0};
+    doAnswer(inv -> {
+      OutputStream out = inv.getArgument(1);
+      callCount[0]++;
+      if (callCount[0] == 1) {
+        out.write("{\"result\":\"ok\"}".getBytes(StandardCharsets.UTF_8));
+      } else {
+        out.write(new byte[limit]); // triggers overflow when added to accumulated size
+      }
+      return 0;
+    }).when(mockRpcServer).handleRequest(any(InputStream.class), any(OutputStream.class));
+
+    MockHttpServletResponse resp = doPost("[{\"id\":1},{\"id\":2},{\"id\":3}]");
+    assertEquals(200, resp.getStatus());
+    JsonNode body = MAPPER.readTree(resp.getContentAsString());
+    assertTrue("overflow response must be an array", body.isArray());
+    assertEquals(-32003, body.get(0).get("error").get("code").asInt());
+    assertEquals("third sub-request must not be executed after overflow", 2, callCount[0]);
+  }
+
   // --- normal path ---
 
   @Test