Set the limit for incoming data to prevent a CVE report in NIST

5 files changed, 35 insertions, 10 deletions

diff --git a/src/main/java/org/yaml/snakeyaml/LoaderOptions.java b/src/main/java/org/yaml/snakeyaml/LoaderOptions.java
index 1f6b81aa..bcb97135 100644
--- a/src/main/java/org/yaml/snakeyaml/LoaderOptions.java
+++ b/src/main/java/org/yaml/snakeyaml/LoaderOptions.java
@@ -23,6 +23,7 @@ public class LoaderOptions {
   private boolean processComments = false;
   private boolean enumCaseSensitive = true;
   private int nestingDepthLimit = 50;
+  private int codePointLimit = 100 * 1024;
 
   public boolean isAllowDuplicateKeys() {
     return allowDuplicateKeys;
@@ -130,4 +131,17 @@ public class LoaderOptions {
   public void setNestingDepthLimit(int nestingDepthLimit) {
     this.nestingDepthLimit = nestingDepthLimit;
   }
+
+  public int getCodePointLimit() {
+    return codePointLimit;
+  }
+
+  /**
+   * The max amount of code points in the input YAML document. Please be aware that byte limit
+   * depends on the encoding.
+   * @param codePointLimit - the max allowed size of the YAML data
+   */
+  public void setCodePointLimit(int codePointLimit) {
+    this.codePointLimit = codePointLimit;
+  }
 }
diff --git a/src/main/java/org/yaml/snakeyaml/Yaml.java b/src/main/java/org/yaml/snakeyaml/Yaml.java
index 42a4c977..0cfd2bb2 100644
--- a/src/main/java/org/yaml/snakeyaml/Yaml.java
+++ b/src/main/java/org/yaml/snakeyaml/Yaml.java
@@ -471,7 +471,7 @@ public class Yaml {
   }
 
   private Object loadFromReader(StreamReader sreader, Class<?> type) {
-    Composer composer = new Composer(new ParserImpl(sreader, loadingConfig.isProcessComments()),
+    Composer composer = new Composer(new ParserImpl(sreader, loadingConfig),
         resolver, loadingConfig);
     constructor.setComposer(composer);
     return constructor.getSingleData(type);
diff --git a/src/main/java/org/yaml/snakeyaml/parser/ParserImpl.java b/src/main/java/org/yaml/snakeyaml/parser/ParserImpl.java
index a359aa83..3a8eac28 100644
--- a/src/main/java/org/yaml/snakeyaml/parser/ParserImpl.java
+++ b/src/main/java/org/yaml/snakeyaml/parser/ParserImpl.java
@@ -19,6 +19,7 @@ import java.util.List;
 import java.util.Map;
 import org.yaml.snakeyaml.DumperOptions;
 import org.yaml.snakeyaml.DumperOptions.Version;
+import org.yaml.snakeyaml.LoaderOptions;
 import org.yaml.snakeyaml.comments.CommentType;
 import org.yaml.snakeyaml.error.Mark;
 import org.yaml.snakeyaml.error.YAMLException;
@@ -135,10 +136,15 @@ public class ParserImpl implements Parser {
     this(new ScannerImpl(reader));
   }
 
+  @Deprecated
   public ParserImpl(StreamReader reader, boolean parseComments) {
     this(new ScannerImpl(reader).setParseComments(parseComments));
   }
 
+  public ParserImpl(StreamReader reader, LoaderOptions options) {
+    this(new ScannerImpl(reader, options));
+  }
+
   public ParserImpl(Scanner scanner) {
     this.scanner = scanner;
     currentEvent = null;
diff --git a/src/main/java/org/yaml/snakeyaml/reader/StreamReader.java b/src/main/java/org/yaml/snakeyaml/reader/StreamReader.java
index b026f952..6799e76e 100644
--- a/src/main/java/org/yaml/snakeyaml/reader/StreamReader.java
+++ b/src/main/java/org/yaml/snakeyaml/reader/StreamReader.java
@@ -44,11 +44,6 @@ public class StreamReader {
    */
   private int pointer = 0;
   private boolean eof;
-  /**
-   * index is only required to implement 1024 key length restriction
-   * http://yaml.org/spec/1.1/#simple key/ It must count code points, but it counts characters (to
-   * be fixed)
-   */
   private int index = 0; // in code points
   private int line = 0;
   private int column = 0; // in code points
diff --git a/src/main/java/org/yaml/snakeyaml/scanner/ScannerImpl.java b/src/main/java/org/yaml/snakeyaml/scanner/ScannerImpl.java
index 51a723b6..1bd73569 100644
--- a/src/main/java/org/yaml/snakeyaml/scanner/ScannerImpl.java
+++ b/src/main/java/org/yaml/snakeyaml/scanner/ScannerImpl.java
@@ -23,6 +23,7 @@ import java.util.List;
 import java.util.Map;
 import java.util.regex.Pattern;
 import org.yaml.snakeyaml.DumperOptions;
+import org.yaml.snakeyaml.LoaderOptions;
 import org.yaml.snakeyaml.comments.CommentType;
 import org.yaml.snakeyaml.error.Mark;
 import org.yaml.snakeyaml.error.YAMLException;
@@ -178,6 +179,7 @@ public final class ScannerImpl implements Scanner {
 
   // A flag that indicates if comments should be parsed
   private boolean parseComments;
+  private final LoaderOptions loaderOptions;
 
   // Variables related to simple keys treatment. See PyYAML.
 
@@ -214,12 +216,17 @@ public final class ScannerImpl implements Scanner {
   private final Map<Integer, SimpleKey> possibleSimpleKeys;
 
   public ScannerImpl(StreamReader reader) {
+    this(reader, new LoaderOptions());
+  }
+
+  public ScannerImpl(StreamReader reader,  LoaderOptions options) {
     this.parseComments = false;
     this.reader = reader;
     this.tokens = new ArrayList<Token>(100);
     this.indents = new ArrayStack<Integer>(10);
     // The order in possibleSimpleKeys is kept for nextPossibleSimpleKey()
     this.possibleSimpleKeys = new LinkedHashMap<Integer, SimpleKey>();
+    this.loaderOptions = options;
     fetchStreamStart();// Add the STREAM-START token.
   }
 
@@ -319,6 +326,9 @@ public final class ScannerImpl implements Scanner {
    * Fetch one or more tokens from the StreamReader.
    */
   private void fetchMoreTokens() {
+    if (reader.getIndex() > loaderOptions.getCodePointLimit()) {
+      throw new YAMLException("The incoming YAML document exceeds the limit: " + loaderOptions.getCodePointLimit());
+    }
     // Eat whitespaces and process comments until we reach the next token.
     scanToNextToken();
     // Remove obsolete possible simple keys.
@@ -1175,7 +1185,7 @@ public final class ScannerImpl implements Scanner {
     // whitespace, then this is the start of a plain scalar.
     return Constant.NULL_BL_T_LINEBR.hasNo(c, "-?:,[]{}#&*!|>'\"%@`")
         || (Constant.NULL_BL_T_LINEBR.hasNo(reader.peek(1))
-            && (c == '-' || (this.flowLevel == 0 && "?:".indexOf(c) != -1)));
+        && (c == '-' || (this.flowLevel == 0 && "?:".indexOf(c) != -1)));
   }
 
   // Scanners.
@@ -1820,7 +1830,7 @@ public final class ScannerImpl implements Scanner {
       }
     }
     // Pass several results back together.
-    return new Object[] {chunks.toString(), maxIndent, endMark};
+    return new Object[]{chunks.toString(), maxIndent, endMark};
   }
 
   private Object[] scanBlockScalarBreaks(int indent) {
@@ -1850,7 +1860,7 @@ public final class ScannerImpl implements Scanner {
       }
     }
     // Return both the assembled intervening string and the end-mark.
-    return new Object[] {chunks.toString(), endMark};
+    return new Object[]{chunks.toString(), endMark};
   }
 
   /**
@@ -2039,7 +2049,7 @@ public final class ScannerImpl implements Scanner {
         c = reader.peek(length);
         if (Constant.NULL_BL_T_LINEBR.has(c)
             || (c == ':' && Constant.NULL_BL_T_LINEBR.has(reader.peek(length + 1),
-                flowLevel != 0 ? ",[]{}" : ""))
+            flowLevel != 0 ? ",[]{}" : ""))
             || (this.flowLevel != 0 && ",?[]{}".indexOf(c) != -1)) {
           break;
         }