[+] IProtocolPiece.GetNextPiece

[+] IProtocolPiece.ReallocateDrainBuffer
[*] Harden protocol piece use after destroy

Include/Aurora/IO/Protocol/IProtocolPiece.hpp

@@ -11,6 +11,19 @@ namespace Aurora::IO::Protocol
 {
     struct IProtocolPiece
     {
+        /**
+         * @brief 
+         * @return 
+        */
+        virtual AuSPtr<IProtocolPiece> GetNextPiece() = 0;
+
+        /**
+         * @brief 
+         * @param uOutputLength 
+         * @return 
+        */
+        virtual bool ReallocateDrainBuffer(AuUInt32 uOutputLength) = 0;
+
         /**
          * @brief Removes this piece from the stack
         */

Source/IO/Protocol/AuProtocolPiece.cpp

@@ -14,6 +14,11 @@ namespace Aurora::IO::Protocol
 {
     AuSPtr<IStreamWriter> ProtocolPiece::ToNextWriter()
     {
+        if (!this->pParent)
+        {
+            return {};
+        }
+
         if (this->pParent->pTopPiece.get() == this)
         {
             if (this->pParent->pDrainBuffer)
@@ -32,8 +37,36 @@ namespace Aurora::IO::Protocol
         return this->pWriteInteface ? this->pWriteInteface->GetStreamWriter() : AuSPtr<IStreamWriter> {};
     }
 
+    AuSPtr<IProtocolPiece> ProtocolPiece::GetNextPiece()
+    {
+        return this->pNext;
+    }
+
+    bool ProtocolPiece::ReallocateDrainBuffer(AuUInt32 uOutputLength)
+    {
+        if (!this->pParent)
+        {
+            return false;
+        }
+
+        if (!this->pNext)
+        {
+            if (this->pParent->pDrainBuffer)
+            {
+                return false;
+            }
+        }
+
+        return this->outputBuffer.Resize(uOutputLength);
+    }
+
     void ProtocolPiece::Remove()
     {
+        if (!this->pParent)
+        {
+            return;
+        }
+
         auto &pBottomPiece = this->pParent->pBottomPiece;
         
         // fix chain
@@ -66,6 +99,8 @@ namespace Aurora::IO::Protocol
         {
             this->pParent->pBottomPiece = this->pNext;
         }
+
+        this->pParent = nullptr;
     }
 
     AuSPtr<Memory::ByteBuffer> ProtocolPiece::GetNextPieceBuffer()

Source/IO/Protocol/AuProtocolPiece.hpp

@@ -23,6 +23,9 @@ namespace Aurora::IO::Protocol
         AuSPtr<IStreamWriter> pOuputWriter;
         bool bMultipleTick {};
 
+
+        AuSPtr<IProtocolPiece> GetNextPiece() override;
+        bool ReallocateDrainBuffer(AuUInt32 uOutputLength) override;
         void Remove() override;
         AuSPtr<IStreamWriter> ToNextWriter() override;
         AuSPtr<IStreamWriter> ToInputWriter() override;