2 files changed, 160 insertions, 19 deletions

diff --git a/src/core/file_sys/content_archive.cpp b/src/core/file_sys/content_archive.cpp
index 7cfb6f36b..f0d376bf5 100644
--- a/src/core/file_sys/content_archive.cpp
+++ b/src/core/file_sys/content_archive.cpp
@@ -12,6 +12,7 @@
 #include "core/crypto/aes_util.h"
 #include "core/crypto/ctr_encryption_layer.h"
 #include "core/file_sys/content_archive.h"
+#include "core/file_sys/nca_patch.h"
 #include "core/file_sys/partition_filesystem.h"
 #include "core/file_sys/romfs.h"
 #include "core/file_sys/vfs_offset.h"
@@ -68,10 +69,31 @@ struct RomFSSuperblock {
 };
 static_assert(sizeof(RomFSSuperblock) == 0x200, "RomFSSuperblock has incorrect size.");
 
+struct BKTRHeader {
+    u64_le offset;
+    u64_le size;
+    u32_le magic;
+    INSERT_PADDING_BYTES(0x4);
+    u32_le number_entries;
+    INSERT_PADDING_BYTES(0x4);
+};
+static_assert(sizeof(BKTRHeader) == 0x20, "BKTRHeader has incorrect size.");
+
+struct BKTRSuperblock {
+    NCASectionHeaderBlock header_block;
+    IVFCHeader ivfc;
+    INSERT_PADDING_BYTES(0x18);
+    BKTRHeader relocation;
+    BKTRHeader subsection;
+    INSERT_PADDING_BYTES(0xC0);
+};
+static_assert(sizeof(BKTRSuperblock) == 0x200, "BKTRSuperblock has incorrect size.");
+
 union NCASectionHeader {
     NCASectionRaw raw;
     PFS0Superblock pfs0;
     RomFSSuperblock romfs;
+    BKTRSuperblock bktr;
 };
 static_assert(sizeof(NCASectionHeader) == 0x200, "NCASectionHeader has incorrect size.");
 
@@ -104,7 +126,7 @@ boost::optional<Core::Crypto::Key128> NCA::GetKeyAreaKey(NCASectionCryptoType ty
     Core::Crypto::Key128 out;
     if (type == NCASectionCryptoType::XTS)
         std::copy(key_area.begin(), key_area.begin() + 0x10, out.begin());
-    else if (type == NCASectionCryptoType::CTR)
+    else if (type == NCASectionCryptoType::CTR || type == NCASectionCryptoType::BKTR)
         std::copy(key_area.begin() + 0x20, key_area.begin() + 0x30, out.begin());
     else
         LOG_CRITICAL(Crypto, "Called GetKeyAreaKey on invalid NCASectionCryptoType type={:02X}",
@@ -154,6 +176,9 @@ VirtualFile NCA::Decrypt(NCASectionHeader s_header, VirtualFile in, u64 starting
         LOG_DEBUG(Crypto, "called with mode=NONE");
         return in;
     case NCASectionCryptoType::CTR:
+    // During normal BKTR decryption, this entire function is skipped. This is for the metadata,
+    // which uses the same CTR as usual.
+    case NCASectionCryptoType::BKTR:
         LOG_DEBUG(Crypto, "called with mode=CTR, starting_offset={:016X}", starting_offset);
         {
             boost::optional<Core::Crypto::Key128> key = boost::none;
@@ -190,7 +215,9 @@ VirtualFile NCA::Decrypt(NCASectionHeader s_header, VirtualFile in, u64 starting
     }
 }
 
-NCA::NCA(VirtualFile file_) : file(std::move(file_)) {
+NCA::NCA(VirtualFile file_, VirtualFile bktr_base_romfs_)
+    : file(std::move(file_)),
+      bktr_base_romfs(bktr_base_romfs_ ? std::move(bktr_base_romfs_) : nullptr) {
     status = Loader::ResultStatus::Success;
 
     if (file == nullptr) {
@@ -270,17 +297,15 @@ NCA::NCA(VirtualFile file_) : file(std::move(file_)) {
         auto section = sections[i];
 
         if (section.raw.header.filesystem_type == NCASectionFilesystemType::ROMFS) {
+            const size_t base_offset =
+                header.section_tables[i].media_offset * MEDIA_OFFSET_MULTIPLIER;
             const size_t romfs_offset =
-                header.section_tables[i].media_offset * MEDIA_OFFSET_MULTIPLIER +
-                section.romfs.ivfc.levels[IVFC_MAX_LEVEL - 1].offset;
+                base_offset + section.romfs.ivfc.levels[IVFC_MAX_LEVEL - 1].offset;
             const size_t romfs_size = section.romfs.ivfc.levels[IVFC_MAX_LEVEL - 1].size;
-            auto dec =
-                Decrypt(section, std::make_shared<OffsetVfsFile>(file, romfs_size, romfs_offset),
-                        romfs_offset);
-            if (dec != nullptr) {
-                files.push_back(std::move(dec));
-                romfs = files.back();
-            } else {
+            auto raw = std::make_shared<OffsetVfsFile>(file, romfs_size, romfs_offset);
+            auto dec = Decrypt(section, raw, romfs_offset);
+
+            if (dec == nullptr) {
                 if (status != Loader::ResultStatus::Success)
                     return;
                 if (has_rights_id)
@@ -289,6 +314,120 @@ NCA::NCA(VirtualFile file_) : file(std::move(file_)) {
                     status = Loader::ResultStatus::ErrorIncorrectKeyAreaKey;
                 return;
             }
+
+            if (section.raw.header.crypto_type == NCASectionCryptoType::BKTR) {
+                if (section.bktr.relocation.magic != Common::MakeMagic('B', 'K', 'T', 'R') ||
+                    section.bktr.subsection.magic != Common::MakeMagic('B', 'K', 'T', 'R')) {
+                    status = Loader::ResultStatus::ErrorBadBKTRHeader;
+                    return;
+                }
+
+                if (section.bktr.relocation.offset + section.bktr.relocation.size !=
+                    section.bktr.subsection.offset) {
+                    status = Loader::ResultStatus::ErrorBKTRSubsectionNotAfterRelocation;
+                    return;
+                }
+
+                const u64 size =
+                    MEDIA_OFFSET_MULTIPLIER * (header.section_tables[i].media_end_offset -
+                                               header.section_tables[i].media_offset);
+                if (section.bktr.subsection.offset + section.bktr.subsection.size != size) {
+                    status = Loader::ResultStatus::ErrorBKTRSubsectionNotAtEnd;
+                    return;
+                }
+
+                const u64 offset = section.romfs.ivfc.levels[IVFC_MAX_LEVEL - 1].offset;
+                RelocationBlock relocation_block{};
+                if (dec->ReadObject(&relocation_block, section.bktr.relocation.offset - offset) !=
+                    sizeof(RelocationBlock)) {
+                    status = Loader::ResultStatus::ErrorBadRelocationBlock;
+                    return;
+                }
+                SubsectionBlock subsection_block{};
+                if (dec->ReadObject(&subsection_block, section.bktr.subsection.offset - offset) !=
+                    sizeof(RelocationBlock)) {
+                    status = Loader::ResultStatus::ErrorBadSubsectionBlock;
+                    return;
+                }
+
+                std::vector<RelocationBucketRaw> relocation_buckets_raw(
+                    (section.bktr.relocation.size - sizeof(RelocationBlock)) /
+                    sizeof(RelocationBucketRaw));
+                if (dec->ReadBytes(relocation_buckets_raw.data(),
+                                   section.bktr.relocation.size - sizeof(RelocationBlock),
+                                   section.bktr.relocation.offset + sizeof(RelocationBlock) -
+                                       offset) !=
+                    section.bktr.relocation.size - sizeof(RelocationBlock)) {
+                    status = Loader::ResultStatus::ErrorBadRelocationBuckets;
+                    return;
+                }
+
+                std::vector<SubsectionBucketRaw> subsection_buckets_raw(
+                    (section.bktr.subsection.size - sizeof(SubsectionBlock)) /
+                    sizeof(SubsectionBucketRaw));
+                if (dec->ReadBytes(subsection_buckets_raw.data(),
+                                   section.bktr.subsection.size - sizeof(SubsectionBlock),
+                                   section.bktr.subsection.offset + sizeof(SubsectionBlock) -
+                                       offset) !=
+                    section.bktr.subsection.size - sizeof(SubsectionBlock)) {
+                    status = Loader::ResultStatus::ErrorBadSubsectionBuckets;
+                    return;
+                }
+
+                std::vector<RelocationBucket> relocation_buckets(relocation_buckets_raw.size());
+                std::transform(relocation_buckets_raw.begin(), relocation_buckets_raw.end(),
+                               relocation_buckets.begin(), &ConvertRelocationBucketRaw);
+                std::vector<SubsectionBucket> subsection_buckets(subsection_buckets_raw.size());
+                std::transform(subsection_buckets_raw.begin(), subsection_buckets_raw.end(),
+                               subsection_buckets.begin(), &ConvertSubsectionBucketRaw);
+
+                u32 ctr_low;
+                std::memcpy(&ctr_low, section.raw.section_ctr.data(), sizeof(ctr_low));
+                subsection_buckets.back().entries.push_back(
+                    {section.bktr.relocation.offset, {0}, ctr_low});
+                subsection_buckets.back().entries.push_back({size, {0}, 0});
+
+                boost::optional<Core::Crypto::Key128> key = boost::none;
+                if (encrypted) {
+                    if (has_rights_id) {
+                        status = Loader::ResultStatus::Success;
+                        key = GetTitlekey();
+                        if (key == boost::none) {
+                            status = Loader::ResultStatus::ErrorMissingTitlekey;
+                            return;
+                        }
+                    } else {
+                        key = GetKeyAreaKey(NCASectionCryptoType::BKTR);
+                        if (key == boost::none) {
+                            status = Loader::ResultStatus::ErrorMissingKeyAreaKey;
+                            return;
+                        }
+                    }
+                }
+
+                if (bktr_base_romfs == nullptr) {
+                    status = Loader::ResultStatus::ErrorMissingBKTRBaseRomFS;
+                    return;
+                }
+
+                auto bktr = std::make_shared<BKTR>(
+                    bktr_base_romfs, std::make_shared<OffsetVfsFile>(file, romfs_size, base_offset),
+                    relocation_block, relocation_buckets, subsection_block, subsection_buckets,
+                    encrypted, encrypted ? key.get() : Core::Crypto::Key128{}, base_offset,
+                    romfs_offset - base_offset, section.raw.section_ctr);
+
+                // BKTR applies to entire IVFC, so make an offset version to level 6
+
+                files.push_back(std::make_shared<OffsetVfsFile>(
+                    bktr, romfs_size, section.romfs.ivfc.levels[IVFC_MAX_LEVEL - 1].offset));
+                romfs = files.back();
+            } else {
+                files.push_back(std::move(dec));
+                romfs = files.back();
+                const u64 raw_size =
+                    MEDIA_OFFSET_MULTIPLIER * (header.section_tables[i].media_end_offset -
+                                               header.section_tables[i].media_offset);
+            }
         } else if (section.raw.header.filesystem_type == NCASectionFilesystemType::PFS0) {
             u64 offset = (static_cast<u64>(header.section_tables[i].media_offset) *
                           MEDIA_OFFSET_MULTIPLIER) +
@@ -349,11 +488,17 @@ NCAContentType NCA::GetType() const {
 }
 
 u64 NCA::GetTitleId() const {
+    if (is_update || status == Loader::ResultStatus::ErrorMissingBKTRBaseRomFS)
+        return header.title_id | 0x800;
     if (status != Loader::ResultStatus::Success)
         return {};
     return header.title_id;
 }
 
+bool NCA::IsUpdate() const {
+    return is_update;
+}
+
 VirtualFile NCA::GetRomFS() const {
     return romfs;
 }
@@ -366,10 +511,6 @@ VirtualFile NCA::GetBaseFile() const {
     return file;
 }
 
-bool NCA::IsUpdate() const {
-    return is_update;
-}
-
 bool NCA::ReplaceFileWithSubdirectory(VirtualFile file, VirtualDir dir) {
     return false;
 }
diff --git a/src/core/file_sys/content_archive.h b/src/core/file_sys/content_archive.h
index 0ea666cac..104226f3a 100644
--- a/src/core/file_sys/content_archive.h
+++ b/src/core/file_sys/content_archive.h
@@ -79,7 +79,7 @@ bool IsValidNCA(const NCAHeader& header);
 // After construction, use GetStatus to determine if the file is valid and ready to be used.
 class NCA : public ReadOnlyVfsDirectory {
 public:
-    explicit NCA(VirtualFile file);
+    explicit NCA(VirtualFile file, VirtualFile bktr_base_romfs = nullptr);
     Loader::ResultStatus GetStatus() const;
 
     std::vector<std::shared_ptr<VfsFile>> GetFiles() const override;
@@ -89,14 +89,13 @@ public:
 
     NCAContentType GetType() const;
     u64 GetTitleId() const;
+    bool IsUpdate() const;
 
     VirtualFile GetRomFS() const;
     VirtualDir GetExeFS() const;
 
     VirtualFile GetBaseFile() const;
 
-    bool IsUpdate() const;
-
 protected:
     bool ReplaceFileWithSubdirectory(VirtualFile file, VirtualDir dir) override;
 
@@ -112,14 +111,15 @@ private:
     VirtualFile romfs = nullptr;
     VirtualDir exefs = nullptr;
     VirtualFile file;
+    VirtualFile bktr_base_romfs;
 
     NCAHeader header{};
     bool has_rights_id{};
-    bool is_update{};
 
     Loader::ResultStatus status{};
 
     bool encrypted;
+    bool is_update;
 
     Core::Crypto::KeyManager keys;
 };