From 0a4886e7ab1c396d102b843d9221eae22b07a1af Mon Sep 17 00:00:00 2001
From: Jan <jan@laupetin.net>
Date: Thu, 18 Mar 2021 11:48:41 +0100
Subject: [PATCH] Make sure zone memory can allocate blocks that are larger
 than its default maximum size and throw an error if allocation failed

---
 src/ZoneWriting/Writing/InMemoryZoneData.cpp | 17 ++++++++++++-----
 src/ZoneWriting/Writing/ZoneWriter.cpp       |  6 ++++++
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/src/ZoneWriting/Writing/InMemoryZoneData.cpp b/src/ZoneWriting/Writing/InMemoryZoneData.cpp
index 35d8df1c..79a0a7ef 100644
--- a/src/ZoneWriting/Writing/InMemoryZoneData.cpp
+++ b/src/ZoneWriting/Writing/InMemoryZoneData.cpp
@@ -1,6 +1,6 @@
 #include "InMemoryZoneData.h"
 
-#include <cassert>
+#include <stdexcept>
 
 InMemoryZoneData::InMemoryZoneData()
     : m_total_size(0)
@@ -12,15 +12,22 @@ InMemoryZoneData::MemoryBuffer::MemoryBuffer(const size_t size)
     : m_data(std::make_unique<char[]>(size)),
       m_size(0)
 {
+    if (!m_data)
+        throw std::runtime_error("Failed to allocate memory for memory buffer.");
 }
 
 void* InMemoryZoneData::GetBufferOfSize(const size_t size)
 {
-    assert(size <= BUFFER_SIZE);
-
-    if(m_buffers.back().m_size + size > BUFFER_SIZE)
+    if(size > BUFFER_SIZE)
     {
-        m_buffers.emplace_back(BUFFER_SIZE);
+        m_buffers.emplace_back(size);
+    }
+    else
+    {
+        if (m_buffers.back().m_size + size > BUFFER_SIZE)
+        {
+            m_buffers.emplace_back(BUFFER_SIZE);
+        }
     }
 
     auto& backBuffer = m_buffers.back();
diff --git a/src/ZoneWriting/Writing/ZoneWriter.cpp b/src/ZoneWriting/Writing/ZoneWriter.cpp
index 70e262ae..0f5a96f6 100644
--- a/src/ZoneWriting/Writing/ZoneWriter.cpp
+++ b/src/ZoneWriting/Writing/ZoneWriter.cpp
@@ -1,6 +1,7 @@
 #include "ZoneWriter.h"
 
 #include <iostream>
+#include <stdexcept>
 
 #include "WritingException.h"
 #include "WritingFileStream.h"
@@ -74,6 +75,11 @@ bool ZoneWriter::WriteZone(std::ostream& stream)
         std::cout << "Writing fastfile failed: " << e.Message() << std::endl;
         return false;
     }
+    catch (std::runtime_error& e)
+    {
+        std::cout << "Writing fastfile failed: " << e.what() << std::endl;
+        return false;
+    }
 
     endStream->Flush();