3 files changed, 139 insertions, 43 deletions

diff --git a/src/core/loader/3dsx.cpp b/src/core/loader/3dsx.cpp
index 14aeebebb..055661363 100644
--- a/src/core/loader/3dsx.cpp
+++ b/src/core/loader/3dsx.cpp
@@ -19,7 +19,7 @@
 
 namespace Loader {
 
-/**
+/*
  * File layout:
  * - File header
  * - Code, rodata and data relocation table headers
@@ -39,13 +39,16 @@ namespace Loader {
  * The entrypoint is always the start of the code segment.
  * The BSS section must be cleared manually by the application.
  */
+
 enum THREEDSX_Error {
     ERROR_NONE = 0,
     ERROR_READ = 1,
     ERROR_FILE = 2,
     ERROR_ALLOC = 3
 };
+
 static const u32 RELOCBUFSIZE = 512;
+static const unsigned int NUM_SEGMENTS = 3;
 
 // File header
 #pragma pack(1)
@@ -98,7 +101,10 @@ static u32 TranslateAddr(u32 addr, const THREEloadinfo *loadinfo, u32* offsets)
     return loadinfo->seg_addrs[2] + addr - offsets[1];
 }
 
-static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
+using Kernel::SharedPtr;
+using Kernel::CodeSet;
+
+static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr, SharedPtr<CodeSet>* out_codeset)
 {
     if (!file.IsOpen())
         return ERROR_FILE;
@@ -116,15 +122,13 @@ static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
     loadinfo.seg_sizes[1] = (hdr.rodata_seg_size + 0xFFF) &~0xFFF;
     loadinfo.seg_sizes[2] = (hdr.data_seg_size + 0xFFF) &~0xFFF;
     u32 offsets[2] = { loadinfo.seg_sizes[0], loadinfo.seg_sizes[0] + loadinfo.seg_sizes[1] };
-    u32 data_load_size = (hdr.data_seg_size - hdr.bss_size + 0xFFF) &~0xFFF;
-    u32 bss_load_size = loadinfo.seg_sizes[2] - data_load_size;
-    u32 n_reloc_tables = hdr.reloc_hdr_size / 4;
-    std::vector<u8> all_mem(loadinfo.seg_sizes[0] + loadinfo.seg_sizes[1] + loadinfo.seg_sizes[2] + 3 * n_reloc_tables);
+    u32 n_reloc_tables = hdr.reloc_hdr_size / sizeof(u32);
+    std::vector<u8> program_image(loadinfo.seg_sizes[0] + loadinfo.seg_sizes[1] + loadinfo.seg_sizes[2]);
 
     loadinfo.seg_addrs[0] = base_addr;
     loadinfo.seg_addrs[1] = loadinfo.seg_addrs[0] + loadinfo.seg_sizes[0];
     loadinfo.seg_addrs[2] = loadinfo.seg_addrs[1] + loadinfo.seg_sizes[1];
-    loadinfo.seg_ptrs[0] = &all_mem[0];
+    loadinfo.seg_ptrs[0] = program_image.data();
     loadinfo.seg_ptrs[1] = loadinfo.seg_ptrs[0] + loadinfo.seg_sizes[0];
     loadinfo.seg_ptrs[2] = loadinfo.seg_ptrs[1] + loadinfo.seg_sizes[1];
 
@@ -132,10 +136,9 @@ static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
     file.Seek(hdr.header_size, SEEK_SET);
 
     // Read the relocation headers
-    u32* relocs = (u32*)(loadinfo.seg_ptrs[2] + hdr.data_seg_size);
-
-    for (unsigned current_segment : {0, 1, 2}) {
-        size_t size = n_reloc_tables * 4;
+    std::vector<u32> relocs(n_reloc_tables * NUM_SEGMENTS);
+    for (unsigned int current_segment = 0; current_segment < NUM_SEGMENTS; ++current_segment) {
+        size_t size = n_reloc_tables * sizeof(u32);
         if (file.ReadBytes(&relocs[current_segment * n_reloc_tables], size) != size)
             return ERROR_READ;
     }
@@ -152,7 +155,7 @@ static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
     memset((char*)loadinfo.seg_ptrs[2] + hdr.data_seg_size - hdr.bss_size, 0, hdr.bss_size);
 
     // Relocate the segments
-    for (unsigned current_segment : {0, 1, 2}) {
+    for (unsigned int current_segment = 0; current_segment < NUM_SEGMENTS; ++current_segment) {
         for (unsigned current_segment_reloc_table = 0; current_segment_reloc_table < n_reloc_tables; current_segment_reloc_table++) {
             u32 n_relocs = relocs[current_segment * n_reloc_tables + current_segment_reloc_table];
             if (current_segment_reloc_table >= 2) {
@@ -160,7 +163,7 @@ static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
                 file.Seek(n_relocs*sizeof(THREEDSX_Reloc), SEEK_CUR);
                 continue;
             }
-            static THREEDSX_Reloc reloc_table[RELOCBUFSIZE];
+            THREEDSX_Reloc reloc_table[RELOCBUFSIZE];
 
             u32* pos = (u32*)loadinfo.seg_ptrs[current_segment];
             const u32* end_pos = pos + (loadinfo.seg_sizes[current_segment] / 4);
@@ -179,7 +182,7 @@ static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
                     pos += table.skip;
                     s32 num_patches = table.patch;
                     while (0 < num_patches && pos < end_pos) {
-                        u32 in_addr = (char*)pos - (char*)&all_mem[0];
+                        u32 in_addr = (u8*)pos - program_image.data();
                         u32 addr = TranslateAddr(*pos, &loadinfo, offsets);
                         LOG_TRACE(Loader, "Patching %08X <-- rel(%08X,%d) (%08X)\n",
                                   base_addr + in_addr, addr, current_segment_reloc_table, *pos);
@@ -201,14 +204,29 @@ static THREEDSX_Error Load3DSXFile(FileUtil::IOFile& file, u32 base_addr)
         }
     }
 
-    // Write the data
-    memcpy(Memory::GetPointer(base_addr), &all_mem[0], loadinfo.seg_sizes[0] + loadinfo.seg_sizes[1] + loadinfo.seg_sizes[2]);
+    // Create the CodeSet
+    SharedPtr<CodeSet> code_set = CodeSet::Create("", 0);
+
+    code_set->code.offset = loadinfo.seg_ptrs[0] - program_image.data();
+    code_set->code.addr   = loadinfo.seg_addrs[0];
+    code_set->code.size   = loadinfo.seg_sizes[0];
+
+    code_set->rodata.offset = loadinfo.seg_ptrs[1] - program_image.data();
+    code_set->rodata.addr   = loadinfo.seg_addrs[1];
+    code_set->rodata.size   = loadinfo.seg_sizes[1];
 
-    LOG_DEBUG(Loader, "CODE:   %u pages\n", loadinfo.seg_sizes[0] / 0x1000);
-    LOG_DEBUG(Loader, "RODATA: %u pages\n", loadinfo.seg_sizes[1] / 0x1000);
-    LOG_DEBUG(Loader, "DATA:   %u pages\n", data_load_size / 0x1000);
-    LOG_DEBUG(Loader, "BSS:    %u pages\n", bss_load_size / 0x1000);
+    code_set->data.offset = loadinfo.seg_ptrs[2] - program_image.data();
+    code_set->data.addr   = loadinfo.seg_addrs[2];
+    code_set->data.size   = loadinfo.seg_sizes[2];
 
+    code_set->entrypoint = code_set->code.addr;
+    code_set->memory = std::make_shared<std::vector<u8>>(std::move(program_image));
+
+    LOG_DEBUG(Loader, "code size:   0x%X", loadinfo.seg_sizes[0]);
+    LOG_DEBUG(Loader, "rodata size: 0x%X", loadinfo.seg_sizes[1]);
+    LOG_DEBUG(Loader, "data size:   0x%X (including 0x%X of bss)", loadinfo.seg_sizes[2], hdr.bss_size);
+
+    *out_codeset = code_set;
     return ERROR_NONE;
 }
 
@@ -231,16 +249,19 @@ ResultStatus AppLoader_THREEDSX::Load() {
     if (!file->IsOpen())
         return ResultStatus::Error;
 
-    Kernel::g_current_process = Kernel::Process::Create(filename, 0);
+    SharedPtr<CodeSet> codeset;
+    if (Load3DSXFile(*file, Memory::PROCESS_IMAGE_VADDR, &codeset) != ERROR_NONE)
+        return ResultStatus::Error;
+    codeset->name = filename;
+
+    Kernel::g_current_process = Kernel::Process::Create(std::move(codeset));
     Kernel::g_current_process->svc_access_mask.set();
     Kernel::g_current_process->address_mappings = default_address_mappings;
 
     // Attach the default resource limit (APPLICATION) to the process
     Kernel::g_current_process->resource_limit = Kernel::ResourceLimit::GetForCategory(Kernel::ResourceLimitCategory::APPLICATION);
 
-    Load3DSXFile(*file, Memory::PROCESS_IMAGE_VADDR);
-
-    Kernel::g_current_process->Run(Memory::PROCESS_IMAGE_VADDR, 48, Kernel::DEFAULT_STACK_SIZE);
+    Kernel::g_current_process->Run(48, Kernel::DEFAULT_STACK_SIZE);
 
     is_loaded = true;
     return ResultStatus::Success;
diff --git a/src/core/loader/elf.cpp b/src/core/loader/elf.cpp
index a7eea78aa..ca3c18a9f 100644
--- a/src/core/loader/elf.cpp
+++ b/src/core/loader/elf.cpp
@@ -16,6 +16,9 @@
 #include "core/loader/elf.h"
 #include "core/memory.h"
 
+using Kernel::SharedPtr;
+using Kernel::CodeSet;
+
 ////////////////////////////////////////////////////////////////////////////////////////////////////
 // ELF Header Constants
 
@@ -97,6 +100,12 @@ enum ElfSectionFlags
 #define PT_LOPROC  0x70000000
 #define PT_HIPROC  0x7FFFFFFF
 
+// Segment flags
+#define PF_X               0x1
+#define PF_W               0x2
+#define PF_R               0x4
+#define PF_MASKPROC 0xF0000000
+
 typedef unsigned int   Elf32_Addr;
 typedef unsigned short Elf32_Half;
 typedef unsigned int   Elf32_Off;
@@ -193,7 +202,7 @@ public:
     ElfMachine GetMachine() const { return (ElfMachine)(header->e_machine); }
     u32 GetEntryPoint() const { return entryPoint; }
     u32 GetFlags() const { return (u32)(header->e_flags); }
-    void LoadInto(u32 vaddr);
+    SharedPtr<CodeSet> LoadInto(u32 vaddr);
     bool LoadSymbols();
 
     int GetNumSegments() const { return (int)(header->e_phnum); }
@@ -249,7 +258,7 @@ const char *ElfReader::GetSectionName(int section) const {
     return nullptr;
 }
 
-void ElfReader::LoadInto(u32 vaddr) {
+SharedPtr<CodeSet> ElfReader::LoadInto(u32 vaddr) {
     LOG_DEBUG(Loader, "String section: %i", header->e_shstrndx);
 
     // Should we relocate?
@@ -267,19 +276,61 @@ void ElfReader::LoadInto(u32 vaddr) {
     u32 segment_addr[32];
     u32 base_addr = relocate ? vaddr : 0;
 
-    for (unsigned i = 0; i < header->e_phnum; i++) {
-        Elf32_Phdr* p = segments + i;
-        LOG_DEBUG(Loader, "Type: %i Vaddr: %08x Filesz: %i Memsz: %i ", p->p_type, p->p_vaddr,
+    u32 total_image_size = 0;
+    for (unsigned int i = 0; i < header->e_phnum; ++i) {
+        Elf32_Phdr* p = &segments[i];
+        if (p->p_type == PT_LOAD) {
+            total_image_size += (p->p_memsz + 0xFFF) & ~0xFFF;
+        }
+    }
+
+    std::vector<u8> program_image(total_image_size);
+    size_t current_image_position = 0;
+
+    SharedPtr<CodeSet> codeset = CodeSet::Create("", 0);
+
+    for (unsigned int i = 0; i < header->e_phnum; ++i) {
+        Elf32_Phdr* p = &segments[i];
+        LOG_DEBUG(Loader, "Type: %i Vaddr: %08X Filesz: %8X Memsz: %8X ", p->p_type, p->p_vaddr,
                   p->p_filesz, p->p_memsz);
 
         if (p->p_type == PT_LOAD) {
-            segment_addr[i] = base_addr + p->p_vaddr;
-            memcpy(Memory::GetPointer(segment_addr[i]), GetSegmentPtr(i), p->p_filesz);
-            LOG_DEBUG(Loader, "Loadable Segment Copied to %08x, size %08x", segment_addr[i],
-                      p->p_memsz);
+            CodeSet::Segment* codeset_segment;
+            u32 permission_flags = p->p_flags & (PF_R | PF_W | PF_X);
+            if (permission_flags == (PF_R | PF_X)) {
+                codeset_segment = &codeset->code;
+            } else if (permission_flags == (PF_R)) {
+                codeset_segment = &codeset->rodata;
+            } else if (permission_flags == (PF_R | PF_W)) {
+                codeset_segment = &codeset->data;
+            } else {
+                LOG_ERROR(Loader, "Unexpected ELF PT_LOAD segment id %u with flags %X", i, p->p_flags);
+                continue;
+            }
+
+            if (codeset_segment->size != 0) {
+                LOG_ERROR(Loader, "ELF has more than one segment of the same type. Skipping extra segment (id %i)", i);
+                continue;
+            }
+
+            u32 segment_addr = base_addr + p->p_vaddr;
+            u32 aligned_size = (p->p_memsz + 0xFFF) & ~0xFFF;
+
+            codeset_segment->offset = current_image_position;
+            codeset_segment->addr = segment_addr;
+            codeset_segment->size = aligned_size;
+
+            memcpy(&program_image[current_image_position], GetSegmentPtr(i), p->p_filesz);
+            current_image_position += aligned_size;
         }
     }
+
+    codeset->entrypoint = base_addr + header->e_entry;
+    codeset->memory = std::make_shared<std::vector<u8>>(std::move(program_image));
+
     LOG_DEBUG(Loader, "Done loading.");
+
+    return codeset;
 }
 
 SectionID ElfReader::GetSectionByName(const char *name, int firstSection) const {
@@ -352,18 +403,18 @@ ResultStatus AppLoader_ELF::Load() {
     if (file->ReadBytes(&buffer[0], size) != size)
         return ResultStatus::Error;
 
-    Kernel::g_current_process = Kernel::Process::Create(filename, 0);
+    ElfReader elf_reader(&buffer[0]);
+    SharedPtr<CodeSet> codeset = elf_reader.LoadInto(Memory::PROCESS_IMAGE_VADDR);
+    codeset->name = filename;
+
+    Kernel::g_current_process = Kernel::Process::Create(std::move(codeset));
     Kernel::g_current_process->svc_access_mask.set();
     Kernel::g_current_process->address_mappings = default_address_mappings;
 
     // Attach the default resource limit (APPLICATION) to the process
     Kernel::g_current_process->resource_limit = Kernel::ResourceLimit::GetForCategory(Kernel::ResourceLimitCategory::APPLICATION);
 
-    ElfReader elf_reader(&buffer[0]);
-    elf_reader.LoadInto(Memory::PROCESS_IMAGE_VADDR);
-    // TODO: Fill application title
-
-    Kernel::g_current_process->Run(elf_reader.GetEntryPoint(), 48, Kernel::DEFAULT_STACK_SIZE);
+    Kernel::g_current_process->Run(48, Kernel::DEFAULT_STACK_SIZE);
 
     is_loaded = true;
     return ResultStatus::Success;
diff --git a/src/core/loader/ncch.cpp b/src/core/loader/ncch.cpp
index 2b26b31cf..87603d198 100644
--- a/src/core/loader/ncch.cpp
+++ b/src/core/loader/ncch.cpp
@@ -118,6 +118,9 @@ FileType AppLoader_NCCH::IdentifyType(FileUtil::IOFile& file) {
 }
 
 ResultStatus AppLoader_NCCH::LoadExec() const {
+    using Kernel::SharedPtr;
+    using Kernel::CodeSet;
+
     if (!is_loaded)
         return ResultStatus::ErrorNotLoaded;
 
@@ -126,7 +129,30 @@ ResultStatus AppLoader_NCCH::LoadExec() const {
         std::string process_name = Common::StringFromFixedZeroTerminatedBuffer(
                 (const char*)exheader_header.codeset_info.name, 8);
         u64 program_id = *reinterpret_cast<u64_le const*>(&ncch_header.program_id[0]);
-        Kernel::g_current_process = Kernel::Process::Create(process_name, program_id);
+
+        SharedPtr<CodeSet> codeset = CodeSet::Create(process_name, program_id);
+
+        codeset->code.offset = 0;
+        codeset->code.addr = exheader_header.codeset_info.text.address;
+        codeset->code.size = exheader_header.codeset_info.text.num_max_pages * Memory::PAGE_SIZE;
+
+        codeset->rodata.offset = codeset->code.offset + codeset->code.size;
+        codeset->rodata.addr = exheader_header.codeset_info.ro.address;
+        codeset->rodata.size = exheader_header.codeset_info.ro.num_max_pages * Memory::PAGE_SIZE;
+
+        // TODO(yuriks): Not sure if the bss size is added to the page-aligned .data size or just
+        //               to the regular size. Playing it safe for now.
+        u32 bss_page_size = (exheader_header.codeset_info.bss_size + 0xFFF) & ~0xFFF;
+        code.resize(code.size() + bss_page_size, 0);
+
+        codeset->data.offset = codeset->rodata.offset + codeset->rodata.size;
+        codeset->data.addr = exheader_header.codeset_info.data.address;
+        codeset->data.size = exheader_header.codeset_info.data.num_max_pages * Memory::PAGE_SIZE + bss_page_size;
+
+        codeset->entrypoint = codeset->code.addr;
+        codeset->memory = std::make_shared<std::vector<u8>>(std::move(code));
+
+        Kernel::g_current_process = Kernel::Process::Create(std::move(codeset));
 
         // Attach a resource limit to the process based on the resource limit category
         Kernel::g_current_process->resource_limit = Kernel::ResourceLimit::GetForCategory(
@@ -137,11 +163,9 @@ ResultStatus AppLoader_NCCH::LoadExec() const {
         std::copy_n(exheader_header.arm11_kernel_caps.descriptors, kernel_caps.size(), begin(kernel_caps));
         Kernel::g_current_process->ParseKernelCaps(kernel_caps.data(), kernel_caps.size());
 
-        Memory::WriteBlock(entry_point, &code[0], code.size());
-
         s32 priority = exheader_header.arm11_system_local_caps.priority;
         u32 stack_size = exheader_header.codeset_info.stack_size;
-        Kernel::g_current_process->Run(entry_point, priority, stack_size);
+        Kernel::g_current_process->Run(priority, stack_size);
         return ResultStatus::Success;
     }
     return ResultStatus::Error;