Add option to configure to download system files from Nintendo Update Service (#6269)

Co-authored-by: B3n30 <benediktthomas@gmail.com>

src/core/file_sys/cia_container.cpp

@@ -16,8 +16,6 @@
 
 namespace FileSys {
 
-constexpr u32 CIA_SECTION_ALIGNMENT = 0x40;
-
 Loader::ResultStatus CIAContainer::Load(const FileBackend& backend) {
     std::vector<u8> header_data(sizeof(Header));
 

src/core/file_sys/cia_container.h

@@ -29,6 +29,7 @@ constexpr std::size_t CIA_CONTENT_BITS_SIZE = (CIA_CONTENT_MAX_COUNT / 8);
 constexpr std::size_t CIA_HEADER_SIZE = 0x2020;
 constexpr std::size_t CIA_DEPENDENCY_SIZE = 0x300;
 constexpr std::size_t CIA_METADATA_SIZE = 0x400;
+constexpr u32 CIA_SECTION_ALIGNMENT = 0x40;
 
 /**
  * Helper which implements an interface to read and write CTR Installable Archive (CIA) files.
@@ -69,7 +70,6 @@ public:
 
     void Print() const;
 
-private:
     struct Header {
         u32_le header_size;
         u16_le type;
@@ -87,10 +87,14 @@ private:
             // The bits in the content index are arranged w/ index 0 as the MSB, 7 as the LSB, etc.
             return (content_present[index >> 3] & (0x80 >> (index & 7))) != 0;
         }
+        void SetContentPresent(u16 index) {
+            content_present[index >> 3] |= (0x80 >> (index & 7));
+        }
     };
 
     static_assert(sizeof(Header) == CIA_HEADER_SIZE, "CIA Header structure size is wrong");
 
+private:
     struct Metadata {
         std::array<u64_le, 0x30> dependencies;
         std::array<u8, 0x180> reserved;

src/core/hle/service/am/am.cpp

@@ -9,6 +9,7 @@
 #include <cryptopp/aes.h>
 #include <cryptopp/modes.h>
 #include <fmt/format.h>
+#include "common/alignment.h"
 #include "common/common_paths.h"
 #include "common/file_util.h"
 #include "common/logging/log.h"
@@ -31,6 +32,9 @@
 #include "core/hle/service/fs/fs_user.h"
 #include "core/loader/loader.h"
 #include "core/loader/smdh.h"
+#ifdef ENABLE_WEB_SERVICE
+#include "web_service/nus_download.h"
+#endif
 
 namespace Service::AM {
 
@@ -138,6 +142,8 @@ ResultCode CIAFile::WriteTitleMetadata() {
             decryption_state->content[i].SetKeyWithIV(title_key->data(), title_key->size(),
                                                       ctr.data());
         }
+    } else {
+        LOG_ERROR(Service_AM, "Can't get title key from ticket");
     }
 
     install_state = CIAInstallState::TMDLoaded;
@@ -180,6 +186,11 @@ ResultVal<std::size_t> CIAFile::WriteContentData(u64 offset, std::size_t length,
                                  buffer + (range_min - offset) + available_to_write);
 
             if ((tmd.GetContentTypeByIndex(i) & FileSys::TMDContentTypeFlag::Encrypted) != 0) {
+                if (decryption_state->content.size() <= i) {
+                    // TODO: There is probably no correct error to return here. What error should be
+                    // returned?
+                    return FileSys::ERROR_INSUFFICIENT_SPACE;
+                }
                 decryption_state->content[i].ProcessData(temp.data(), temp.data(), temp.size());
             }
 
@@ -235,7 +246,7 @@ ResultVal<std::size_t> CIAFile::Write(u64 offset, std::size_t length, bool flush
         std::size_t buf_offset = buf_loaded - offset;
         std::size_t buf_copy_size =
             std::min(length, static_cast<std::size_t>(container.GetContentOffset() - offset)) -
-            buf_loaded;
+            buf_offset;
         std::size_t buf_max_size = std::min(offset + length, container.GetContentOffset());
         data.resize(buf_max_size);
         memcpy(data.data() + copy_offset, buffer + buf_offset, buf_copy_size);
@@ -418,6 +429,99 @@ InstallStatus InstallCIA(const std::string& path,
     return InstallStatus::ErrorInvalid;
 }
 
+InstallStatus InstallFromNus(u64 title_id, int version) {
+#ifdef ENABLE_WEB_SERVICE
+    LOG_DEBUG(Service_AM, "Downloading {:X}", title_id);
+
+    CIAFile install_file{GetTitleMediaType(title_id)};
+
+    std::string path = fmt::format("/ccs/download/{:016X}/tmd", title_id);
+    if (version != -1) {
+        path += fmt::format(".{}", version);
+    }
+    auto tmd_response = WebService::NUS::Download(path);
+    if (!tmd_response) {
+        LOG_ERROR(Service_AM, "Failed to download tmd for {:016X}", title_id);
+        return InstallStatus::ErrorFileNotFound;
+    }
+    FileSys::TitleMetadata tmd;
+    tmd.Load(*tmd_response);
+
+    path = fmt::format("/ccs/download/{:016X}/cetk", title_id);
+    auto cetk_response = WebService::NUS::Download(path);
+    if (!cetk_response) {
+        LOG_ERROR(Service_AM, "Failed to download cetk for {:016X}", title_id);
+        return InstallStatus::ErrorFileNotFound;
+    }
+
+    std::vector<u8> content;
+    const auto content_count = tmd.GetContentCount();
+    for (std::size_t i = 0; i < content_count; ++i) {
+        const std::string filename = fmt::format("{:08x}", tmd.GetContentIDByIndex(i));
+        path = fmt::format("/ccs/download/{:016X}/{}", title_id, filename);
+        const auto temp_response = WebService::NUS::Download(path);
+        if (!temp_response) {
+            LOG_ERROR(Service_AM, "Failed to download content for {:016X}", title_id);
+            return InstallStatus::ErrorFileNotFound;
+        }
+        content.insert(content.end(), temp_response->begin(), temp_response->end());
+    }
+
+    FileSys::CIAContainer::Header fake_header{
+        .header_size = sizeof(FileSys::CIAContainer::Header),
+        .type = 0,
+        .version = 0,
+        .cert_size = 0,
+        .tik_size = static_cast<u32_le>(cetk_response->size()),
+        .tmd_size = static_cast<u32_le>(tmd_response->size()),
+        .meta_size = 0,
+    };
+    for (u16 i = 0; i < content_count; ++i) {
+        fake_header.SetContentPresent(i);
+    }
+    std::vector<u8> header_data(sizeof(fake_header));
+    std::memcpy(header_data.data(), &fake_header, sizeof(fake_header));
+
+    std::size_t current_offset = 0;
+    const auto write_to_cia_file_aligned = [&install_file, &current_offset](std::vector<u8>& data) {
+        const u64 offset =
+            Common::AlignUp(current_offset + data.size(), FileSys::CIA_SECTION_ALIGNMENT);
+        data.resize(offset - current_offset, 0);
+        const auto result = install_file.Write(current_offset, data.size(), true, data.data());
+        if (result.Failed()) {
+            LOG_ERROR(Service_AM, "CIA file installation aborted with error code {:08x}",
+                      result.Code().raw);
+            return InstallStatus::ErrorAborted;
+        }
+        current_offset += data.size();
+        return InstallStatus::Success;
+    };
+
+    auto result = write_to_cia_file_aligned(header_data);
+    if (result != InstallStatus::Success) {
+        return result;
+    }
+
+    result = write_to_cia_file_aligned(*cetk_response);
+    if (result != InstallStatus::Success) {
+        return result;
+    }
+
+    result = write_to_cia_file_aligned(*tmd_response);
+    if (result != InstallStatus::Success) {
+        return result;
+    }
+
+    result = write_to_cia_file_aligned(content);
+    if (result != InstallStatus::Success) {
+        return result;
+    }
+    return InstallStatus::Success;
+#else
+    return InstallStatus::ErrorFileNotFound;
+#endif
+}
+
 Service::FS::MediaType GetTitleMediaType(u64 titleId) {
     u16 platform = static_cast<u16>(titleId >> 48);
     u16 category = static_cast<u16>((titleId >> 32) & 0xFFFF);

src/core/hle/service/am/am.h

@@ -110,6 +110,13 @@ private:
 InstallStatus InstallCIA(const std::string& path,
                          std::function<ProgressCallback>&& update_callback = nullptr);
 
+/**
+ * Downloads and installs title form the Nintendo Update Service.
+ * @param title_id the title_id to download
+ * @returns  whether the install was successful or error code
+ */
+InstallStatus InstallFromNus(u64 title_id, int version = -1);
+
 /**
  * Get the mediatype for an installed title
  * @param titleId the installed title ID

src/core/hw/aes/key.cpp

@@ -88,7 +88,7 @@ struct KeySlot {
 };
 
 std::array<KeySlot, KeySlotID::MaxKeySlotID> key_slots;
-std::array<std::optional<AESKey>, 6> common_key_y_slots;
+std::array<std::optional<AESKey>, MaxCommonKeySlot> common_key_y_slots;
 
 enum class FirmwareType : u32 {
     ARM9 = 0,  // uses NDMA
@@ -494,9 +494,9 @@ void LoadPresetKeys() {
 
 } // namespace
 
-void InitKeys() {
+void InitKeys(bool force) {
     static bool initialized = false;
-    if (initialized)
+    if (initialized && !force)
         return;
     initialized = true;
     HW::RSA::InitSlots();

src/core/hw/aes/key.h

@@ -48,11 +48,13 @@ enum KeySlotID : std::size_t {
     MaxKeySlotID = 0x40,
 };
 
+constexpr std::size_t MaxCommonKeySlot = 6;
+
 constexpr std::size_t AES_BLOCK_SIZE = 16;
 
 using AESKey = std::array<u8, AES_BLOCK_SIZE>;
 
-void InitKeys();
+void InitKeys(bool force = false);
 
 void SetGeneratorConstant(const AESKey& key);
 void SetKeyX(std::size_t slot_id, const AESKey& key);