Skip to content

fix: enforce request body size limit during buffering in AiProxyPlugin #6897

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wy471x wants to merge 4 commits into
base: master
Choose a base branch
from
+29 −12
Open

fix: enforce request body size limit during buffering in AiProxyPlugin #6897

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
cf83bb7
fix: enforce request body size limit during buffering in AiProxyPlugin
wy471x Aug 6, 2026
b6b9986
test: add unit test for DataBufferLimitException → 413 in AiProxyPlugin
wy471x Aug 6, 2026
75d19ca
Merge branch 'master' into fix_AiProxyPlugin-validates-body-size
Aias00 Aug 14, 2026
43981cf
Merge branch 'master' into fix_AiProxyPlugin-validates-body-size
wy471x Aug 14, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 9 additions & 12 deletions ...ugin-ai-proxy/src/main/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPlugin.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferLimitException;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
Expand All @@ -63,7 +64,7 @@ public class AiProxyPlugin extends AbstractShenyuPlugin {
/**
* Maximum request body size: 5MB.
*/
private static final long MAX_REQUEST_BODY_SIZE_BYTES = 5 * 1024 * 1024L;
private static final int MAX_REQUEST_BODY_SIZE_BYTES = 5 * 1024 * 1024;

private final AiModelFactoryRegistry aiModelFactoryRegistry;

Expand Down Expand Up @@ -100,18 +101,8 @@ protected Mono<Void> doExecute(
CacheKeyUtils.INST.getKey(
selector.getId(), Constants.DEFAULT_RULE));

return DataBufferUtils.join(exchange.getRequest().getBody())
return DataBufferUtils.join(exchange.getRequest().getBody(), MAX_REQUEST_BODY_SIZE_BYTES)

@nikhiln64 nikhiln64 Aug 9, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Context

Nice catch, this is a real fix. Buffering the whole body with the single argument DataBufferUtils.join before checking readableByteCount meant the memory was already spent by the time the limit was enforced, and passing the maxSize into join so it aborts mid stream and onErrorResume maps it to 413 is the right shape.

▎ Example: The test mocks DataBufferUtils.join statically and matches the size argument with anyInt(), so it only proves the DataBufferLimitException to 413 mapping and never pins the limit the plugin is responsible for. A regression that weakened the limit, say passing Integer.MAX_VALUE or a mis scaled constant, would still pass green.

Suggestion

An ArgumentCaptor or an eq(MAX_REQUEST_BODY_SIZE_BYTES) matcher so the test asserts join is called with exactly the 5MB constant would close that gap. If it is easy, one unmocked test that feeds a real body over the limit through doExecute and asserts 413 with no downstream call would make the guarantee end to end.

NOTE

Not blocking, the production change looks correct as is.

.flatMap(dataBuffer -> {
// Validate actual body size after reading, not just Content-Length header
final int actualSize = dataBuffer.readableByteCount();
if (actualSize > MAX_REQUEST_BODY_SIZE_BYTES) {
DataBufferUtils.release(dataBuffer);
LOG.warn("[AiProxy] Request body size {} exceeds maximum allowed size {}",
actualSize, MAX_REQUEST_BODY_SIZE_BYTES);
exchange.getResponse().setStatusCode(HttpStatus.PAYLOAD_TOO_LARGE);
return exchange.getResponse().setComplete();
}

final String requestBody = dataBuffer.toString(StandardCharsets.UTF_8);
DataBufferUtils.release(dataBuffer);

Expand Down Expand Up @@ -152,6 +143,12 @@ protected Mono<Void> doExecute(
return handleStreamRequest(exchange, selector, requestBody, primaryConfig, selectorHandle);
}
return handleNonStreamRequest(exchange, selector, requestBody, primaryConfig, selectorHandle);
})
.onErrorResume(DataBufferLimitException.class, e -> {
LOG.warn("[AiProxy] Request body exceeds maximum allowed size {} bytes",
MAX_REQUEST_BODY_SIZE_BYTES);
exchange.getResponse().setStatusCode(HttpStatus.PAYLOAD_TOO_LARGE);
return exchange.getResponse().setComplete();
});
}

Expand Down
20 changes: 20 additions & 0 deletions ...-ai-proxy/src/test/java/org/apache/shenyu/plugin/ai/proxy/enhanced/AiProxyPluginTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.context.ApplicationContext;
import org.springframework.core.io.buffer.DataBufferLimitException;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.mock.http.server.reactive.MockServerHttpRequest;
Expand All @@ -57,6 +59,7 @@

import java.util.Optional;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
Expand Down Expand Up @@ -339,4 +342,21 @@ public void testExecutorServiceError() {
verify(configService).resolveAdminFallbackConfig(primaryConfig, handle);
verify(executorService).execute(any(), any(), any());
}

@Test
public void testRequestBodyExceedsMaxSize() {
final AiProxyHandle handle = new AiProxyHandle();
aiProxyPluginHandler.getSelectorCachedHandle()
.cachedHandle(CacheKeyUtils.INST.getKey(SELECTOR_ID, Constants.DEFAULT_RULE), handle);

try (MockedStatic<DataBufferUtils> dataBufferUtilsMock = mockStatic(DataBufferUtils.class)) {
dataBufferUtilsMock.when(() -> DataBufferUtils.join(any(), anyInt()))
.thenReturn(Mono.error(new DataBufferLimitException("Request body exceeds limit")));

StepVerifier.create(plugin.doExecute(exchange, mock(ShenyuPluginChain.class), selector, rule))
.verifyComplete();

assertEquals(HttpStatus.PAYLOAD_TOO_LARGE, exchange.getResponse().getStatusCode());
}
}
}
Loading