From 71834e1018bb69a68fb7f4afe29e779e0f175fbc Mon Sep 17 00:00:00 2001
From: BodgeMaster <>
Date: Fri, 14 Oct 2022 18:08:49 +0200
Subject: [PATCH] lib/nbt: Build an end tag object into compound tag objects to
 prevent creation of endless compounds

---
 src/lib/nbt.cpp | 26 ++++++++++++++++++++++----
 src/lib/nbt.hpp |  2 ++
 2 files changed, 24 insertions(+), 4 deletions(-)

diff --git a/src/lib/nbt.cpp b/src/lib/nbt.cpp
index 054d036..e676ca0 100644
--- a/src/lib/nbt.cpp
+++ b/src/lib/nbt.cpp
@@ -940,23 +940,27 @@ namespace NBT {
 
         Compound::Compound() {
             this->type = TagType::COMPOUND;
+            this->endPointer = new End();
         }
 
         Compound::Compound(tiny_utf8::string name) {
             this->type = TagType::COMPOUND;
             this->name = name;
+            this->endPointer = new End();
         }
 
         Compound::Compound(tiny_utf8::string name, std::vector<Generic*> data) {
             this->type = TagType::COMPOUND;
             this->name = name;
             this->tags = data;
+            this->endPointer = new End();
         }
 
         Compound::~Compound() {
             for (uint64_t i=0; i<this->tags.size(); i++) {
                 delete this->tags.at(i);
             }
+            delete this->endPointer;
         }
 
         ErrorOrVoid Compound::serializeWithoutHeader(std::vector<uint8_t>* rawData) {
@@ -966,20 +970,29 @@ namespace NBT {
                     return result;
                 }
             }
+            this->endPointer->serialize(rawData);
 
             return ErrorOrVoid();
         }
 
         ErrorOr<Generic*> Compound::getElementPointer(uint64_t position) {
-            if (position >= this->tags.size()) {
+            if (position > this->tags.size()) {
                 return ErrorOr<Generic*>(true, ErrorCodes::OUT_OF_RANGE);
             }
 
+            if (position == this->tags.size()) {
+                return this->endPointer;
+            }
+
             return ErrorOr<Generic*>(this->tags.at(position));
         }
 
         ErrorOrVoid Compound::setElementPointerAt(uint64_t position, Generic* pointer) {
-            if (position >= this->tags.size()) {
+            // built-in end tag
+            if (position == this->tags.size()) {
+                return ErrorOrVoid(true, ErrorCodes::NOT_ALLOWED);
+            }
+            if (position > this->tags.size()) {
                 return ErrorOrVoid(true, ErrorCodes::OUT_OF_RANGE);
             }
 
@@ -994,7 +1007,11 @@ namespace NBT {
         }
 
         ErrorOrVoid Compound::deleteElement(uint64_t position) {
-            if (position >= this->tags.size()) {
+            // built-in end tag
+            if (position == this->tags.size()) {
+                return ErrorOrVoid(true, ErrorCodes::NOT_ALLOWED);
+            }
+            if (position > this->tags.size()) {
                 return ErrorOrVoid(true, ErrorCodes::OUT_OF_RANGE);
             }
 
@@ -1005,7 +1022,8 @@ namespace NBT {
         }
 
         uint64_t Compound::length() {
-            return this->tags.size();
+            // account for built-in end tag
+            return this->tags.size()+1;
         }
 
         Int32Array::Int32Array() {
diff --git a/src/lib/nbt.hpp b/src/lib/nbt.hpp
index 5ab4ac8..88c8cd3 100644
--- a/src/lib/nbt.hpp
+++ b/src/lib/nbt.hpp
@@ -244,6 +244,8 @@ namespace NBT {
         class Compound: public Generic {
             private:
                 std::vector<Generic*> tags;
+                // built-in end tag
+                End* endPointer;
             public:
                 Compound();
                 Compound(tiny_utf8::string name);