5 files changed, 145 insertions, 34 deletions

diff --git a/src/core/file_sys/patch_manager.cpp b/src/core/file_sys/patch_manager.cpp
index 4b3b5e665..57b7741f8 100644
--- a/src/core/file_sys/patch_manager.cpp
+++ b/src/core/file_sys/patch_manager.cpp
@@ -2,6 +2,7 @@
 // Licensed under GPLv2 or any later version
 // Refer to the license.txt file included.
 
+#include <algorithm>
 #include <array>
 #include <cstddef>
 
@@ -18,6 +19,7 @@
 namespace FileSys {
 
 constexpr u64 SINGLE_BYTE_MODULUS = 0x100;
+constexpr u64 DLC_BASE_TITLE_ID_MASK = 0xFFFFFFFFFFFFE000;
 
 std::string FormatTitleVersion(u32 version, TitleVersionFormat format) {
     std::array<u8, sizeof(u32)> bytes{};
@@ -32,9 +34,10 @@ std::string FormatTitleVersion(u32 version, TitleVersionFormat format) {
     return fmt::format("v{}.{}.{}", bytes[3], bytes[2], bytes[1]);
 }
 
-constexpr std::array<const char*, 2> PATCH_TYPE_NAMES{
+constexpr std::array<const char*, 3> PATCH_TYPE_NAMES{
     "Update",
     "LayeredFS",
+    "DLC",
 };
 
 std::string FormatPatchTypeName(PatchType type) {
@@ -139,6 +142,7 @@ std::map<PatchType, std::string> PatchManager::GetPatchVersionNames() const {
     std::map<PatchType, std::string> out;
     const auto installed = Service::FileSystem::GetUnionContents();
 
+    // Game Updates
     const auto update_tid = GetUpdateTitleID(title_id);
     PatchManager update{update_tid};
     auto [nacp, discard_icon_file] = update.GetControlMetadata();
@@ -157,10 +161,34 @@ std::map<PatchType, std::string> PatchManager::GetPatchVersionNames() const {
         }
     }
 
+    // LayeredFS
     const auto lfs_dir = Service::FileSystem::GetModificationLoadRoot(title_id);
     if (lfs_dir != nullptr && lfs_dir->GetSize() > 0)
         out.insert_or_assign(PatchType::LayeredFS, "");
 
+    // DLC
+    const auto dlc_entries = installed->ListEntriesFilter(TitleType::AOC, ContentRecordType::Data);
+    std::vector<RegisteredCacheEntry> dlc_match;
+    dlc_match.reserve(dlc_entries.size());
+    std::copy_if(dlc_entries.begin(), dlc_entries.end(), std::back_inserter(dlc_match),
+                 [this, &installed](const RegisteredCacheEntry& entry) {
+                     return (entry.title_id & DLC_BASE_TITLE_ID_MASK) == title_id &&
+                            installed->GetEntry(entry)->GetStatus() ==
+                                Loader::ResultStatus::Success;
+                 });
+    if (!dlc_match.empty()) {
+        // Ensure sorted so DLC IDs show in order.
+        std::sort(dlc_match.begin(), dlc_match.end());
+
+        std::string list;
+        for (size_t i = 0; i < dlc_match.size() - 1; ++i)
+            list += fmt::format("{}, ", dlc_match[i].title_id & 0x7FF);
+
+        list += fmt::format("{}", dlc_match.back().title_id & 0x7FF);
+
+        out.insert_or_assign(PatchType::DLC, std::move(list));
+    }
+
     return out;
 }
 
diff --git a/src/core/file_sys/patch_manager.h b/src/core/file_sys/patch_manager.h
index 464f17515..3a2a9d212 100644
--- a/src/core/file_sys/patch_manager.h
+++ b/src/core/file_sys/patch_manager.h
@@ -27,6 +27,7 @@ std::string FormatTitleVersion(u32 version,
 enum class PatchType {
     Update,
     LayeredFS,
+    DLC,
 };
 
 std::string FormatPatchTypeName(PatchType type);
diff --git a/src/core/file_sys/romfs_factory.cpp b/src/core/file_sys/romfs_factory.cpp
index d027a8d59..4994c2532 100644
--- a/src/core/file_sys/romfs_factory.cpp
+++ b/src/core/file_sys/romfs_factory.cpp
@@ -39,36 +39,35 @@ ResultVal<VirtualFile> RomFSFactory::OpenCurrentProcess() {
 }
 
 ResultVal<VirtualFile> RomFSFactory::Open(u64 title_id, StorageId storage, ContentRecordType type) {
+    std::shared_ptr<NCA> res;
+
     switch (storage) {
-    case StorageId::NandSystem: {
-        const auto res = Service::FileSystem::GetSystemNANDContents()->GetEntry(title_id, type);
-        if (res == nullptr) {
-            // TODO(DarkLordZach): Find the right error code to use here
-            return ResultCode(-1);
-        }
-        const auto romfs = res->GetRomFS();
-        if (romfs == nullptr) {
-            // TODO(DarkLordZach): Find the right error code to use here
-            return ResultCode(-1);
-        }
-        return MakeResult<VirtualFile>(romfs);
-    }
-    case StorageId::NandUser: {
-        const auto res = Service::FileSystem::GetUserNANDContents()->GetEntry(title_id, type);
-        if (res == nullptr) {
-            // TODO(DarkLordZach): Find the right error code to use here
-            return ResultCode(-1);
-        }
-        const auto romfs = res->GetRomFS();
-        if (romfs == nullptr) {
-            // TODO(DarkLordZach): Find the right error code to use here
-            return ResultCode(-1);
-        }
-        return MakeResult<VirtualFile>(romfs);
-    }
+    case StorageId::None:
+        res = Service::FileSystem::GetUnionContents()->GetEntry(title_id, type);
+        break;
+    case StorageId::NandSystem:
+        res = Service::FileSystem::GetSystemNANDContents()->GetEntry(title_id, type);
+        break;
+    case StorageId::NandUser:
+        res = Service::FileSystem::GetUserNANDContents()->GetEntry(title_id, type);
+        break;
+    case StorageId::SdCard:
+        res = Service::FileSystem::GetSDMCContents()->GetEntry(title_id, type);
+        break;
     default:
         UNIMPLEMENTED_MSG("Unimplemented storage_id={:02X}", static_cast<u8>(storage));
     }
+
+    if (res == nullptr) {
+        // TODO(DarkLordZach): Find the right error code to use here
+        return ResultCode(-1);
+    }
+    const auto romfs = res->GetRomFS();
+    if (romfs == nullptr) {
+        // TODO(DarkLordZach): Find the right error code to use here
+        return ResultCode(-1);
+    }
+    return MakeResult<VirtualFile>(romfs);
 }
 
 } // namespace FileSys
diff --git a/src/core/hle/service/aoc/aoc_u.cpp b/src/core/hle/service/aoc/aoc_u.cpp
index d9eeac9ec..cfc28fa0c 100644
--- a/src/core/hle/service/aoc/aoc_u.cpp
+++ b/src/core/hle/service/aoc/aoc_u.cpp
@@ -2,22 +2,57 @@
 // Licensed under GPLv2 or any later version
 // Refer to the license.txt file included.
 
+#include <algorithm>
+#include <numeric>
+#include <vector>
 #include "common/logging/log.h"
+#include "core/file_sys/content_archive.h"
+#include "core/file_sys/nca_metadata.h"
+#include "core/file_sys/partition_filesystem.h"
+#include "core/file_sys/registered_cache.h"
 #include "core/hle/ipc_helpers.h"
+#include "core/hle/kernel/process.h"
 #include "core/hle/service/aoc/aoc_u.h"
+#include "core/hle/service/filesystem/filesystem.h"
+#include "core/loader/loader.h"
 
 namespace Service::AOC {
 
-AOC_U::AOC_U() : ServiceFramework("aoc:u") {
+constexpr u64 DLC_BASE_TITLE_ID_MASK = 0xFFFFFFFFFFFFE000;
+constexpr u64 DLC_BASE_TO_AOC_ID_MASK = 0x1000;
+
+static bool CheckAOCTitleIDMatchesBase(u64 base, u64 aoc) {
+    return (aoc & DLC_BASE_TITLE_ID_MASK) == base;
+}
+
+static std::vector<u64> AccumulateAOCTitleIDs() {
+    std::vector<u64> add_on_content;
+    const auto rcu = FileSystem::GetUnionContents();
+    const auto list =
+        rcu->ListEntriesFilter(FileSys::TitleType::AOC, FileSys::ContentRecordType::Data);
+    std::transform(list.begin(), list.end(), std::back_inserter(add_on_content),
+                   [](const FileSys::RegisteredCacheEntry& rce) { return rce.title_id; });
+    add_on_content.erase(
+        std::remove_if(
+            add_on_content.begin(), add_on_content.end(),
+            [&rcu](u64 tid) {
+                return rcu->GetEntry(tid, FileSys::ContentRecordType::Data)->GetStatus() !=
+                       Loader::ResultStatus::Success;
+            }),
+        add_on_content.end());
+    return add_on_content;
+}
+
+AOC_U::AOC_U() : ServiceFramework("aoc:u"), add_on_content(AccumulateAOCTitleIDs()) {
     static const FunctionInfo functions[] = {
         {0, nullptr, "CountAddOnContentByApplicationId"},
         {1, nullptr, "ListAddOnContentByApplicationId"},
         {2, &AOC_U::CountAddOnContent, "CountAddOnContent"},
         {3, &AOC_U::ListAddOnContent, "ListAddOnContent"},
         {4, nullptr, "GetAddOnContentBaseIdByApplicationId"},
-        {5, nullptr, "GetAddOnContentBaseId"},
+        {5, &AOC_U::GetAddOnContentBaseId, "GetAddOnContentBaseId"},
         {6, nullptr, "PrepareAddOnContentByApplicationId"},
-        {7, nullptr, "PrepareAddOnContent"},
+        {7, &AOC_U::PrepareAddOnContent, "PrepareAddOnContent"},
         {8, nullptr, "GetAddOnContentListChangedEvent"},
     };
     RegisterHandlers(functions);
@@ -28,15 +63,59 @@ AOC_U::~AOC_U() = default;
 void AOC_U::CountAddOnContent(Kernel::HLERequestContext& ctx) {
     IPC::ResponseBuilder rb{ctx, 4};
     rb.Push(RESULT_SUCCESS);
-    rb.Push<u64>(0);
-    LOG_WARNING(Service_AOC, "(STUBBED) called");
+
+    const auto current = Core::System::GetInstance().CurrentProcess()->GetTitleID();
+    rb.Push<u32>(std::count_if(add_on_content.begin(), add_on_content.end(), [&current](u64 tid) {
+        return (tid & DLC_BASE_TITLE_ID_MASK) == current;
+    }));
 }
 
 void AOC_U::ListAddOnContent(Kernel::HLERequestContext& ctx) {
+    IPC::RequestParser rp{ctx};
+
+    const auto offset = rp.PopRaw<u32>();
+    auto count = rp.PopRaw<u32>();
+
+    const auto current = Core::System::GetInstance().CurrentProcess()->GetTitleID();
+
+    std::vector<u32> out;
+    for (size_t i = 0; i < add_on_content.size(); ++i) {
+        if ((add_on_content[i] & DLC_BASE_TITLE_ID_MASK) == current)
+            out.push_back(static_cast<u32>(add_on_content[i] & 0x7FF));
+    }
+
+    if (out.size() <= offset) {
+        IPC::ResponseBuilder rb{ctx, 2};
+        // TODO(DarkLordZach): Find the correct error code.
+        rb.Push(ResultCode(-1));
+        return;
+    }
+
+    count = std::min<size_t>(out.size() - offset, count);
+    std::rotate(out.begin(), out.begin() + offset, out.end());
+    out.resize(count);
+
+    ctx.WriteBuffer(out);
+
+    IPC::ResponseBuilder rb{ctx, 2};
+    rb.Push(RESULT_SUCCESS);
+}
+
+void AOC_U::GetAddOnContentBaseId(Kernel::HLERequestContext& ctx) {
     IPC::ResponseBuilder rb{ctx, 4};
     rb.Push(RESULT_SUCCESS);
-    rb.Push<u64>(0);
-    LOG_WARNING(Service_AOC, "(STUBBED) called");
+    rb.Push(Core::System::GetInstance().CurrentProcess()->GetTitleID() | DLC_BASE_TO_AOC_ID_MASK);
+}
+
+void AOC_U::PrepareAddOnContent(Kernel::HLERequestContext& ctx) {
+    IPC::RequestParser rp{ctx};
+
+    const auto aoc_id = rp.PopRaw<u32>();
+
+    LOG_WARNING(Service_AOC, "(STUBBED) called with aoc_id={:08X}", aoc_id);
+
+    IPC::ResponseBuilder rb{ctx, 2};
+    rb.Push(RESULT_SUCCESS);
 }
 
 void InstallInterfaces(SM::ServiceManager& service_manager) {
diff --git a/src/core/hle/service/aoc/aoc_u.h b/src/core/hle/service/aoc/aoc_u.h
index 29ce8f488..b3c7cab7a 100644
--- a/src/core/hle/service/aoc/aoc_u.h
+++ b/src/core/hle/service/aoc/aoc_u.h
@@ -16,6 +16,10 @@ public:
 private:
     void CountAddOnContent(Kernel::HLERequestContext& ctx);
     void ListAddOnContent(Kernel::HLERequestContext& ctx);
+    void GetAddOnContentBaseId(Kernel::HLERequestContext& ctx);
+    void PrepareAddOnContent(Kernel::HLERequestContext& ctx);
+
+    std::vector<u64> add_on_content;
 };
 
 /// Registers all AOC services with the specified service manager.