Add proper error handling for memcpy

src/common/common_funcs.h

@@ -4,12 +4,15 @@
 
 #pragma once
 
+#include <cstring>
 #include <string>
+#include <type_traits>
 
 #if !defined(ARCHITECTURE_x86_64)
 #include <cstdlib> // for exit
 #endif
 #include "common/common_types.h"
+#include "common/logging/log.h"
 
 #define ARRAY_SIZE(a) (sizeof(a) / sizeof(a[0]))
 
@@ -17,6 +20,12 @@
 #define CONCAT2(x, y) DO_CONCAT2(x, y)
 #define DO_CONCAT2(x, y) x##y
 
+#define TRY(x, fail)                                                                               \
+    if (!(x)) {                                                                                    \
+        fail;                                                                                      \
+        return false;                                                                              \
+    }
+
 // helper macro to properly align structure members.
 // Calling INSERT_PADDING_BYTES will add a new member variable with a name like "pad121",
 // depending on the current source line to make sure variable names are unique.
@@ -60,3 +69,13 @@ __declspec(dllimport) void __stdcall DebugBreak(void);
 // This function might change the error code.
 // Defined in Misc.cpp.
 std::string GetLastErrorMsg();
+
+template <typename T>
+bool CheckedMemcpy(void* dest, T& container, std::ptrdiff_t offset, std::size_t size) {
+    static_assert(std::is_same_v<typename T::value_type, u8>, "Only works with u8");
+    if (container.size() < offset + size) {
+        return false;
+    }
+    std::memcpy(dest, container.data() + offset, size);
+    return true;
+}

src/core/data_container.cpp

@@ -4,6 +4,7 @@
 
 #include <cmath>
 #include "common/assert.h"
+#include "common/common_funcs.h"
 #include "core/data_container.h"
 
 namespace Core {
@@ -20,32 +21,53 @@ DPFSContainer::DPFSContainer(DPFSDescriptor descriptor_, u8 level1_selector_,
     ASSERT_MSG(descriptor.version == 0x10000, "DPFS Version is not correct");
 }
 
-u8 DPFSContainer::GetBit(u8 level, u8 selector, u64 index) const {
+bool DPFSContainer::GetBit(u8& out, u8 level, u8 selector, u64 index) const {
     ASSERT_MSG(level <= 2 && selector <= 1, "Level or selector invalid");
-    return (data[(descriptor.levels[level].offset + selector * descriptor.levels[level].size) / 4 +
-                 index / 32] >>
-            (31 - (index % 32))) &
-           static_cast<u32_le>(1);
+
+    const auto word =
+        (descriptor.levels[level].offset + selector * descriptor.levels[level].size) / 4 +
+        index / 32;
+    if (data.size() <= word) {
+        LOG_ERROR(Core, "Out of bound: level {} selector {} index {}", level, selector, index);
+        return false;
+    }
+
+    out = (data[word] >> (31 - (index % 32))) & static_cast<u32_le>(1);
+    return true;
 }
 
-u8 DPFSContainer::GetByte(u8 level, u8 selector, u64 index) const {
+bool DPFSContainer::GetByte(u8& out, u8 level, u8 selector, u64 index) const {
     ASSERT_MSG(level <= 2 && selector <= 1, "Level or selector invalid");
-    return reinterpret_cast<const u8*>(
-        data.data())[descriptor.levels[level].offset + selector * descriptor.levels[level].size +
-                     index];
+
+    const auto byte =
+        descriptor.levels[level].offset + selector * descriptor.levels[level].size + index;
+    if (data.size() * 4 <= byte) {
+        LOG_ERROR(Core, "Out of bound: level {} selector {} index {}", level, selector, index);
+        return false;
+    }
+
+    out =
+        reinterpret_cast<const u8*>(data.data())[descriptor.levels[level].offset +
+                                                 selector * descriptor.levels[level].size + index];
+    return true;
 }
 
-std::vector<u8> DPFSContainer::GetLevel3Data() const {
+bool DPFSContainer::GetLevel3Data(std::vector<u8>& out) const {
     std::vector<u8> level3_data(descriptor.levels[2].size);
     for (std::size_t i = 0; i < level3_data.size(); i++) {
         auto level2_bit_index = i / std::pow(2, descriptor.levels[2].block_size);
         auto level1_bit_index =
             (level2_bit_index / 8) / std::pow(2, descriptor.levels[1].block_size);
-        auto level2_selector = GetBit(0, level1_selector, level1_bit_index);
-        auto level3_selector = GetBit(1, level2_selector, level2_bit_index);
-        level3_data[i] = GetByte(2, level3_selector, i);
+
+        u8 level2_selector, level3_selector;
+        if (!GetBit(level2_selector, 0, level1_selector, level1_bit_index) ||
+            !GetBit(level3_selector, 1, level2_selector, level2_bit_index) ||
+            !GetByte(level3_data[i], 2, level3_selector, i)) {
+            return false;
+        }
     }
-    return level3_data;
+    out = std::move(level3_data);
+    return true;
 }
 
 DataContainer::DataContainer(std::vector<u8> data_) : data(std::move(data_)) {
@@ -75,7 +97,8 @@ bool DataContainer::IsGood() const {
 
 bool DataContainer::InitAsDISA() {
     DISAHeader header;
-    std::memcpy(&header, data.data() + 0x100, sizeof(header));
+    TRY(CheckedMemcpy(&header, data, 0x100, sizeof(header)),
+        LOG_ERROR(Core, "File size is too small"));
 
     if (header.version != 0x40000) {
         LOG_ERROR(Core, "DISA Version {:x} is not correct", header.version);
@@ -103,7 +126,8 @@ bool DataContainer::InitAsDISA() {
 
 bool DataContainer::InitAsDIFF() {
     DIFFHeader header;
-    std::memcpy(&header, data.data() + 0x100, sizeof(header));
+    TRY(CheckedMemcpy(&header, data, 0x100, sizeof(header)),
+        LOG_ERROR(Core, "File size is too small"));
 
     if (header.version != 0x30000) {
         LOG_ERROR(Core, "DIFF Version {:x} is not correct", header.version);
@@ -123,52 +147,75 @@ bool DataContainer::InitAsDIFF() {
     return true;
 }
 
-std::vector<u8> DataContainer::GetPartitionData(u8 index) const {
+bool DataContainer::GetPartitionData(std::vector<u8>& out, u8 index) const {
     auto partition_descriptor_offset = partition_table_offset + partition_descriptors[index].offset;
 
     DIFIHeader difi;
-    std::memcpy(&difi, data.data() + partition_descriptor_offset, sizeof(difi));
-    ASSERT_MSG(difi.magic == MakeMagic('D', 'I', 'F', 'I'), "DIFI Magic is not correct");
-    ASSERT_MSG(difi.version == 0x10000, "DIFI Version is not correct");
+    TRY(CheckedMemcpy(&difi, data, partition_descriptor_offset, sizeof(difi)),
+        LOG_ERROR(Core, "File size is too small"));
+
+    if (difi.magic != MakeMagic('D', 'I', 'F', 'I') || difi.version != 0x10000) {
+        LOG_ERROR(Core, "Invalid magic {:08x} or version {}", difi.magic, difi.version);
+        return false;
+    }
 
     ASSERT_MSG(difi.ivfc.size >= sizeof(IVFCDescriptor), "IVFC descriptor size is too small");
     IVFCDescriptor ivfc_descriptor;
-    std::memcpy(&ivfc_descriptor, data.data() + partition_descriptor_offset + difi.ivfc.offset,
-                sizeof(ivfc_descriptor));
+    TRY(CheckedMemcpy(&ivfc_descriptor, data, partition_descriptor_offset + difi.ivfc.offset,
+                      sizeof(ivfc_descriptor)),
+        LOG_ERROR(Core, "File size is too small"));
 
     if (difi.enable_external_IVFC_level_4) {
-        std::vector<u8> result(
+        if (data.size() < partitions[index].offset + difi.external_IVFC_level_4_offset +
+                              ivfc_descriptor.levels[3].size) {
+            LOG_ERROR(Core, "File size is too small");
+            return false;
+        }
+        out = std::vector<u8>(
             data.data() + partitions[index].offset + difi.external_IVFC_level_4_offset,
             data.data() + partitions[index].offset + difi.external_IVFC_level_4_offset +
                 ivfc_descriptor.levels[3].size);
-        return result;
+        return true;
     }
 
     // Unwrap DPFS Tree
     ASSERT_MSG(difi.dpfs.size >= sizeof(DPFSDescriptor), "DPFS descriptor size is too small");
     DPFSDescriptor dpfs_descriptor;
-    std::memcpy(&dpfs_descriptor, data.data() + partition_descriptor_offset + difi.dpfs.offset,
-                sizeof(dpfs_descriptor));
+    TRY(CheckedMemcpy(&dpfs_descriptor, data, partition_descriptor_offset + difi.dpfs.offset,
+                      sizeof(dpfs_descriptor)),
+        LOG_ERROR(Core, "File size is too small"));
 
     std::vector<u32_le> partition_data(partitions[index].size / 4);
-    std::memcpy(partition_data.data(), data.data() + partitions[index].offset,
-                partitions[index].size);
+    TRY(CheckedMemcpy(partition_data.data(), data, partitions[index].offset,
+                      partitions[index].size),
+        LOG_ERROR(Core, "File size is too small"));
 
     DPFSContainer dpfs_container(dpfs_descriptor, difi.dpfs_level1_selector,
                                  std::move(partition_data));
-    auto ivfc_data = dpfs_container.GetLevel3Data();
 
-    std::vector<u8> result(ivfc_data.data() + ivfc_descriptor.levels[3].offset,
-                           ivfc_data.data() + ivfc_descriptor.levels[3].offset +
-                               ivfc_descriptor.levels[3].size);
-    return result;
+    std::vector<u8> ivfc_data;
+    if (!dpfs_container.GetLevel3Data(ivfc_data)) {
+        return false;
+    }
+
+    if (ivfc_data.size() < ivfc_descriptor.levels[3].offset + ivfc_descriptor.levels[3].size) {
+        LOG_ERROR(Core, "IVFC data size is too small");
+        return false;
+    }
+
+    out = std::vector<u8>(ivfc_data.data() + ivfc_descriptor.levels[3].offset,
+                          ivfc_data.data() + ivfc_descriptor.levels[3].offset +
+                              ivfc_descriptor.levels[3].size);
+    return true;
 }
 
-std::vector<std::vector<u8>> DataContainer::GetIVFCLevel4Data() const {
+bool DataContainer::GetIVFCLevel4Data(std::vector<std::vector<u8>>& out) const {
     if (partition_count == 1) {
-        return {GetPartitionData(0)};
+        out.resize(1);
+        return GetPartitionData(out[0], 0);
     } else {
-        return {GetPartitionData(0), GetPartitionData(1)};
+        out.resize(2);
+        return GetPartitionData(out[0], 0) && GetPartitionData(out[1], 1);
     }
 }
 

src/core/data_container.h

@@ -94,11 +94,11 @@ public:
     explicit DPFSContainer(DPFSDescriptor descriptor, u8 level1_selector, std::vector<u32_le> data);
 
     /// Unwraps the DPFS Tree, returning actual data in Level3.
-    std::vector<u8> GetLevel3Data() const;
+    bool GetLevel3Data(std::vector<u8>& out) const;
 
 private:
-    u8 GetBit(u8 level, u8 selector, u64 index) const;
-    u8 GetByte(u8 level, u8 selector, u64 index) const;
+    bool GetBit(u8& out, u8 level, u8 selector, u64 index) const;
+    bool GetByte(u8& out, u8 level, u8 selector, u64 index) const;
 
     DPFSDescriptor descriptor;
     u8 level1_selector;
@@ -114,7 +114,7 @@ public:
     ~DataContainer();
 
     /// Unwraps the whole container, returning the data in IVFC Level 4 of all partitions.
-    std::vector<std::vector<u8>> GetIVFCLevel4Data() const;
+    bool GetIVFCLevel4Data(std::vector<std::vector<u8>>& out) const;
 
     bool IsGood() const;
 
@@ -123,7 +123,7 @@ private:
     bool InitAsDIFF();
 
     /// Unwraps the whole container, returning the data in IVFC Level 4 of a partition.
-    std::vector<u8> GetPartitionData(u8 index) const;
+    bool GetPartitionData(std::vector<u8>& out, u8 index) const;
 
     bool is_good = false;
     std::vector<u8> data;

src/core/importer.cpp

@@ -164,7 +164,12 @@ bool SDMCImporter::ImportSavegame(u64 id, [[maybe_unused]] const ProgressCallbac
         return false;
     }
 
-    SDSavegame save(std::move(container.GetIVFCLevel4Data()));
+    std::vector<std::vector<u8>> container_data;
+    if (!container.GetIVFCLevel4Data(container_data)) {
+        return false;
+    }
+
+    SDSavegame save(std::move(container_data));
     if (!save.IsGood()) {
         return false;
     }
@@ -329,7 +334,12 @@ bool SDMCImporter::ImportSysdata(u64 id, [[maybe_unused]] const ProgressCallback
             return false;
         }
 
-        SDSavegame save(std::move(container.GetIVFCLevel4Data()));
+        std::vector<std::vector<u8>> container_data;
+        if (!container.GetIVFCLevel4Data(container_data)) {
+            return false;
+        }
+
+        SDSavegame save(std::move(container_data));
         if (!save.IsGood()) {
             return false;
         }

src/core/inner_fat.cpp

@@ -4,6 +4,7 @@
 
 #include <fmt/format.h>
 #include "common/assert.h"
+#include "common/common_funcs.h"
 #include "common/file_util.h"
 #include "core/data_container.h"
 #include "core/decryptor.h"
@@ -95,23 +96,27 @@ SDSavegame::SDSavegame(std::vector<std::vector<u8>> partitions) {
 SDSavegame::~SDSavegame() = default;
 
 bool SDSavegame::Init() {
-    auto header_iter = duplicate_data ? data.data() : partitionA.data();
+    const auto& header_vector = duplicate_data ? data : partitionA;
 
     // Read header
-    std::memcpy(&header, header_iter, sizeof(header));
+    TRY(CheckedMemcpy(&header, header_vector, 0, sizeof(header)),
+        LOG_ERROR(Core, "File size is too small"));
+
     if (header.magic != MakeMagic('S', 'A', 'V', 'E') || header.version != 0x40000) {
         LOG_ERROR(Core, "File is invalid, decryption errors may have happened.");
         return false;
     }
 
     // Read filesystem information
-    std::memcpy(&fs_info, header_iter + header.filesystem_information_offset, sizeof(fs_info));
+    TRY(CheckedMemcpy(&fs_info, header_vector, header.filesystem_information_offset,
+                      sizeof(fs_info)),
+        LOG_ERROR(Core, "File size is too small"));
 
     // Read data region
     if (duplicate_data) {
         data_region.resize(fs_info.data_region_block_count * fs_info.data_region_block_size);
-        std::memcpy(data_region.data(), data.data() + fs_info.data_region_offset,
-                    data_region.size());
+        TRY(CheckedMemcpy(data_region.data(), data, fs_info.data_region_offset, data_region.size()),
+            LOG_ERROR(Core, "File size is too small"));
     } else {
         data_region = std::move(partitionB);
     }
@@ -121,31 +126,36 @@ bool SDSavegame::Init() {
     // so it should be safe to directly read the bytes.
 
     // Read directory entry table
-    auto directory_entry_table_iter =
-        header_iter + (duplicate_data ? fs_info.data_region_offset +
-                                            fs_info.directory_entry_table.duplicate.block_index *
-                                                fs_info.data_region_block_size
-                                      : fs_info.directory_entry_table.non_duplicate);
-
     directory_entry_table.resize(fs_info.maximum_directory_count + 2); // including head and root
-    std::memcpy(directory_entry_table.data(), directory_entry_table_iter,
-                directory_entry_table.size() * sizeof(DirectoryEntryTableEntry));
+
+    auto directory_entry_table_pos =
+        duplicate_data
+            ? fs_info.data_region_offset + fs_info.directory_entry_table.duplicate.block_index *
+                                               fs_info.data_region_block_size
+            : fs_info.directory_entry_table.non_duplicate;
+
+    TRY(CheckedMemcpy(directory_entry_table.data(), header_vector, directory_entry_table_pos,
+                      directory_entry_table.size() * sizeof(DirectoryEntryTableEntry)),
+        LOG_ERROR(Core, "File is too small"));
 
     // Read file entry table
-    auto file_entry_table_iter =
-        header_iter + (duplicate_data ? fs_info.data_region_offset +
-                                            fs_info.file_entry_table.duplicate.block_index *
-                                                fs_info.data_region_block_size
-                                      : fs_info.file_entry_table.non_duplicate);
-
     file_entry_table.resize(fs_info.maximum_file_count + 1); // including head
-    std::memcpy(file_entry_table.data(), file_entry_table_iter,
-                file_entry_table.size() * sizeof(FileEntryTableEntry));
+
+    auto file_entry_table_pos =
+        duplicate_data
+            ? fs_info.data_region_offset +
+                  fs_info.file_entry_table.duplicate.block_index * fs_info.data_region_block_size
+            : fs_info.file_entry_table.non_duplicate;
+
+    TRY(CheckedMemcpy(file_entry_table.data(), header_vector, file_entry_table_pos,
+                      file_entry_table.size() * sizeof(FileEntryTableEntry)),
+        LOG_ERROR(Core, "File is too small"));
 
     // Read file allocation table
     fat.resize(fs_info.file_allocation_table_entry_count);
-    std::memcpy(fat.data(), header_iter + fs_info.file_allocation_table_offset,
-                fat.size() * sizeof(FATNode));
+    TRY(CheckedMemcpy(fat.data(), header_vector, fs_info.file_allocation_table_offset,
+                      fat.size() * sizeof(FATNode)),
+        LOG_ERROR(Core, "File size is too small"));
 
     return true;
 }
@@ -189,6 +199,11 @@ bool SDSavegame::ExtractFile(const std::string& path, std::size_t index) const {
 
         const std::size_t size = fs_info.data_region_block_size * (last_block - block + 1);
         const std::size_t to_write = std::min(file_size, size);
+
+        if (data_region.size() < fs_info.data_region_block_size * block + to_write) {
+            LOG_ERROR(Core, "Out of bound block: {} to_write: {}", block, to_write);
+            return false;
+        }
         if (file.WriteBytes(data_region.data() + fs_info.data_region_block_size * block,
                             to_write) != to_write) {
             LOG_ERROR(Core, "Write data failed (file: {})", path + name);
@@ -257,36 +272,50 @@ bool SDExtdata::Init() {
     if (!vsxe_container.IsGood()) {
         return false;
     }
-    auto vsxe = vsxe_container.GetIVFCLevel4Data()[0];
+
+    std::vector<std::vector<u8>> container_data;
+    if (!vsxe_container.GetIVFCLevel4Data(container_data)) {
+        return false;
+    }
+    const auto& vsxe = container_data[0];
 
     // Read header
-    std::memcpy(&header, vsxe.data(), sizeof(header));
+    TRY(CheckedMemcpy(&header, vsxe, 0, sizeof(header)), LOG_ERROR(Core, "File size is too small"));
+
     if (header.magic != MakeMagic('V', 'S', 'X', 'E') || header.version != 0x30000) {
         LOG_ERROR(Core, "File is invalid, decryption errors may have happened.");
         return false;
     }
 
     // Read filesystem information
-    std::memcpy(&fs_info, vsxe.data() + header.filesystem_information_offset, sizeof(fs_info));
+    TRY(CheckedMemcpy(&fs_info, vsxe, header.filesystem_information_offset, sizeof(fs_info)),
+        LOG_ERROR(Core, "File size is too small"));
 
     // Read data region
-    data_region.resize(fs_info.data_region_block_count * fs_info.data_region_block_size);
-    std::memcpy(data_region.data(), vsxe.data() + fs_info.data_region_offset, data_region.size());
+    TRY(CheckedMemcpy(data_region.data(), vsxe, fs_info.data_region_offset, data_region.size()),
+        LOG_ERROR(Core, "File size is too small"));
 
     // Read directory entry table
     directory_entry_table.resize(fs_info.maximum_directory_count + 2); // including head and root
-    std::memcpy(directory_entry_table.data(),
-                vsxe.data() + fs_info.data_region_offset +
-                    fs_info.directory_entry_table.duplicate.block_index *
-                        fs_info.data_region_block_size,
-                directory_entry_table.size() * sizeof(DirectoryEntryTableEntry));
+
+    const auto directory_entry_table_pos =
+        fs_info.data_region_offset +
+        fs_info.directory_entry_table.duplicate.block_index * fs_info.data_region_block_size;
+
+    TRY(CheckedMemcpy(directory_entry_table.data(), vsxe, directory_entry_table_pos,
+                      directory_entry_table.size() * sizeof(DirectoryEntryTableEntry)),
+        LOG_ERROR(Core, "File size is too small"));
 
     // Read file entry table
     file_entry_table.resize(fs_info.maximum_file_count + 1); // including head
-    std::memcpy(file_entry_table.data(),
-                vsxe.data() + fs_info.data_region_offset +
-                    fs_info.file_entry_table.duplicate.block_index * fs_info.data_region_block_size,
-                file_entry_table.size() * sizeof(FileEntryTableEntry));
+
+    const auto file_entry_table_pos =
+        fs_info.data_region_offset +
+        fs_info.file_entry_table.duplicate.block_index * fs_info.data_region_block_size;
+
+    TRY(CheckedMemcpy(file_entry_table.data(), vsxe, file_entry_table_pos,
+                      file_entry_table.size() * sizeof(FileEntryTableEntry)),
+        LOG_ERROR(Core, "File size is too small"));
 
     // File allocation table isn't needed here, as the only files allocated by them are
     // directory/file entry tables which we already read above.
@@ -346,8 +375,12 @@ bool SDExtdata::ExtractFile(const std::string& path, std::size_t index) const {
         return false;
     }
 
-    auto data = container.GetIVFCLevel4Data()[0];
-    if (file.WriteBytes(data.data(), data.size()) != data.size()) {
+    std::vector<std::vector<u8>> data;
+    if (!container.GetIVFCLevel4Data(data)) {
+        return false;
+    }
+
+    if (file.WriteBytes(data[0].data(), data[0].size()) != data[0].size()) {
         LOG_ERROR(Core, "Write data failed (file: {})", path + name);
         return false;
     }

src/frontend/utilities.cpp

@@ -38,12 +38,6 @@ UtilitiesDialog::UtilitiesDialog(QWidget* parent)
         ui->sdDecryptionDisabledLabel->setVisible(!checked);
         ui->sdDecryption->setEnabled(checked);
 
-        ui->savedataExtractionLabel->setVisible(false);
-        ui->savedataExtractionDisabledLabel->setVisible(false);
-        ui->savedataExtractionLabel->setVisible(checked);
-        ui->savedataExtractionDisabledLabel->setVisible(!checked);
-        ui->savedataExtraction->setEnabled(checked);
-
         ui->extdataExtractionLabel->setVisible(false);
         ui->extdataExtractionDisabledLabel->setVisible(false);
         ui->extdataExtractionLabel->setVisible(checked);
@@ -206,7 +200,12 @@ void UtilitiesDialog::SaveDataExtractionTool() {
                 return false;
             }
 
-            Core::SDSavegame save(std::move(container.GetIVFCLevel4Data()));
+            std::vector<std::vector<u8>> container_data;
+            if (!container.GetIVFCLevel4Data(container_data)) {
+                return false;
+            }
+
+            Core::SDSavegame save(std::move(container_data));
             if (!save.IsGood()) {
                 return false;
             }
@@ -228,7 +227,12 @@ void UtilitiesDialog::SaveDataExtractionTool() {
                 return false;
             }
 
-            Core::SDSavegame save(std::move(container.GetIVFCLevel4Data()));
+            std::vector<std::vector<u8>> container_data;
+            if (!container.GetIVFCLevel4Data(container_data)) {
+                return false;
+            }
+
+            Core::SDSavegame save(std::move(container_data));
             if (!save.IsGood()) {
                 return false;
             }