From 15660bd8570735139d91d0165a2614747f570202 Mon Sep 17 00:00:00 2001 From: Lioncash Date: Mon, 3 Aug 2020 14:14:39 -0400 Subject: [PATCH] aes_util: Allow SetIV to be non-allocating In a few places, the data to be set as the IV is already within an array. We shouldn't require this data to be heap-allocated if it doesn't need to be. This allows certain callers to reduce heap churn. --- src/core/crypto/aes_util.cpp | 21 ++++++++++++--------- src/core/crypto/aes_util.h | 9 ++++++--- src/core/crypto/ctr_encryption_layer.cpp | 9 ++++----- src/core/crypto/ctr_encryption_layer.h | 9 ++++++--- src/core/crypto/partition_data_manager.cpp | 5 ++--- src/core/file_sys/content_archive.cpp | 7 ++++--- src/core/file_sys/nca_patch.cpp | 3 ++- 7 files changed, 36 insertions(+), 27 deletions(-) diff --git a/src/core/crypto/aes_util.cpp b/src/core/crypto/aes_util.cpp index 4be76bb43..330996b24 100644 --- a/src/core/crypto/aes_util.cpp +++ b/src/core/crypto/aes_util.cpp @@ -2,6 +2,7 @@ // Licensed under GPLv2 or any later version // Refer to the license.txt file included. +#include #include #include "common/assert.h" #include "common/logging/log.h" @@ -10,8 +11,10 @@ namespace Core::Crypto { namespace { -std::vector CalculateNintendoTweak(std::size_t sector_id) { - std::vector out(0x10); +using NintendoTweak = std::array; + +NintendoTweak CalculateNintendoTweak(std::size_t sector_id) { + NintendoTweak out{}; for (std::size_t i = 0xF; i <= 0xF; --i) { out[i] = sector_id & 0xFF; sector_id >>= 8; @@ -63,13 +66,6 @@ AESCipher::~AESCipher() { mbedtls_cipher_free(&ctx->decryption_context); } -template -void AESCipher::SetIV(std::vector iv) { - ASSERT_MSG((mbedtls_cipher_set_iv(&ctx->encryption_context, iv.data(), iv.size()) || - mbedtls_cipher_set_iv(&ctx->decryption_context, iv.data(), iv.size())) == 0, - "Failed to set IV on mbedtls ciphers."); -} - template void AESCipher::Transcode(const u8* src, std::size_t size, u8* dest, Op op) const { auto* const context = op == Op::Encrypt ? &ctx->encryption_context : &ctx->decryption_context; @@ -124,6 +120,13 @@ void AESCipher::XTSTranscode(const u8* src, std::size_t size, u8* } } +template +void AESCipher::SetIVImpl(const u8* data, std::size_t size) { + ASSERT_MSG((mbedtls_cipher_set_iv(&ctx->encryption_context, data, size) || + mbedtls_cipher_set_iv(&ctx->decryption_context, data, size)) == 0, + "Failed to set IV on mbedtls ciphers."); +} + template class AESCipher; template class AESCipher; } // namespace Core::Crypto diff --git a/src/core/crypto/aes_util.h b/src/core/crypto/aes_util.h index edc4ab910..e2a304186 100644 --- a/src/core/crypto/aes_util.h +++ b/src/core/crypto/aes_util.h @@ -6,7 +6,6 @@ #include #include -#include #include "common/common_types.h" #include "core/file_sys/vfs.h" @@ -32,10 +31,12 @@ class AESCipher { public: AESCipher(Key key, Mode mode); - ~AESCipher(); - void SetIV(std::vector iv); + template + void SetIV(const ContiguousContainer& container) { + SetIVImpl(std::data(container), std::size(container)); + } template void Transcode(const Source* src, std::size_t size, Dest* dest, Op op) const { @@ -59,6 +60,8 @@ public: std::size_t sector_size, Op op); private: + void SetIVImpl(const u8* data, std::size_t size); + std::unique_ptr ctx; }; } // namespace Core::Crypto diff --git a/src/core/crypto/ctr_encryption_layer.cpp b/src/core/crypto/ctr_encryption_layer.cpp index 902841c77..5c84bb0a4 100644 --- a/src/core/crypto/ctr_encryption_layer.cpp +++ b/src/core/crypto/ctr_encryption_layer.cpp @@ -2,6 +2,7 @@ // Licensed under GPLv2 or any later version // Refer to the license.txt file included. +#include #include #include "common/assert.h" #include "core/crypto/ctr_encryption_layer.h" @@ -10,8 +11,7 @@ namespace Core::Crypto { CTREncryptionLayer::CTREncryptionLayer(FileSys::VirtualFile base_, Key128 key_, std::size_t base_offset) - : EncryptionLayer(std::move(base_)), base_offset(base_offset), cipher(key_, Mode::CTR), - iv(16, 0) {} + : EncryptionLayer(std::move(base_)), base_offset(base_offset), cipher(key_, Mode::CTR) {} std::size_t CTREncryptionLayer::Read(u8* data, std::size_t length, std::size_t offset) const { if (length == 0) @@ -39,9 +39,8 @@ std::size_t CTREncryptionLayer::Read(u8* data, std::size_t length, std::size_t o return read + Read(data + read, length - read, offset + read); } -void CTREncryptionLayer::SetIV(const std::vector& iv_) { - const auto length = std::min(iv_.size(), iv.size()); - iv.assign(iv_.cbegin(), iv_.cbegin() + length); +void CTREncryptionLayer::SetIV(const IVData& iv_) { + iv = iv_; } void CTREncryptionLayer::UpdateIV(std::size_t offset) const { diff --git a/src/core/crypto/ctr_encryption_layer.h b/src/core/crypto/ctr_encryption_layer.h index a7bf810f4..a2429f001 100644 --- a/src/core/crypto/ctr_encryption_layer.h +++ b/src/core/crypto/ctr_encryption_layer.h @@ -4,7 +4,8 @@ #pragma once -#include +#include + #include "core/crypto/aes_util.h" #include "core/crypto/encryption_layer.h" #include "core/crypto/key_manager.h" @@ -14,18 +15,20 @@ namespace Core::Crypto { // Sits on top of a VirtualFile and provides CTR-mode AES decription. class CTREncryptionLayer : public EncryptionLayer { public: + using IVData = std::array; + CTREncryptionLayer(FileSys::VirtualFile base, Key128 key, std::size_t base_offset); std::size_t Read(u8* data, std::size_t length, std::size_t offset) const override; - void SetIV(const std::vector& iv); + void SetIV(const IVData& iv); private: std::size_t base_offset; // Must be mutable as operations modify cipher contexts. mutable AESCipher cipher; - mutable std::vector iv; + mutable IVData iv{}; void UpdateIV(std::size_t offset) const; }; diff --git a/src/core/crypto/partition_data_manager.cpp b/src/core/crypto/partition_data_manager.cpp index 7ed71ac3a..b31a81560 100644 --- a/src/core/crypto/partition_data_manager.cpp +++ b/src/core/crypto/partition_data_manager.cpp @@ -346,10 +346,9 @@ FileSys::VirtualFile PartitionDataManager::GetPackage2Raw(Package2Type type) con } static bool AttemptDecrypt(const std::array& key, Package2Header& header) { - const std::vector iv(header.header_ctr.begin(), header.header_ctr.end()); Package2Header temp = header; AESCipher cipher(key, Mode::CTR); - cipher.SetIV(iv); + cipher.SetIV(header.header_ctr); cipher.Transcode(&temp.header_ctr, sizeof(Package2Header) - 0x100, &temp.header_ctr, Op::Decrypt); if (temp.magic == Common::MakeMagic('P', 'K', '2', '1')) { @@ -388,7 +387,7 @@ void PartitionDataManager::DecryptPackage2(const std::array& packa auto c = a->ReadAllBytes(); AESCipher cipher(package2_keys[revision], Mode::CTR); - cipher.SetIV({header.section_ctr[1].begin(), header.section_ctr[1].end()}); + cipher.SetIV(header.section_ctr[1]); cipher.Transcode(c.data(), c.size(), c.data(), Op::Decrypt); const auto ini_file = std::make_shared(c); diff --git a/src/core/file_sys/content_archive.cpp b/src/core/file_sys/content_archive.cpp index 473245d5a..5039341c7 100644 --- a/src/core/file_sys/content_archive.cpp +++ b/src/core/file_sys/content_archive.cpp @@ -495,9 +495,10 @@ VirtualFile NCA::Decrypt(const NCASectionHeader& s_header, VirtualFile in, u64 s auto out = std::make_shared(std::move(in), *key, starting_offset); - std::vector iv(16); - for (u8 i = 0; i < 8; ++i) - iv[i] = s_header.raw.section_ctr[0x8 - i - 1]; + Core::Crypto::CTREncryptionLayer::IVData iv{}; + for (std::size_t i = 0; i < 8; ++i) { + iv[i] = s_header.raw.section_ctr[8 - i - 1]; + } out->SetIV(iv); return std::static_pointer_cast(out); } diff --git a/src/core/file_sys/nca_patch.cpp b/src/core/file_sys/nca_patch.cpp index 0090cc6c4..fe7375e84 100644 --- a/src/core/file_sys/nca_patch.cpp +++ b/src/core/file_sys/nca_patch.cpp @@ -3,6 +3,7 @@ // Refer to the license.txt file included. #include +#include #include #include @@ -66,7 +67,7 @@ std::size_t BKTR::Read(u8* data, std::size_t length, std::size_t offset) const { Core::Crypto::AESCipher cipher(key, Core::Crypto::Mode::CTR); // Calculate AES IV - std::vector iv(16); + std::array iv{}; auto subsection_ctr = subsection.ctr; auto offset_iv = section_offset + base_offset; for (std::size_t i = 0; i < section_ctr.size(); ++i)