Add limiter for max bytes and nesting level

pom.xml

@@ -5,7 +5,7 @@
 
     <groupId>com.viaversion</groupId>
     <artifactId>opennbt</artifactId>
-    <version>2.0</version>
+    <version>2.1</version>
     <packaging>jar</packaging>
 
     <name>OpenNBT</name>

src/main/java/com/github/steveice10/opennbt/NBTIO.java

@@ -1,6 +1,7 @@
 package com.github.steveice10.opennbt;
 
 import com.github.steveice10.opennbt.tag.builtin.CompoundTag;
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 
 import java.io.DataInput;
 import java.io.DataInputStream;
@@ -151,7 +152,18 @@ public class NBTIO {
      * @throws java.io.IOException If an I/O error occurs.
      */
     public static CompoundTag readTag(InputStream in) throws IOException {
-        return readTag(in, false);
+        return readTag(in, TagLimiter.noop());
+    }
+
+    /**
+     * Reads a big endian NBT tag.
+     *
+     * @param in Input stream to read from.
+     * @return The read tag, or null if the tag is an end tag.
+     * @throws java.io.IOException If an I/O error occurs.
+     */
+    public static CompoundTag readTag(InputStream in, TagLimiter tagLimiter) throws IOException {
+        return readTag((DataInput) new DataInputStream(in), tagLimiter);
     }
 
     /**
@@ -174,6 +186,18 @@ public class NBTIO {
      * @throws java.io.IOException If an I/O error occurs.
      */
     public static CompoundTag readTag(DataInput in) throws IOException {
+        return readTag(in, TagLimiter.noop());
+    }
+
+    /**
+     * Reads an NBT tag.
+     *
+     * @param in Data input to read from.
+     * @param tagLimiter taglimiter
+     * @return The read tag, or null if the tag is an end tag.
+     * @throws java.io.IOException If an I/O error occurs.
+     */
+    public static CompoundTag readTag(DataInput in, TagLimiter tagLimiter) throws IOException {
         int id = in.readByte();
         if(id != CompoundTag.ID) {
             throw new IOException(String.format("Expected root tag to be a CompoundTag, was %s", id));
@@ -183,7 +207,7 @@ public class NBTIO {
         in.skipBytes(in.readUnsignedShort());
 
         CompoundTag tag = new CompoundTag();
-        tag.read(in);
+        tag.read(in, tagLimiter);
         return tag;
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/ByteArrayTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -77,8 +79,10 @@ public class ByteArrayTag extends Tag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countInt();
         this.value = new byte[in.readInt()];
+        tagLimiter.countBytes(this.value.length);
         in.readFully(this.value);
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/ByteTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -46,7 +48,8 @@ public class ByteTag extends NumberTag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countByte();
         this.value = in.readByte();
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/CompoundTag.java

@@ -2,6 +2,7 @@ package com.github.steveice10.opennbt.tag.builtin;
 
 import com.github.steveice10.opennbt.tag.TagCreateException;
 import com.github.steveice10.opennbt.tag.TagRegistry;
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 import com.google.common.base.Preconditions;
 import org.jetbrains.annotations.Nullable;
 
@@ -179,10 +180,13 @@ public class CompoundTag extends Tag implements Iterable<Entry<String, Tag>> {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
         try {
+            tagLimiter.checkLevel(nestingLevel);
+            int newNestingLevel = nestingLevel + 1;
             int id;
             while (true) {
+                tagLimiter.countByte();
                 id = in.readByte();
                 if (id == 0) {
                     // End tag
@@ -190,8 +194,10 @@ public class CompoundTag extends Tag implements Iterable<Entry<String, Tag>> {
                 }
 
                 String name = in.readUTF();
+                tagLimiter.countBytes(2 * name.length());
+
                 Tag tag = TagRegistry.createInstance(id);
-                tag.read(in);
+                tag.read(in, tagLimiter, newNestingLevel);
                 this.value.put(name, tag);
             }
         } catch(TagCreateException e) {

src/main/java/com/github/steveice10/opennbt/tag/builtin/DoubleTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -46,7 +48,8 @@ public class DoubleTag extends NumberTag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countDouble();
         this.value = in.readDouble();
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/FloatTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -46,7 +48,8 @@ public class FloatTag extends NumberTag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countFloat();
         this.value = in.readFloat();
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/IntArrayTag.java

@@ -1,5 +1,6 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 import com.google.common.base.Preconditions;
 
 import java.io.DataInput;
@@ -77,8 +78,10 @@ public class IntArrayTag extends Tag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countInt();
         this.value = new int[in.readInt()];
+        tagLimiter.countBytes(4 * this.value.length);
         for(int index = 0; index < this.value.length; index++) {
             this.value[index] = in.readInt();
         }

src/main/java/com/github/steveice10/opennbt/tag/builtin/IntTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -46,7 +48,8 @@ public class IntTag extends NumberTag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countInt();
         this.value = in.readInt();
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/ListTag.java

@@ -2,6 +2,7 @@ package com.github.steveice10.opennbt.tag.builtin;
 
 import com.github.steveice10.opennbt.tag.TagCreateException;
 import com.github.steveice10.opennbt.tag.TagRegistry;
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 import com.google.common.base.Preconditions;
 import org.jetbrains.annotations.Nullable;
 
@@ -138,9 +139,11 @@ public class ListTag extends Tag implements Iterable<Tag> {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
         this.type = null;
 
+        tagLimiter.checkLevel(nestingLevel);
+        tagLimiter.countBytes(1 + 4);
         int id = in.readByte();
         if(id != 0) {
             this.type = TagRegistry.getClassFor(id);
@@ -150,6 +153,7 @@ public class ListTag extends Tag implements Iterable<Tag> {
         }
 
         int count = in.readInt();
+        int newNestingLevel = nestingLevel + 1;
         for(int index = 0; index < count; index++) {
             Tag tag;
             try {
@@ -158,7 +162,7 @@ public class ListTag extends Tag implements Iterable<Tag> {
                 throw new IOException("Failed to create tag.", e);
             }
 
-            tag.read(in);
+            tag.read(in, tagLimiter, newNestingLevel);
             this.add(tag);
         }
     }

src/main/java/com/github/steveice10/opennbt/tag/builtin/LongArrayTag.java

@@ -1,5 +1,6 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 import com.google.common.base.Preconditions;
 
 import java.io.DataInput;
@@ -77,8 +78,10 @@ public class LongArrayTag extends Tag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countInt();
         this.value = new long[in.readInt()];
+        tagLimiter.countBytes(8 * this.value.length);
         for(int index = 0; index < this.value.length; index++) {
             this.value[index] = in.readLong();
         }

src/main/java/com/github/steveice10/opennbt/tag/builtin/LongTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -46,7 +48,8 @@ public class LongTag extends NumberTag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countLong();
         this.value = in.readLong();
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/ShortTag.java

@@ -1,5 +1,7 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
+
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -46,7 +48,8 @@ public class ShortTag extends NumberTag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
+        tagLimiter.countShort();
         this.value = in.readShort();
     }
 

src/main/java/com/github/steveice10/opennbt/tag/builtin/StringTag.java

@@ -1,5 +1,6 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 import com.google.common.base.Preconditions;
 
 import java.io.DataInput;
@@ -46,8 +47,9 @@ public class StringTag extends Tag {
     }
 
     @Override
-    public void read(DataInput in) throws IOException {
+    public void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException {
         this.value = in.readUTF();
+        tagLimiter.countBytes(2 * value.length()); // More or less, ignoring the length reading
     }
 
     @Override

src/main/java/com/github/steveice10/opennbt/tag/builtin/Tag.java

@@ -1,5 +1,6 @@
 package com.github.steveice10.opennbt.tag.builtin;
 
+import com.github.steveice10.opennbt.tag.limiter.TagLimiter;
 import java.io.DataInput;
 import java.io.DataOutput;
 import java.io.IOException;
@@ -25,7 +26,30 @@ public abstract class Tag implements Cloneable {
      * @param in Stream to write to.
      * @throws java.io.IOException If an I/O error occurs.
      */
-    public abstract void read(DataInput in) throws IOException;
+    public final void read(DataInput in) throws IOException {
+        this.read(in, TagLimiter.noop(), 0);
+    }
+
+    /**
+     * Reads this tag from an input stream.
+     *
+     * @param in Stream to write to.
+     * @param tagLimiter taglimiter
+     * @throws java.io.IOException If an I/O error occurs.
+     */
+    public final void read(DataInput in, TagLimiter tagLimiter) throws IOException {
+        this.read(in, tagLimiter, 0);
+    }
+
+    /**
+     * Reads this tag from an input stream.
+     *
+     * @param in Stream to write to.
+     * @param tagLimiter taglimiter
+     * @param nestingLevel current level of nesting
+     * @throws java.io.IOException If an I/O error occurs.
+     */
+    public abstract void read(DataInput in, TagLimiter tagLimiter, int nestingLevel) throws IOException;
 
     /**
      * Writes this tag to an output stream.

src/main/java/com/github/steveice10/opennbt/tag/limiter/NoopTagLimiter.java

@@ -0,0 +1,29 @@
+package com.github.steveice10.opennbt.tag.limiter;
+
+final class NoopTagLimiter implements TagLimiter {
+
+    static final TagLimiter INSTANCE = new NoopTagLimiter();
+
+    @Override
+    public void countBytes(int bytes) {
+    }
+
+    @Override
+    public void checkLevel(int nestedLevel) {
+    }
+
+    @Override
+    public int maxBytes() {
+        return Integer.MAX_VALUE;
+    }
+
+    @Override
+    public int maxLevels() {
+        return Integer.MAX_VALUE;
+    }
+
+    @Override
+    public int bytes() {
+        return 0;
+    }
+}

src/main/java/com/github/steveice10/opennbt/tag/limiter/TagLimiter.java

@@ -0,0 +1,85 @@
+package com.github.steveice10.opennbt.tag.limiter;
+
+public interface TagLimiter {
+
+    /**
+     * Returns a new tag limiter with the given max bytes and nesting levels.
+     *
+     * @param maxBytes  max amount of bytes to be read before an exception is thrown when reading nbt
+     * @param maxLevels max levels of nesting before an exception is thrown when reading nbt
+     * @return tag limiter
+     */
+    static TagLimiter create(int maxBytes, int maxLevels) {
+        return new TagLimiterImpl(maxBytes, maxLevels);
+    }
+
+    /**
+     * Returns a noop tag limiter.
+     *
+     * @return noop tag limiter
+     */
+    static TagLimiter noop() {
+        return NoopTagLimiter.INSTANCE;
+    }
+
+    /**
+     * Counts the given number of bytes and throws an exception if the max bytes count is exceeded.
+     *
+     * @param bytes bytes to count
+     * @throws IllegalArgumentException if max bytes count is exceeded
+     */
+    void countBytes(int bytes);
+
+    /**
+     * Checks the current level of nesting and throws an exception if it exceeds the max levels.
+     *
+     * @param nestedLevel current level of nesting
+     * @throws IllegalArgumentException if max level count is exceeded
+     */
+    void checkLevel(int nestedLevel);
+
+    default void countByte() {
+        this.countBytes(1);
+    }
+
+    default void countShort() {
+        this.countBytes(2);
+    }
+
+    default void countInt() {
+        this.countBytes(4);
+    }
+
+    default void countFloat() {
+        this.countBytes(4);
+    }
+
+    default void countLong() {
+        this.countBytes(8);
+    }
+
+    default void countDouble() {
+        this.countBytes(8);
+    }
+
+    /**
+     * Returns the max number of bytes to be read before an exception is thrown when reading nbt.
+     *
+     * @return max bytes
+     */
+    int maxBytes();
+
+    /**
+     * Returns the max number of levels of nesting before an exception is thrown when reading nbt.
+     *
+     * @return max nesting levels
+     */
+    int maxLevels();
+
+    /**
+     * Returns the amount of currently read bytes.
+     *
+     * @return currently read bytes
+     */
+    int bytes();
+}

src/main/java/com/github/steveice10/opennbt/tag/limiter/TagLimiterImpl.java

@@ -0,0 +1,43 @@
+package com.github.steveice10.opennbt.tag.limiter;
+
+final class TagLimiterImpl implements TagLimiter {
+
+    private final int maxBytes;
+    private final int maxLevels;
+    private int bytes;
+
+    TagLimiterImpl(int maxBytes, int maxLevels) {
+        this.maxBytes = maxBytes;
+        this.maxLevels = maxLevels;
+    }
+
+    @Override
+    public void countBytes(int bytes) {
+        this.bytes += bytes;
+        if (this.bytes >= maxBytes) {
+            throw new IllegalArgumentException("NBT data larger than expected (capped at " + this.maxBytes + ")");
+        }
+    }
+
+    @Override
+    public void checkLevel(int nestedLevel) {
+        if (nestedLevel >= this.maxLevels) {
+            throw new IllegalArgumentException("Nesting level higher than expected (capped at " + this.maxLevels + ")");
+        }
+    }
+
+    @Override
+    public int maxBytes() {
+        return maxBytes;
+    }
+
+    @Override
+    public int maxLevels() {
+        return maxLevels;
+    }
+
+    @Override
+    public int bytes() {
+        return bytes;
+    }
+}