[Proxy] Limit replay buffer size in AdminProxyHandler

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<ByteBuffer> 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

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,

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<ByteBuffer> 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<ByteBuffer> 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);
+        }
+    }
 }