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.
This commit is contained in:
		| @@ -2,6 +2,7 @@ | ||||
| // Licensed under GPLv2 or any later version | ||||
| // Refer to the license.txt file included. | ||||
|  | ||||
| #include <array> | ||||
| #include <mbedtls/cipher.h> | ||||
| #include "common/assert.h" | ||||
| #include "common/logging/log.h" | ||||
| @@ -10,8 +11,10 @@ | ||||
|  | ||||
| namespace Core::Crypto { | ||||
| namespace { | ||||
| std::vector<u8> CalculateNintendoTweak(std::size_t sector_id) { | ||||
|     std::vector<u8> out(0x10); | ||||
| using NintendoTweak = std::array<u8, 16>; | ||||
|  | ||||
| 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<Key, KeySize>::~AESCipher() { | ||||
|     mbedtls_cipher_free(&ctx->decryption_context); | ||||
| } | ||||
|  | ||||
| template <typename Key, std::size_t KeySize> | ||||
| void AESCipher<Key, KeySize>::SetIV(std::vector<u8> 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 <typename Key, std::size_t KeySize> | ||||
| void AESCipher<Key, KeySize>::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<Key, KeySize>::XTSTranscode(const u8* src, std::size_t size, u8* | ||||
|     } | ||||
| } | ||||
|  | ||||
| template <typename Key, std::size_t KeySize> | ||||
| void AESCipher<Key, KeySize>::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<Key128>; | ||||
| template class AESCipher<Key256>; | ||||
| } // namespace Core::Crypto | ||||
|   | ||||
| @@ -6,7 +6,6 @@ | ||||
|  | ||||
| #include <memory> | ||||
| #include <type_traits> | ||||
| #include <vector> | ||||
| #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<u8> iv); | ||||
|     template <typename ContiguousContainer> | ||||
|     void SetIV(const ContiguousContainer& container) { | ||||
|         SetIVImpl(std::data(container), std::size(container)); | ||||
|     } | ||||
|  | ||||
|     template <typename Source, typename Dest> | ||||
|     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<CipherContext> ctx; | ||||
| }; | ||||
| } // namespace Core::Crypto | ||||
|   | ||||
| @@ -2,6 +2,7 @@ | ||||
| // Licensed under GPLv2 or any later version | ||||
| // Refer to the license.txt file included. | ||||
|  | ||||
| #include <algorithm> | ||||
| #include <cstring> | ||||
| #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<u8>& 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 { | ||||
|   | ||||
| @@ -4,7 +4,8 @@ | ||||
|  | ||||
| #pragma once | ||||
|  | ||||
| #include <vector> | ||||
| #include <array> | ||||
|  | ||||
| #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<u8, 16>; | ||||
|  | ||||
|     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<u8>& iv); | ||||
|     void SetIV(const IVData& iv); | ||||
|  | ||||
| private: | ||||
|     std::size_t base_offset; | ||||
|  | ||||
|     // Must be mutable as operations modify cipher contexts. | ||||
|     mutable AESCipher<Key128> cipher; | ||||
|     mutable std::vector<u8> iv; | ||||
|     mutable IVData iv{}; | ||||
|  | ||||
|     void UpdateIV(std::size_t offset) const; | ||||
| }; | ||||
|   | ||||
| @@ -346,10 +346,9 @@ FileSys::VirtualFile PartitionDataManager::GetPackage2Raw(Package2Type type) con | ||||
| } | ||||
|  | ||||
| static bool AttemptDecrypt(const std::array<u8, 16>& key, Package2Header& header) { | ||||
|     const std::vector<u8> iv(header.header_ctr.begin(), header.header_ctr.end()); | ||||
|     Package2Header temp = header; | ||||
|     AESCipher<Key128> 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<Key128, 0x20>& packa | ||||
|     auto c = a->ReadAllBytes(); | ||||
|  | ||||
|     AESCipher<Key128> 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<FileSys::VectorVfsFile>(c); | ||||
|   | ||||
| @@ -495,9 +495,10 @@ VirtualFile NCA::Decrypt(const NCASectionHeader& s_header, VirtualFile in, u64 s | ||||
|  | ||||
|             auto out = std::make_shared<Core::Crypto::CTREncryptionLayer>(std::move(in), *key, | ||||
|                                                                           starting_offset); | ||||
|             std::vector<u8> 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<VfsFile>(out); | ||||
|         } | ||||
|   | ||||
| @@ -3,6 +3,7 @@ | ||||
| // Refer to the license.txt file included. | ||||
|  | ||||
| #include <algorithm> | ||||
| #include <array> | ||||
| #include <cstddef> | ||||
| #include <cstring> | ||||
|  | ||||
| @@ -66,7 +67,7 @@ std::size_t BKTR::Read(u8* data, std::size_t length, std::size_t offset) const { | ||||
|     Core::Crypto::AESCipher<Core::Crypto::Key128> cipher(key, Core::Crypto::Mode::CTR); | ||||
|  | ||||
|     // Calculate AES IV | ||||
|     std::vector<u8> iv(16); | ||||
|     std::array<u8, 16> iv{}; | ||||
|     auto subsection_ctr = subsection.ctr; | ||||
|     auto offset_iv = section_offset + base_offset; | ||||
|     for (std::size_t i = 0; i < section_ctr.size(); ++i) | ||||
|   | ||||
		Reference in New Issue
	
	Block a user
	 Lioncash
					Lioncash