From 81a597ab0d6c9c9bc834263bdbc07e9c144db0e1 Mon Sep 17 00:00:00 2001
From: Ted John <ted@brambles.org>
Date: Mon, 11 Jul 2016 19:09:16 +0100
Subject: [PATCH] append salt bytes inside encoded data chunk

---
 src/core/Guard.cpp              | 13 +++++
 src/core/Guard.hpp              |  1 +
 src/object/ObjectRepository.cpp | 84 +++++++++++++++++++--------------
 3 files changed, 63 insertions(+), 35 deletions(-)

diff --git a/src/core/Guard.cpp b/src/core/Guard.cpp
index 73c7ef2301..f88d871846 100644
--- a/src/core/Guard.cpp
+++ b/src/core/Guard.cpp
@@ -32,6 +32,19 @@ namespace Guard
             Console::Error::WriteLine(message);
         }
 
+#if DEBUG
+        Debug::Break();
+#endif
+        assert(false);
+    }
+
+    void Fail(const char * message)
+    {
+        if (message != nullptr)
+        {
+            Console::Error::WriteLine(message);
+        }
+
 #if DEBUG
         Debug::Break();
 #endif
diff --git a/src/core/Guard.hpp b/src/core/Guard.hpp
index 904aaf6f99..c4ba8cdefa 100644
--- a/src/core/Guard.hpp
+++ b/src/core/Guard.hpp
@@ -22,6 +22,7 @@
 namespace Guard
 {
     void Assert(bool expression, const char * message = nullptr);
+    void Fail(const char * message = nullptr);
 
     template<typename T>
     void ArgumentNotNull(T * argument, const char * message = nullptr)
diff --git a/src/object/ObjectRepository.cpp b/src/object/ObjectRepository.cpp
index 4570bd11b5..0a005b1b39 100644
--- a/src/object/ObjectRepository.cpp
+++ b/src/object/ObjectRepository.cpp
@@ -487,47 +487,68 @@ private:
         }
     }
 
-    static void SaveObject(const utf8 * path, const rct_object_entry * entry, const void * data, size_t dataSize)
+    static void SaveObject(const utf8 * path,
+                           const rct_object_entry * entry,
+                           const void * data, size_t dataSize,
+                           bool fixChecksum = true)
     {
-        uint8 objectType = entry->flags & 0x0F;
-
-        // Encode data
-        sawyercoding_chunk_header chunkHeader;
-        chunkHeader.encoding = object_entry_group_encoding[objectType];
-        chunkHeader.length = dataSize;
-        uint8 * encodedDataBuffer = Memory::Allocate<uint8>(0x600000);
-        size_t encodedDataSize = sawyercoding_write_chunk_buffer(encodedDataBuffer, (uint8 *)data, chunkHeader);
-
-        void * extraBytes = nullptr;
-        size_t extraBytesCount = 0;
-
-        int realChecksum = object_calculate_checksum(entry, data, dataSize);
-        if (realChecksum != entry->checksum)
+        if (fixChecksum)
         {
-            char objectName[9];
-            object_entry_get_name_fixed(objectName, sizeof(objectName), entry);
-            log_verbose("[%s] Incorrect checksum, adding salt bytes...", objectName);
-
-            // Calculate the value of extra bytes that can be appended to the data so that the
-            // data is then valid for the object's checksum
-            extraBytes = CalculateExtraBytesToFixChecksum(realChecksum, entry->checksum, &extraBytesCount);
-#ifdef DEBUG
+            int realChecksum = object_calculate_checksum(entry, data, dataSize);
+            if (realChecksum != entry->checksum)
             {
-                // Check that CalculateExtraBytesToFixChecksum did the job
+                char objectName[9];
+                object_entry_get_name_fixed(objectName, sizeof(objectName), entry);
+                log_verbose("[%s] Incorrect checksum, adding salt bytes...", objectName);
+
+                // Calculate the value of extra bytes that can be appended to the data so that the
+                // data is then valid for the object's checksum
+                size_t extraBytesCount = 0;
+                void * extraBytes = CalculateExtraBytesToFixChecksum(realChecksum, entry->checksum, &extraBytesCount);
+
+                // Create new data blob with appended bytes
                 size_t newDataSize = dataSize + extraBytesCount;
                 void * newData = Memory::Allocate<void>(newDataSize);
                 void * newDataSaltOffset = (void *)((uintptr_t)newData + dataSize);
                 Memory::Copy(newData, data, dataSize);
                 Memory::Copy(newDataSaltOffset, extraBytes, extraBytesCount);
 
-                int actualChecksum = object_calculate_checksum(entry, newData, newDataSize);
-                Guard::Assert(actualChecksum == realChecksum, "CalculateExtraBytesToFixChecksum failed to fix checksum.");
+                try
+                {
+                    int newRealChecksum = object_calculate_checksum(entry, newData, newDataSize);
+                    if (newRealChecksum != entry->checksum)
+                    {
+                        Guard::Fail("CalculateExtraBytesToFixChecksum failed to fix checksum.");
 
-                Memory::Free(newData);
+                        // Save old data form
+                        SaveObject(path, entry, data, dataSize, false);
+                    }
+                    else
+                    {
+                        // Save new data form
+                        SaveObject(path, entry, newData, newDataSize, false);
+                    }
+                    Memory::Free(newData);
+                    Memory::Free(extraBytes);
+                }
+                catch (Exception ex)
+                {
+                    Memory::Free(newData);
+                    Memory::Free(extraBytes);
+                    throw;
+                }
+                return;
             }
-#endif
         }
 
+        // Encode data
+        uint8 objectType = entry->flags & 0x0F;
+        sawyercoding_chunk_header chunkHeader;
+        chunkHeader.encoding = object_entry_group_encoding[objectType];
+        chunkHeader.length = dataSize;
+        uint8 * encodedDataBuffer = Memory::Allocate<uint8>(0x600000);
+        size_t encodedDataSize = sawyercoding_write_chunk_buffer(encodedDataBuffer, (uint8 *)data, chunkHeader);
+
         // Save to file
         try
         {
@@ -535,17 +556,10 @@ private:
             fs.Write(entry, sizeof(rct_object_entry));
             fs.Write(encodedDataBuffer, encodedDataSize);
 
-            if (extraBytes != nullptr)
-            {
-                fs.Write(extraBytes, extraBytesCount);
-            }
-
-            Memory::Free(extraBytes);
             Memory::Free(encodedDataBuffer);
         }
         catch (Exception ex)
         {
-            Memory::Free(extraBytes);
             Memory::Free(encodedDataBuffer);
             throw;
         }