fix: handle max body size

dhis2 · Feb 6, 2025 · 035597c · 035597c
1 parent 1d64c79
commit 035597c
Show file tree

Hide file tree

Showing 2 changed files with 44 additions and 0 deletions.
diff --git a/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/route/RouteService.java b/dhis-2/dhis-services/dhis-service-core/src/main/java/org/hisp/dhis/route/RouteService.java
@@ -29,6 +29,8 @@
 
 import static org.hisp.dhis.config.HibernateEncryptionConfig.AES_128_STRING_ENCRYPTOR;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
 import io.netty.handler.timeout.ReadTimeoutException;
 import jakarta.servlet.http.HttpServletRequest;
 import java.io.IOException;
@@ -57,6 +59,7 @@
 import org.jasypt.encryption.pbe.PBEStringCleanablePasswordEncryptor;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
+import org.springframework.core.io.buffer.DataBufferLimitException;
 import org.springframework.http.HttpHeaders;
 import org.springframework.http.HttpMethod;
 import org.springframework.http.HttpStatus;
@@ -86,6 +89,8 @@ public class RouteService {
 
   @Autowired @Getter @Setter private ClientHttpConnector clientHttpConnector;
 
+  @Autowired @Getter @Setter private ObjectMapper objectMapper;
+
   private static final Set<String> ALLOWED_REQUEST_HEADERS =
       Set.of(
           "accept",
@@ -246,6 +251,19 @@ public ResponseEntity<String> execute(
             .onErrorReturn(
                 throwable -> throwable.getCause() instanceof ReadTimeoutException,
                 new ResponseEntity<>(HttpStatus.GATEWAY_TIMEOUT))
+            .onErrorResume(
+                throwable -> throwable.getCause() instanceof DataBufferLimitException,
+                    throwable -> {
+                      try {
+                        return Mono.just(
+                            new ResponseEntity<>(
+                                objectMapper.writeValueAsString(
+                                    Map.of("message", throwable.getCause().getMessage())),
+                                HttpStatus.BAD_GATEWAY));
+                      } catch (JsonProcessingException e) {
+                        throw new RuntimeException(e);
+                      }
+                    })
             .block();
     HttpHeaders responseHeaders = filterResponseHeaders(response.getHeaders());
     String responseBody = response.getBody();

diff --git a/.../dhis-test-web-api/src/test/java/org/hisp/dhis/webapi/controller/RouteControllerTest.java b/.../dhis-test-web-api/src/test/java/org/hisp/dhis/webapi/controller/RouteControllerTest.java
@@ -31,6 +31,7 @@
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockserver.model.HttpRequest.request;
 
 import com.fasterxml.jackson.core.JsonProcessingException;
@@ -141,6 +142,31 @@ public void afterEach() {
       mockServerClient.reset();
     }
 
+    @Test
+    void testRunRouteWhenResponseBodyExceedsLimit() throws JsonProcessingException {
+      mockServerClient
+          .when(request().withPath("/"))
+          .respond(org.mockserver.model.HttpResponse.response("{}"));
+
+      Map<String, Object> route = new HashMap<>();
+      route.put("name", "route-under-test");
+      route.put("url", "https://dhis2.org");
+
+      HttpResponse postHttpResponse = POST("/routes", jsonMapper.writeValueAsString(route));
+      HttpResponse runHttpResponse =
+          GET(
+              "/routes/{id}/run",
+              postHttpResponse.content().get("response.uid").as(JsonString.class).string());
+
+      String message =
+          runHttpResponse
+              .error(HttpStatus.BAD_GATEWAY)
+              .get("message")
+              .as(JsonString.class)
+              .string();
+      assertTrue(message.startsWith("Exceeded limit on max bytes to buffer : "));
+    }
+
     @Test
     void testRunRouteWhenResponseDurationExceedsRouteResponseTimeout()
         throws JsonProcessingException {