diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/AdminProxyHandler.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/AdminProxyHandler.java index 28a314019ecca..2de7de9f2dd35 100644 --- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/AdminProxyHandler.java +++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/AdminProxyHandler.java @@ -165,26 +165,43 @@ protected HttpClient createHttpClient() throws ServletException { // This class allows the request body to be replayed, the default implementation // does not protected class ReplayableProxyContentProvider extends ProxyInputStreamContentProvider { - private Boolean firstIteratorCalled = false; + static final int MIN_REPLAY_BODY_BUFFER_SIZE = 64; + private boolean bodyBufferAvailable = false; + private boolean bodyBufferMaxSizeReached = false; private final ByteArrayOutputStream bodyBuffer; - protected ReplayableProxyContentProvider(HttpServletRequest request, HttpServletResponse response, Request proxyRequest, InputStream input) { + private final long httpInputMaxReplayBufferSize; + + protected ReplayableProxyContentProvider(HttpServletRequest request, HttpServletResponse response, + Request proxyRequest, InputStream input, + int httpInputMaxReplayBufferSize) { super(request, response, proxyRequest, input); - bodyBuffer = new ByteArrayOutputStream(Math.max(request.getContentLength(), 0)); + bodyBuffer = new ByteArrayOutputStream( + Math.min(Math.max(request.getContentLength(), MIN_REPLAY_BODY_BUFFER_SIZE), + httpInputMaxReplayBufferSize)); + this.httpInputMaxReplayBufferSize = httpInputMaxReplayBufferSize; } @Override public Iterator iterator() { - if (firstIteratorCalled) { + if (bodyBufferAvailable) { return Collections.singleton(ByteBuffer.wrap(bodyBuffer.toByteArray())).iterator(); } else { - firstIteratorCalled = true; + bodyBufferAvailable = true; return super.iterator(); } } @Override protected ByteBuffer onRead(byte[] buffer, int offset, int length) { - bodyBuffer.write(buffer, offset, length); + if (!bodyBufferMaxSizeReached) { + if (bodyBuffer.size() + length < httpInputMaxReplayBufferSize) { + bodyBuffer.write(buffer, offset, length); + } else { + bodyBufferMaxSizeReached = true; + bodyBufferAvailable = false; + bodyBuffer.reset(); + } + } return super.onRead(buffer, offset, length); } } @@ -217,8 +234,10 @@ protected Request copyRequest(HttpRequest oldRequest, URI newURI) { @Override protected ContentProvider proxyRequestContent(HttpServletRequest request, - HttpServletResponse response, Request proxyRequest) throws IOException { - return new ReplayableProxyContentProvider(request, response, proxyRequest, request.getInputStream()); + HttpServletResponse response, Request proxyRequest) + throws IOException { + return new ReplayableProxyContentProvider(request, response, proxyRequest, request.getInputStream(), + config.getHttpInputMaxReplayBufferSize()); } @Override diff --git a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java index 17226329e8d89..5ce8968835100 100644 --- a/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java +++ b/pulsar-proxy/src/main/java/org/apache/pulsar/proxy/server/ProxyConfiguration.java @@ -480,6 +480,16 @@ public class ProxyConfiguration implements PulsarConfiguration { ) private int httpOutputBufferSize = 32*1024; + @FieldContext( + minValue = 1, + category = CATEGORY_HTTP, + doc = "Http input buffer max size.\n\n" + + "The maximum amount of data that will be buffered for incoming http requests " + + "so that the request body can be replayed when the backend broker " + + "issues a redirect response." + ) + private int httpInputMaxReplayBufferSize = 5 * 1024 * 1024; + @FieldContext( minValue = 1, category = CATEGORY_HTTP, diff --git a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/AdminProxyHandlerTest.java b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/AdminProxyHandlerTest.java index ad7f0e5e24615..4dceb926707ef 100644 --- a/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/AdminProxyHandlerTest.java +++ b/pulsar-proxy/src/test/java/org/apache/pulsar/proxy/server/AdminProxyHandlerTest.java @@ -18,33 +18,56 @@ */ package org.apache.pulsar.proxy.server; -import static org.mockito.Mockito.*; - -import org.eclipse.jetty.client.api.Request; -import org.testng.Assert; -import org.testng.annotations.Test; - -import java.io.InputStream; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import java.io.ByteArrayInputStream; import java.io.ByteArrayOutputStream; +import java.io.InputStream; import java.lang.reflect.Field; - +import java.nio.ByteBuffer; +import java.util.Iterator; +import javax.servlet.ServletConfig; +import javax.servlet.ServletContext; +import javax.servlet.ServletException; import javax.servlet.http.HttpServletRequest; import javax.servlet.http.HttpServletResponse; +import org.eclipse.jetty.client.HttpClient; +import org.eclipse.jetty.client.api.Request; +import org.testng.Assert; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; public class AdminProxyHandlerTest { + private AdminProxyHandler adminProxyHandler; + + @BeforeClass + public void setupMocks() throws ServletException { + // given + HttpClient httpClient = mock(HttpClient.class); + adminProxyHandler = new AdminProxyHandler(mock(ProxyConfiguration.class), + mock(BrokerDiscoveryProvider.class)) { + @Override + protected HttpClient createHttpClient() throws ServletException { + return httpClient; + } + }; + ServletConfig servletConfig = mock(ServletConfig.class); + when(servletConfig.getServletName()).thenReturn("AdminProxyHandler"); + when(servletConfig.getServletContext()).thenReturn(mock(ServletContext.class)); + adminProxyHandler.init(servletConfig); + } @Test public void replayableProxyContentProviderTest() throws Exception { - - AdminProxyHandler adminProxyHandler = new AdminProxyHandler(mock(ProxyConfiguration.class), - mock(BrokerDiscoveryProvider.class)); - HttpServletRequest request = mock(HttpServletRequest.class); doReturn(-1).when(request).getContentLength(); try { - AdminProxyHandler.ReplayableProxyContentProvider replayableProxyContentProvider = adminProxyHandler.new ReplayableProxyContentProvider( - request, mock(HttpServletResponse.class), mock(Request.class), mock(InputStream.class)); + AdminProxyHandler.ReplayableProxyContentProvider replayableProxyContentProvider = + adminProxyHandler.new ReplayableProxyContentProvider( + request, mock(HttpServletResponse.class), mock(Request.class), mock(InputStream.class), + 1024); Field field = replayableProxyContentProvider.getClass().getDeclaredField("bodyBuffer"); field.setAccessible(true); Assert.assertEquals(((ByteArrayOutputStream) field.get(replayableProxyContentProvider)).size(), 0); @@ -53,4 +76,68 @@ public void replayableProxyContentProviderTest() throws Exception { } } + + @Test + public void shouldLimitReplayBodyBufferSize() throws Exception { + HttpServletRequest request = mock(HttpServletRequest.class); + int maxRequestBodySize = 1024 * 1024; + int requestBodySize = maxRequestBodySize + 1; + doReturn(requestBodySize).when(request).getContentLength(); + byte[] inputBuffer = new byte[requestBodySize]; + + AdminProxyHandler.ReplayableProxyContentProvider replayableProxyContentProvider = + adminProxyHandler.new ReplayableProxyContentProvider(request, mock(HttpServletResponse.class), + mock(Request.class), new ByteArrayInputStream(inputBuffer), + maxRequestBodySize); + + // when + + // content is consumed + Iterator byteBufferIterator = replayableProxyContentProvider.iterator(); + int consumedBytes = 0; + while (byteBufferIterator.hasNext()) { + ByteBuffer byteBuffer = byteBufferIterator.next(); + consumedBytes += byteBuffer.limit(); + } + + // then + Assert.assertEquals(consumedBytes, requestBodySize); + Field field = replayableProxyContentProvider.getClass().getDeclaredField("bodyBufferMaxSizeReached"); + field.setAccessible(true); + Assert.assertEquals(((boolean) field.get(replayableProxyContentProvider)), true); + } + + @Test + public void shouldReplayBodyBuffer() { + // given + HttpServletRequest request = mock(HttpServletRequest.class); + int maxRequestBodySize = 1024 * 1024; + byte[] inputBuffer = new byte[maxRequestBodySize - 1]; + for (int i = 0; i < inputBuffer.length; i++) { + inputBuffer[i] = (byte) (i & 0xff); + } + doReturn(inputBuffer.length).when(request).getContentLength(); + + AdminProxyHandler.ReplayableProxyContentProvider replayableProxyContentProvider = + adminProxyHandler.new ReplayableProxyContentProvider(request, mock(HttpServletResponse.class), + mock(Request.class), new ByteArrayInputStream(inputBuffer), + maxRequestBodySize); + + ByteBuffer consumeBuffer = ByteBuffer.allocate(maxRequestBodySize); + // content can be consumed multiple times + for (int i = 0; i < 3; i++) { + // when + consumeBuffer.clear(); + Iterator byteBufferIterator = replayableProxyContentProvider.iterator(); + while (byteBufferIterator.hasNext()) { + ByteBuffer byteBuffer = byteBufferIterator.next(); + consumeBuffer.put(byteBuffer); + } + consumeBuffer.flip(); + byte[] consumedBytes = new byte[consumeBuffer.limit()]; + consumeBuffer.get(consumedBytes); + // then + Assert.assertEquals(consumedBytes, inputBuffer); + } + } }