From 7be73f86d4c1bbc99a3764dedad1231f8dba92c6 Mon Sep 17 00:00:00 2001
From: BodgeMaster <>
Date: Fri, 14 Oct 2022 18:30:39 +0200
Subject: [PATCH] lib/nbt: Prevent addition of additional end tags to compounds

---
 src/lib/nbt.cpp | 19 +++++++++++++++----
 src/lib/nbt.hpp |  2 +-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/src/lib/nbt.cpp b/src/lib/nbt.cpp
index e676ca0..6214b02 100644
--- a/src/lib/nbt.cpp
+++ b/src/lib/nbt.cpp
@@ -988,9 +988,16 @@ namespace NBT {
         }
 
         ErrorOrVoid Compound::setElementPointerAt(uint64_t position, Generic* pointer) {
-            // built-in end tag
-            if (position == this->tags.size()) {
-                return ErrorOrVoid(true, ErrorCodes::NOT_ALLOWED);
+            if (position == this->tags.size() || pointer->getTagType() == TagType::END) {
+                if (position == this->tags.size() && pointer->getTagType() == TagType::END) {
+                    delete pointer;
+                    // do nothing, already have one of those
+                } else {
+                    delete pointer;
+                    // End tags may only go at the end and
+                    // the end may only hold an end tag.
+                    return ErrorOrVoid(true, ErrorCodes::NOT_ALLOWED);
+                }
             }
             if (position > this->tags.size()) {
                 return ErrorOrVoid(true, ErrorCodes::OUT_OF_RANGE);
@@ -1002,8 +1009,12 @@ namespace NBT {
             return ErrorOrVoid();
         }
 
-        void Compound::appendPointer(Generic* pointer) {
+        ErrorOrVoid Compound::appendPointer(Generic* pointer) {
+            if (pointer->getTagType() == TagType::END) {
+                return ErrorOrVoid(true, ErrorCodes::NOT_ALLOWED);
+            }
             this->tags.push_back(pointer);
+            return ErrorOrVoid();
         }
 
         ErrorOrVoid Compound::deleteElement(uint64_t position) {
diff --git a/src/lib/nbt.hpp b/src/lib/nbt.hpp
index 88c8cd3..0907671 100644
--- a/src/lib/nbt.hpp
+++ b/src/lib/nbt.hpp
@@ -257,7 +257,7 @@ namespace NBT {
 
                 ErrorOr<Generic*> getElementPointer(uint64_t position);
                 ErrorOrVoid setElementPointerAt(uint64_t position, Generic* pointer);
-                void appendPointer(Generic* pointer);
+                ErrorOrVoid appendPointer(Generic* pointer);
                 ErrorOrVoid deleteElement(uint64_t position);
                 uint64_t length();
         };