From ab44192ab01a27115e14a52d4823722cbd357c0d Mon Sep 17 00:00:00 2001
From: Zach Hilman <zachhilman@gmail.com>
Date: Thu, 16 Aug 2018 17:08:13 -0400
Subject: [PATCH] file_sys: Implement NAX containers

---
 src/core/CMakeLists.txt           |   6 ++
 src/core/file_sys/xts_archive.cpp | 164 ++++++++++++++++++++++++++++++
 src/core/file_sys/xts_archive.h   |  68 +++++++++++++
 3 files changed, 238 insertions(+)
 create mode 100644 src/core/file_sys/xts_archive.cpp
 create mode 100644 src/core/file_sys/xts_archive.h

diff --git a/src/core/CMakeLists.txt b/src/core/CMakeLists.txt
index 31a7bf6fdc..a74270a0f3 100644
--- a/src/core/CMakeLists.txt
+++ b/src/core/CMakeLists.txt
@@ -20,6 +20,8 @@ add_library(core STATIC
     crypto/key_manager.h
     crypto/ctr_encryption_layer.cpp
     crypto/ctr_encryption_layer.h
+    crypto/xts_encryption_layer.cpp
+    crypto/xts_encryption_layer.h
     file_sys/bis_factory.cpp
     file_sys/bis_factory.h
     file_sys/card_image.cpp
@@ -57,6 +59,8 @@ add_library(core STATIC
     file_sys/vfs_real.h
     file_sys/vfs_vector.cpp
     file_sys/vfs_vector.h
+    file_sys/xts_archive.cpp
+    file_sys/xts_archive.h
     frontend/emu_window.cpp
     frontend/emu_window.h
     frontend/framebuffer_layout.cpp
@@ -347,6 +351,8 @@ add_library(core STATIC
     loader/linker.h
     loader/loader.cpp
     loader/loader.h
+    loader/nax.cpp
+    loader/nax.h
     loader/nca.cpp
     loader/nca.h
     loader/nro.cpp
diff --git a/src/core/file_sys/xts_archive.cpp b/src/core/file_sys/xts_archive.cpp
new file mode 100644
index 0000000000..54be31916c
--- /dev/null
+++ b/src/core/file_sys/xts_archive.cpp
@@ -0,0 +1,164 @@
+// Copyright 2018 yuzu emulator team
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#include <array>
+#include <regex>
+#include <string>
+#include <mbedtls/md.h>
+#include <mbedtls/sha256.h>
+#include "common/assert.h"
+#include "common/hex_util.h"
+#include "common/logging/log.h"
+#include "core/crypto/aes_util.h"
+#include "core/crypto/xts_encryption_layer.h"
+#include "core/file_sys/partition_filesystem.h"
+#include "core/file_sys/vfs_offset.h"
+#include "core/file_sys/xts_archive.h"
+#include "core/loader/loader.h"
+
+namespace FileSys {
+
+template <typename SourceData, typename SourceKey, typename Destination>
+static bool CalculateHMAC256(Destination* out, const SourceKey* key, size_t key_length,
+                             const SourceData* data, size_t data_length) {
+    mbedtls_md_context_t context;
+    mbedtls_md_init(&context);
+
+    const auto key_f = reinterpret_cast<const u8*>(key);
+    const std::vector<u8> key_v(key_f, key_f + key_length);
+
+    if (mbedtls_md_setup(&context, mbedtls_md_info_from_type(MBEDTLS_MD_SHA256), 1) ||
+        mbedtls_md_hmac_starts(&context, reinterpret_cast<const u8*>(key), key_length) ||
+        mbedtls_md_hmac_update(&context, reinterpret_cast<const u8*>(data), data_length) ||
+        mbedtls_md_hmac_finish(&context, reinterpret_cast<u8*>(out))) {
+        mbedtls_md_free(&context);
+        return false;
+    }
+
+    mbedtls_md_free(&context);
+    return true;
+}
+
+NAX::NAX(VirtualFile file_) : file(std::move(file_)), header(std::make_unique<NAXHeader>()) {
+    std::string path = FileUtil::SanitizePath(file->GetFullPath());
+    static const std::regex nax_path_regex("/registered/(000000[0-9A-F]{2})/([0-9A-F]{32})\\.nca",
+                                           std::regex_constants::ECMAScript |
+                                               std::regex_constants::icase);
+    std::smatch match;
+    if (!std::regex_search(path, match, nax_path_regex)) {
+        status = Loader::ResultStatus::ErrorBadNAXFilePath;
+        return;
+    }
+
+    std::string two_dir = match[1];
+    std::string nca_id = match[2];
+    std::transform(two_dir.begin(), two_dir.end(), two_dir.begin(), ::toupper);
+    std::transform(nca_id.begin(), nca_id.end(), nca_id.begin(), ::tolower);
+
+    status = Parse(fmt::format("/registered/{}/{}.nca", two_dir, nca_id));
+}
+
+NAX::NAX(VirtualFile file_, std::array<u8, 0x10> nca_id)
+    : file(std::move(file_)), header(std::make_unique<NAXHeader>()) {
+    Core::Crypto::SHA256Hash hash{};
+    mbedtls_sha256(nca_id.data(), nca_id.size(), hash.data(), 0);
+    status = Parse(fmt::format("/registered/000000{:02X}/{}.nca", hash[0],
+                               Common::HexArrayToString(nca_id, false)));
+}
+
+Loader::ResultStatus NAX::Parse(std::string path) {
+    if (file->ReadObject(header.get()) != sizeof(NAXHeader))
+        return Loader::ResultStatus::ErrorBadNAXHeader;
+
+    if (header->magic != Common::MakeMagic('N', 'A', 'X', '0'))
+        return Loader::ResultStatus::ErrorBadNAXHeader;
+
+    if (file->GetSize() < 0x4000 + header->file_size)
+        return Loader::ResultStatus::ErrorIncorrectNAXFileSize;
+
+    keys.DeriveSDSeedLazy();
+    std::array<Core::Crypto::Key256, 2> sd_keys{};
+    const auto sd_keys_res = Core::Crypto::DeriveSDKeys(sd_keys, keys);
+    if (sd_keys_res != Loader::ResultStatus::Success) {
+        return sd_keys_res;
+    }
+
+    const auto enc_keys = header->key_area;
+
+    size_t i = 0;
+    for (; i < 2; ++i) {
+        std::array<Core::Crypto::Key128, 2> nax_keys{};
+        if (!CalculateHMAC256(nax_keys.data(), sd_keys[i].data(), 0x10, path.c_str(),
+                              path.size())) {
+            return Loader::ResultStatus::ErrorNAXKeyHMACFailed;
+        }
+
+        for (size_t j = 0; j < 2; ++j) {
+            Core::Crypto::AESCipher<Core::Crypto::Key128> cipher(nax_keys[j],
+                                                                 Core::Crypto::Mode::ECB);
+            cipher.Transcode(enc_keys[j].data(), 0x10, header->key_area[j].data(),
+                             Core::Crypto::Op::Decrypt);
+        }
+
+        Core::Crypto::SHA256Hash validation{};
+        if (!CalculateHMAC256(validation.data(), &header->magic, 0x60, sd_keys[i].data() + 0x10,
+                              0x10)) {
+            return Loader::ResultStatus::ErrorNAXValidationHMACFailed;
+        }
+        if (header->hmac == validation)
+            break;
+    }
+
+    if (i == 2) {
+        return Loader::ResultStatus::ErrorNAXKeyDerivationFailed;
+    }
+
+    type = static_cast<NAXContentType>(i);
+
+    Core::Crypto::Key256 final_key{};
+    memcpy(final_key.data(), &header->key_area, 0x20);
+    const auto enc_file = std::make_shared<OffsetVfsFile>(file, header->file_size, 0x4000);
+    dec_file = std::make_shared<Core::Crypto::XTSEncryptionLayer>(enc_file, final_key);
+
+    return Loader::ResultStatus::Success;
+}
+
+Loader::ResultStatus NAX::GetStatus() {
+    return status;
+}
+
+VirtualFile NAX::GetDecrypted() {
+    return dec_file;
+}
+
+std::shared_ptr<NCA> NAX::AsNCA() {
+    if (type == NAXContentType::NCA)
+        return std::make_shared<NCA>(GetDecrypted());
+    return nullptr;
+}
+
+NAXContentType NAX::GetContentType() {
+    return type;
+}
+
+std::vector<std::shared_ptr<VfsFile>> NAX::GetFiles() const {
+    return {dec_file};
+}
+
+std::vector<std::shared_ptr<VfsDirectory>> NAX::GetSubdirectories() const {
+    return {};
+}
+
+std::string NAX::GetName() const {
+    return file->GetName();
+}
+
+std::shared_ptr<VfsDirectory> NAX::GetParentDirectory() const {
+    return file->GetContainingDirectory();
+}
+
+bool NAX::ReplaceFileWithSubdirectory(VirtualFile file, VirtualDir dir) {
+    return false;
+}
+} // namespace FileSys
diff --git a/src/core/file_sys/xts_archive.h b/src/core/file_sys/xts_archive.h
new file mode 100644
index 0000000000..4e44f634ad
--- /dev/null
+++ b/src/core/file_sys/xts_archive.h
@@ -0,0 +1,68 @@
+// Copyright 2018 yuzu emulator team
+// Licensed under GPLv2 or any later version
+// Refer to the license.txt file included.
+
+#pragma once
+
+#include <array>
+#include <vector>
+#include "common/common_types.h"
+#include "common/swap.h"
+#include "core/file_sys/content_archive.h"
+#include "core/file_sys/vfs.h"
+#include "core/loader/loader.h"
+
+namespace FileSys {
+
+struct NAXHeader {
+    std::array<u8, 0x20> hmac;
+    u64_le magic;
+    std::array<Core::Crypto::Key128, 2> key_area;
+    u64_le file_size;
+    INSERT_PADDING_BYTES(0x30);
+};
+static_assert(sizeof(NAXHeader) == 0x80, "NAXHeader has incorrect size.");
+
+enum class NAXContentType : u8 {
+    Save = 0,
+    NCA = 1,
+};
+
+class NAX : public ReadOnlyVfsDirectory {
+public:
+    explicit NAX(VirtualFile file);
+    explicit NAX(VirtualFile file, std::array<u8, 0x10> nca_id);
+
+    Loader::ResultStatus GetStatus();
+
+    VirtualFile GetDecrypted();
+
+    std::shared_ptr<NCA> AsNCA();
+
+    NAXContentType GetContentType();
+
+    std::vector<std::shared_ptr<VfsFile>> GetFiles() const override;
+
+    std::vector<std::shared_ptr<VfsDirectory>> GetSubdirectories() const override;
+
+    std::string GetName() const override;
+
+    std::shared_ptr<VfsDirectory> GetParentDirectory() const override;
+
+protected:
+    bool ReplaceFileWithSubdirectory(VirtualFile file, VirtualDir dir) override;
+
+private:
+    Loader::ResultStatus Parse(std::string path);
+
+    std::unique_ptr<NAXHeader> header;
+
+    VirtualFile file;
+    Loader::ResultStatus status;
+    NAXContentType type;
+
+    VirtualFile dec_file;
+
+    Core::Crypto::KeyManager keys;
+};
+} // namespace FileSys