Limit size of parsed data in MultipartParser

2 files changed, 87 insertions, 14 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParser.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParser.java
index 8d03ca74500..a9e24943c0d 100644
--- a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParser.java
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParser.java
@@ -5,6 +5,7 @@ import com.yahoo.container.jdisc.HttpRequest;
 import com.yahoo.vespa.hosted.controller.api.application.v4.EnvironmentResource;
 import org.apache.commons.fileupload.MultipartStream;
 import org.apache.commons.fileupload.ParameterParser;
+import org.apache.commons.fileupload.util.Streams;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -21,6 +22,16 @@ import java.util.Map;
  */
 public class MultipartParser {
 
+    private final long maxDataLength;
+
+    public MultipartParser() {
+        this(2 * (long) Math.pow(1024, 3)); // 2 GB
+    }
+
+    MultipartParser(long maxDataLength) {
+        this.maxDataLength = maxDataLength;
+    }
+
     /**
      * Parses the given multi-part request and returns all the parts indexed by their name.
      * 
@@ -37,10 +48,13 @@ public class MultipartParser {
      */
     public Map<String, byte[]> parse(String contentTypeHeader, InputStream data, URI uri) {
         try {
+            LimitedOutputStream output = new LimitedOutputStream(maxDataLength);
             ParameterParser parameterParser = new ParameterParser();
             Map<String, String> contentType = parameterParser.parse(contentTypeHeader, ';');
-            if (contentType.containsKey("application/zip"))
-                return Map.of(EnvironmentResource.APPLICATION_ZIP, data.readAllBytes());
+            if (contentType.containsKey("application/zip")) {
+                Streams.copy(data, output, false);
+                return Map.of(EnvironmentResource.APPLICATION_ZIP, output.toByteArray());
+            }
             if ( ! contentType.containsKey("multipart/form-data"))
                 throw new IllegalArgumentException("Expected a multipart or application/zip message, but got Content-Type: " + contentTypeHeader);
             String boundary = contentType.get("boundary");
@@ -55,17 +69,17 @@ public class MultipartParser {
                 if (contentDispositionContent == null)
                     throw new IllegalArgumentException("Missing Content-Disposition header in a multipart body part");
                 Map<String, String> contentDisposition = parameterParser.parse(contentDispositionContent, ';');
-                ByteArrayOutputStream output = new ByteArrayOutputStream();
                 multipartStream.readBodyData(output);
                 parts.put(contentDisposition.get("name"), output.toByteArray());
+                output.reset();
                 nextPart = multipartStream.readBoundary();
             }
             return parts;
         }
-        catch(MultipartStream.MalformedStreamException e) {
+        catch (MultipartStream.MalformedStreamException e) {
             throw new IllegalArgumentException("Malformed multipart/form-data request", e);
         } 
-        catch(IOException e) {
+        catch (IOException e) {
             throw new IllegalArgumentException("IO error reading multipart request " + uri, e);
         }
     }
@@ -80,4 +94,34 @@ public class MultipartParser {
         return null;
     }
 
+    /** A {@link java.io.ByteArrayOutputStream} that limits the number of bytes written to it */
+    private static class LimitedOutputStream extends ByteArrayOutputStream {
+
+        private long remaining;
+
+        /** Create a new OutputStream that can fit up to len bytes */
+        private LimitedOutputStream(long len) {
+            this.remaining = len;
+        }
+
+        @Override
+        public synchronized void write(int b) {
+            requireCapacity(1);
+            super.write(b);
+            remaining--;
+        }
+
+        @Override
+        public synchronized void write(byte[] b, int off, int len) {
+            requireCapacity(len);
+            super.write(b, off, len);
+            remaining -= len;
+        }
+
+        private void requireCapacity(int len) {
+            if (len > remaining) throw new IllegalArgumentException("Too many bytes to write");
+        }
+
+    }
+
 }
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParserTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParserTest.java
index 2c81b1a7fd8..12a0a00713c 100644
--- a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParserTest.java
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/restapi/application/MultipartParserTest.java
@@ -11,12 +11,12 @@ import org.junit.Test;
 
 import java.io.ByteArrayInputStream;
 import java.net.URI;
-import java.net.URISyntaxException;
 import java.nio.charset.StandardCharsets;
 import java.util.Map;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
 
 /**
  * @author bratseth
@@ -24,7 +24,7 @@ import static org.junit.Assert.assertTrue;
 public class MultipartParserTest {
 
     @Test
-    public void multipartParserTest() throws URISyntaxException {
+    public void parser() {
         String data =
                 "Content-Type: multipart/form-data; boundary=AaB03x\r\n" +
                 "\r\n" +
@@ -43,13 +43,7 @@ public class MultipartParserTest {
                 "\r\n" +
                 "... contents of file1.txt ...\r\n" +
                 "--AaB03x--\r\n";
-        ByteArrayInputStream dataStream = new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8));
-        HttpRequest request = HttpRequest.createRequest(new MockCurrentContainer(), 
-                                                        new URI("http://foo"), 
-                                                        com.yahoo.jdisc.http.HttpRequest.Method.POST, 
-                                                        dataStream);
-        request.getJDiscRequest().headers().put("Content-Type", "multipart/form-data; boundary=AaB03x");
-        Map<String, byte[]> parts = new MultipartParser().parse(request);
+        Map<String, byte[]> parts = parse(data, Long.MAX_VALUE);
         assertEquals(3, parts.size());
         assertTrue(parts.containsKey("submit-name"));
         assertTrue(parts.containsKey("submit-address"));
@@ -57,6 +51,41 @@ public class MultipartParserTest {
         assertEquals("Larry", new String(parts.get("submit-name"), StandardCharsets.UTF_8));
         assertEquals("... contents of file1.txt ...", new String(parts.get("files"), StandardCharsets.UTF_8));
     }
+
+    @Test
+    public void max_length() {
+        String part1 = "Larry";
+        String part2 = "House 1";
+        String data =
+                "Content-Type: multipart/form-data; boundary=AaB03x\r\n" +
+                "\r\n" +
+                "--AaB03x\r\n" +
+                "Content-Disposition: form-data; name=\"submit-name\"\r\n" +
+                "\r\n" +
+                part1 + "\r\n" +
+                "--AaB03x\r\n" +
+                "Content-Disposition: form-data; name=\"submit-address\"\r\n" +
+                "Content-Type: text/plain\r\n" +
+                "\r\n" +
+                part2 + "\r\n" +
+                "--AaB03x--\r\n";
+        parse(data, part1.length() + part2.length());
+        try {
+            parse(data, part1.length() + part2.length() - 1);
+            fail("Expected exception");
+        } catch (IllegalArgumentException ignored) {
+        }
+    }
+
+    private Map<String, byte[]> parse(String data, long maxLength) {
+        ByteArrayInputStream dataStream = new ByteArrayInputStream(data.getBytes(StandardCharsets.UTF_8));
+        HttpRequest request = HttpRequest.createRequest(new MockCurrentContainer(),
+                                                        URI.create("http://foo"),
+                                                        com.yahoo.jdisc.http.HttpRequest.Method.POST,
+                                                        dataStream);
+        request.getJDiscRequest().headers().put("Content-Type", "multipart/form-data; boundary=AaB03x");
+        return new MultipartParser(maxLength).parse(request);
+    }
     
     private static class MockCurrentContainer implements CurrentContainer {