From 5b20d95920422fe77835d7b94ac0d0f843ace5d4 Mon Sep 17 00:00:00 2001
From: CraftBukkit/Spigot <noreply+git-craftbukkit@papermc.io>
Date: Sun, 27 Jul 2014 20:46:04 +1000
Subject: [PATCH] Apply NBTReadLimiter to more things.

By: md_5 <git@md-5.net>
---
 .../net/minecraft/nbt/NbtIo.java.patch        | 13 +++++++
 .../main/java/org/spigotmc/LimitStream.java   | 39 +++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 paper-server/src/main/java/org/spigotmc/LimitStream.java

diff --git a/paper-server/patches/sources/net/minecraft/nbt/NbtIo.java.patch b/paper-server/patches/sources/net/minecraft/nbt/NbtIo.java.patch
index 661faf49b1..f6dc9e632c 100644
--- a/paper-server/patches/sources/net/minecraft/nbt/NbtIo.java.patch
+++ b/paper-server/patches/sources/net/minecraft/nbt/NbtIo.java.patch
@@ -5,3 +5,16 @@
  package net.minecraft.nbt;
  
  import java.io.BufferedOutputStream;
+@@ -324,6 +325,12 @@
+     }
+ 
+     public static CompoundTag read(DataInput input, NbtAccounter tracker) throws IOException {
++        // Spigot start
++        if ( input instanceof io.netty.buffer.ByteBufInputStream )
++        {
++            input = new DataInputStream(new org.spigotmc.LimitStream((InputStream) input, tracker));
++        }
++        // Spigot end
+         Tag nbtbase = NbtIo.readUnnamedTag(input, tracker);
+ 
+         if (nbtbase instanceof CompoundTag) {
diff --git a/paper-server/src/main/java/org/spigotmc/LimitStream.java b/paper-server/src/main/java/org/spigotmc/LimitStream.java
new file mode 100644
index 0000000000..8de241a921
--- /dev/null
+++ b/paper-server/src/main/java/org/spigotmc/LimitStream.java
@@ -0,0 +1,39 @@
+package org.spigotmc;
+
+import java.io.FilterInputStream;
+import java.io.IOException;
+import java.io.InputStream;
+import net.minecraft.nbt.NbtAccounter;
+
+public class LimitStream extends FilterInputStream
+{
+
+    private final NbtAccounter limit;
+
+    public LimitStream(InputStream is, NbtAccounter limit)
+    {
+        super( is );
+        this.limit = limit;
+    }
+
+    @Override
+    public int read() throws IOException
+    {
+        this.limit.accountBytes( 1 );
+        return super.read();
+    }
+
+    @Override
+    public int read(byte[] b) throws IOException
+    {
+        this.limit.accountBytes( b.length );
+        return super.read( b );
+    }
+
+    @Override
+    public int read(byte[] b, int off, int len) throws IOException
+    {
+        this.limit.accountBytes( len );
+        return super.read( b, off, len );
+    }
+}