RecordParser: Limit the record size (#2635)

* Simpler implementation for the max size limit. fixes #2560

Signed-off-by: Richard Fuchshuber <richardfuch@gmail.com>

* Using IllegalStateException

Signed-off-by: Richard Fuchshuber <richardfuch@gmail.com>

src/main/java/io/vertx/core/parsetools/RecordParser.java

@@ -11,6 +11,7 @@
 
 package io.vertx.core.parsetools;
 
+import io.vertx.codegen.annotations.Fluent;
 import io.vertx.codegen.annotations.VertxGen;
 import io.vertx.core.Handler;
 import io.vertx.core.buffer.Buffer;
@@ -192,6 +193,18 @@ static RecordParser newFixed(int size, ReadStream<Buffer> stream) {
    */
   void fixedSizeMode(int size);
 
+  /**
+   * Set the maximum allowed size for a record when using the delimited mode.
+   * The delimiter itself does not count for the record size.
+   * <p>
+   * If a record is longer than specified, an {@link IllegalStateException} will be thrown.
+   *
+   * @param size the maximum record size
+   * @return  a reference to this, so the API can be used fluently
+   */
+  @Fluent
+  RecordParser maxRecordSize(int size);
+
   /**
    * This method is called to provide the parser with data.
    *

src/main/java/io/vertx/core/parsetools/impl/RecordParserImpl.java

@@ -35,6 +35,7 @@ public class RecordParserImpl implements RecordParser {
   private byte[] delim;
   private int recordSize;
   private Handler<Buffer> output;
+  private int maxRecordSize;
   private Handler<Void> endHandler;
   private Handler<Throwable> exceptionHandler;
 
@@ -152,6 +153,21 @@ public void fixedSizeMode(int size) {
     reset = true;
   }
 
+  /**
+   * Set the maximum allowed size for a record when using the delimited mode.
+   * The delimiter itself does not count for the record size.
+   * <p>
+   * If a record is longer than specified, an {@link IllegalStateException} will be thrown.
+   *
+   * @param size the maximum record size
+   * @return  a reference to this, so the API can be used fluently
+   */
+  public RecordParser maxRecordSize(int size) {
+    Arguments.require(size > 0, "Size must be > 0");
+    maxRecordSize = size;
+    return this;
+  }
+
   private void handleParsing() {
     int len = buff.length();
     do {
@@ -217,6 +233,14 @@ public void handle(Buffer buffer) {
       buff.appendBuffer(buffer);
     }
     handleParsing();
+    if (buff != null && maxRecordSize > 0 && buff.length() > maxRecordSize) {
+      IllegalStateException ex = new IllegalStateException("The current record is too long");
+      if (exceptionHandler != null) {
+        exceptionHandler.handle(ex);
+      } else {
+        throw ex;
+      }
+    }
   }
 
   private void end() {

src/test/java/io/vertx/test/core/RecordParserTest.java

@@ -25,10 +25,9 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static io.vertx.test.core.TestUtils.assertIllegalArgumentException;
 import static io.vertx.test.core.TestUtils.assertNullPointerException;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.*;
 
 /**
  * @author <a href="http://tfox.org">Tim Fox</a>
@@ -45,6 +44,7 @@ public void testIllegalArguments() throws Exception {
     assertNullPointerException(() -> parser.setOutput(null));
     assertNullPointerException(() -> parser.delimitedMode((Buffer) null));
     assertNullPointerException(() -> parser.delimitedMode((String) null));
+    assertIllegalArgumentException(() -> parser.maxRecordSize(-1));
   }
 
   @Test
@@ -278,6 +278,29 @@ public void testSpreadDelimiter() {
       new Integer[] { 18 }, Buffer.buffer("start-ab-c-ddd"));
   }
 
+  @Test
+  public void testDelimitedMaxRecordSize() {
+    doTestDelimitedMaxRecordSize(Buffer.buffer("ABCD\nEFGH\n"), Buffer.buffer("\n"), new Integer[] { 2 },
+      4, null, Buffer.buffer("ABCD"), Buffer.buffer("EFGH"));
+    doTestDelimitedMaxRecordSize(Buffer.buffer("A\nBC10\nDEFGHIJKLM\n"), Buffer.buffer("\n"), new Integer[] { 2 },
+      10, null, Buffer.buffer("A"), Buffer.buffer("BC10"), Buffer.buffer("DEFGHIJKLM"));
+    doTestDelimitedMaxRecordSize(Buffer.buffer("AB\nC\n\nDEFG\n\n"), Buffer.buffer("\n\n"), new Integer[] { 2 },
+      4, null, Buffer.buffer("AB\nC"), Buffer.buffer("DEFG"));
+    doTestDelimitedMaxRecordSize(Buffer.buffer("AB--C---D-"), Buffer.buffer("-"), new Integer[] { 3 },
+      2, null, Buffer.buffer("AB"), Buffer.buffer(""), Buffer.buffer("C"), Buffer.buffer(""), Buffer.buffer(""), Buffer.buffer("D"));
+    try {
+      doTestDelimitedMaxRecordSize(Buffer.buffer("ABCD--"), Buffer.buffer("--"), new Integer[] { 2 },
+        3, null, Buffer.buffer());
+      fail("should throw exception");
+    }
+    catch (IllegalStateException ex) { /*OK*/ }
+    AtomicBoolean handled = new AtomicBoolean();
+    Handler<Throwable> exHandler = throwable -> handled.set(true);
+    doTestDelimitedMaxRecordSize(Buffer.buffer("ABCD--"), Buffer.buffer("--"), new Integer[] { 2 },
+      3, exHandler, Buffer.buffer("ABCD"));
+    assertTrue(handled.get());
+  }
+
   @Test
   public void testWrapReadStream() {
     AtomicBoolean paused = new AtomicBoolean();
@@ -334,4 +357,20 @@ public ReadStream<Buffer> endHandler(Handler<Void> handler) {
     assertEquals(Arrays.asList("first", "second"), records);
     assertEquals(1, ends.get());
   }
+
+  private void doTestDelimitedMaxRecordSize(final Buffer input, Buffer delim, Integer[] chunkSizes, int maxRecordSize,
+                                            Handler<Throwable> exHandler, final Buffer... expected) {
+    final Buffer[] results = new Buffer[expected.length];
+    Handler<Buffer> out = new Handler<Buffer>() {
+      int pos;
+      public void handle(Buffer buff) {
+        results[pos++] = buff;
+      }
+    };
+    RecordParser parser = RecordParser.newDelimited(delim, out);
+    parser.maxRecordSize(maxRecordSize);
+    parser.exceptionHandler(exHandler);
+    feedChunks(input, parser, chunkSizes);
+    checkResults(expected, results);
+  }
 }