Make logic to free buffers easier

src/module/ModuleDataFactory.cpp

@@ -18,6 +18,7 @@
 #include "ModuleDataFactory.h"
 #include "../ElfUtils.h"
 #include "../utils/FileUtils.h"
+#include "utils/OnLeavingScope.h"
 #include "utils/wiiu_zlib.hpp"
 #include <coreinit/cache.h>
 #include <coreinit/debug.h>
@@ -38,21 +39,20 @@ ModuleDataFactory::load(const std::string &path, uint32_t destination_address, u
         DEBUG_FUNCTION_LINE("Failed to load file");
         return {};
     }
+    auto cleanupBuffer = onLeavingScope([buffer]() { free(buffer); });
 
     // Load ELF data
     if (!reader.load(reinterpret_cast<char *>(buffer), fsize)) {
         DEBUG_FUNCTION_LINE("Can't find or process %s", path.c_str());
-        free(buffer);
         return {};
     }
 
     uint32_t sec_num = reader.sections.size();
 
-    auto **destinations = (uint8_t **) malloc(sizeof(uint8_t *) * sec_num);
+    auto destinations = make_unique_nothrow<uint8_t *[]>(sec_num);
-
     if (!destinations) {
-        DEBUG_FUNCTION_LINE("Failed to alloc memory for destinations");
+        DEBUG_FUNCTION_LINE_ERR("Failed alloc memory for destinations array");
-        free(buffer);
+        return {};
     }
 
     uint32_t sizeOfModule = 0;
@@ -69,8 +69,6 @@ ModuleDataFactory::load(const std::string &path, uint32_t destination_address, u
 
     if (sizeOfModule > maximum_size) {
         DEBUG_FUNCTION_LINE("Module is too big.");
-        free(buffer);
-        free(destinations);
         return {};
     }
 
@@ -107,18 +105,23 @@ ModuleDataFactory::load(const std::string &path, uint32_t destination_address, u
                 destinations[psec->get_index()] -= 0x10000000;
             } else if (address >= 0xC0000000) {
                 DEBUG_FUNCTION_LINE("Loading section from 0xC0000000 is NOT supported");
-                free(destinations);
-                free(buffer);
                 return {};
             } else {
                 DEBUG_FUNCTION_LINE("Unhandled case");
-                free(destinations);
-                free(buffer);
                 return {};
             }
 
             const char *p = reader.sections[i]->get_data();
 
+            if (destination < targetAddress) {
+                DEBUG_FUNCTION_LINE_ERR("Tried to underflow buffer. %08X < %08X", destination, targetAddress);
+                OSFatal("CustomRPXLoader: Tried to underflow buffer");
+            }
+            if (destination + sectionSize > destination_address) {
+                DEBUG_FUNCTION_LINE_ERR("Tried to overflow buffer. %08X > %08X", destination + sectionSize, destination_address);
+                OSFatal("CustomRPXLoader: Tried to overflow buffer");
+            }
+
             if (psec->get_type() == SHT_NOBITS) {
                 DEBUG_FUNCTION_LINE("memset section %s %08X [%08X] to 0 (%d bytes)", psec->get_name().c_str(), destination, destination + sectionSize, sectionSize);
                 memset((void *) destination, 0, sectionSize);
@@ -150,13 +153,11 @@ ModuleDataFactory::load(const std::string &path, uint32_t destination_address, u
             DEBUG_FUNCTION_LINE("Linking (%d)... %s", i, psec->get_name().c_str());
             if (!linkSection(reader, psec->get_index(), (uint32_t) destinations[psec->get_index()], offset_text, offset_data, trampolin_data, trampolin_data_length)) {
                 DEBUG_FUNCTION_LINE("elfLink failed");
-                free(destinations);
-                free(buffer);
                 return {};
             }
         }
     }
-    std::vector<RelocationData> relocationData = getImportRelocationData(reader, destinations);
+    std::vector<RelocationData> relocationData = getImportRelocationData(reader, destinations.get());
 
     for (auto const &reloc : relocationData) {
         moduleData.addRelocationData(reloc);
@@ -167,9 +168,6 @@ ModuleDataFactory::load(const std::string &path, uint32_t destination_address, u
     DEBUG_FUNCTION_LINE("ICInvalidateRange %08X - %08X", targetAddress, destination_address);
     ICInvalidateRange((void *) targetAddress, destination_address - targetAddress);
 
-    free(destinations);
-    free(buffer);
-
     moduleData.setEntrypoint(entrypoint);
     DEBUG_FUNCTION_LINE("Saved entrypoint as %08X", entrypoint);
 

src/utils/OnLeavingScope.h

@@ -0,0 +1,113 @@
+/**
+* The contents of this file are based on the article posted at the
+* following location:
+*
+*   http://crascit.com/2015/06/03/on-leaving-scope-part-2/
+*
+* The material in that article has some commonality with the code made
+* available as part of Facebook's folly library at:
+*
+*   https://github.com/facebook/folly/blob/master/folly/ScopeGuard.h
+*
+* Furthermore, similar material is currently part of a draft proposal
+* to the C++ standards committee, referencing the same work by Andrei
+* Alexandresu that led to the folly implementation. The draft proposal
+* can be found at:
+*
+*   http://www.open-std.org/jtc1/sc22/wg21/docs/papers/2014/n4189.pdf
+*
+* With the above in mind, the content below is made available under
+* the same license terms as folly to minimize any legal concerns.
+* Should there be ambiguity over copyright ownership between Facebook
+* and myself for any material included in this file, it should be
+* interpreted that Facebook is the copyright owner for the ambiguous
+* section of code concerned.
+*
+*   Craig Scott
+*   3rd June 2015
+*
+* ----------------------------------------------------------------------
+*
+* Copyright 2015 Craig Scott
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+*    http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+#ifndef CRASCIT_ONLEAVINGSCOPE_H
+#define CRASCIT_ONLEAVINGSCOPE_H
+
+#include <type_traits>
+#include <utility>
+
+
+/**
+* This class is intended to be used only to create a local object on the
+* stack. It accepts a function object in its constructor and will invoke
+* that function object in its destructor. The class provides a move
+* constructor so it can be used with the onLeavingScope factory function,
+* but it cannot be copied.
+*
+* Do no use this function directly, use the onLeavingScope() factory
+* function instead.
+*/
+template<typename Func>
+class OnLeavingScope {
+public:
+    // Prevent copying
+    OnLeavingScope(const OnLeavingScope &) = delete;
+    OnLeavingScope &operator=(const OnLeavingScope &) = delete;
+
+    // Allow moving
+    OnLeavingScope(OnLeavingScope &&other) : m_func(std::move(other.m_func)),
+                                             m_owner(other.m_owner) {
+        other.m_owner = false;
+    }
+
+    OnLeavingScope(const Func &f) : m_func(f),
+                                    m_owner(true) {
+    }
+    OnLeavingScope(Func &&f) : m_func(std::move(f)),
+                               m_owner(true) {
+    }
+    ~OnLeavingScope() {
+        if (m_owner)
+            m_func();
+    }
+
+private:
+    Func m_func;
+    bool m_owner;
+};
+
+
+/**
+* Factory function for creating an OnLeavingScope object. It is intended
+* to be used like so:
+*
+*     auto cleanup = onLeavingScope(...);
+*
+* where the ... could be a lambda function, function object or pointer to
+* a free function to be invoked when the cleanup object goes out of scope.
+* The function object must take no function arguments, but can return any
+* type (the return value is ignored).
+*
+* The \a Func template parameter would rarely, if ever, be manually
+* specified. Normally, it would be deduced automatically by the compiler
+* from the object passed as the function argument.
+*/
+template<typename Func>
+OnLeavingScope<typename std::decay<Func>::type> onLeavingScope(Func &&f) {
+    return OnLeavingScope<typename std::decay<Func>::type>(std::forward<Func>(f));
+}
+
+#endif // CRASCIT_ONLEAVINGSCOPE_H