From d62ca85c5b90793e50110f5799b36f039b92928f Mon Sep 17 00:00:00 2001 From: Michele Rodolfi Date: Mon, 14 Dec 2020 16:50:56 +0100 Subject: [PATCH] Optimize memory usage and clean interface --- src/crypto.cpp | 56 ++++------- src/crypto.h | 193 +++++++++++++++++++++----------------- src/horcrux.cpp | 36 ++++--- src/horcrux.h | 4 +- src/io.cpp | 10 +- src/io.h | 2 + src/main.cpp | 5 +- test/integration-test.cpp | 81 +++++++++++----- 8 files changed, 223 insertions(+), 164 deletions(-) diff --git a/src/crypto.cpp b/src/crypto.cpp index 0f0e078..6b3a442 100644 --- a/src/crypto.cpp +++ b/src/crypto.cpp @@ -13,8 +13,10 @@ namespace horcrux { -void AES256_CBC::init(Mode mode, const raw_data& key, +void AES256_CBC::init(Mode new_mode, const raw_data& key, const raw_data& iv) { + // force mode change + mode = new_mode; if (key.size() != kKeySize) { throw std::runtime_error("Wrong key size"); } @@ -46,7 +48,6 @@ size_t AES256_CBC::process_start(Mode mode, auto current = begin; size_t output_offset = 0; raw_data iv(kIvSize); - if (mode == Mode::kEncrypt) { // Make sure output is large enough to add IV output.resize(kIvSize); @@ -122,6 +123,7 @@ size_t AES256_CBC::process_final(Mode mode, return len; } +/* this function does not use process_start, hence it does not handle the IV */ size_t AES256_CBC::process_all(Mode mode, raw_data::const_iterator begin, raw_data::const_iterator end, @@ -134,54 +136,32 @@ size_t AES256_CBC::process_all(Mode mode, return len; } -raw_data AES256_CBC::process(const raw_data& inputdata, bool begin, bool end) { - //if (mode == Mode::kEncrypt) - // return encrypt(inputdata, begin, end); - //else - // return decrypt(inputdata, begin, end); +raw_data AES256_CBC::process(const raw_data& inputdata, const Flags& flags) { raw_data output; size_t len; - if (begin) + if (flags.to_ulong() & flag::kBegin) len = process_start(mode, inputdata.begin(), inputdata.end(), output); else len = process_chunk(mode, inputdata.begin(), inputdata.end(), output, 0); - if (end) + if (flags.to_ulong() & flag::kEnd) process_final(mode, output, len); return output; } -raw_data AES256_CBC::encrypt(const raw_data& plaintext, bool begin, bool end) { - //return encrypt(encryption_key, horcrux::generate_random(kIvSize), - // plaintext); - raw_data output; - size_t len; - if (begin) - len = process_start(Mode::kEncrypt, plaintext.begin(), plaintext.end(), output); - else - len = process_chunk(Mode::kEncrypt, plaintext.begin(), plaintext.end(), output, 0); - if (end) - process_final(Mode::kEncrypt, output, len); - return output; +raw_data AES256_CBC::encrypt(const raw_data& plaintext, + const std::bitset& flags) { + mode = Mode::kEncrypt; + return process(plaintext, flags); } -raw_data AES256_CBC::decrypt(const raw_data& ciphertext, bool begin, bool end) { - //return decrypt(encryption_key, ciphertext); - //raw_data output; - //init(Mode::kDecrypt, key, iv); - //process_all(Mode::kDecrypt, begin, end, output, 0); - //auto len = process_start(Mode::kDecrypt, ciphertext.begin(), ciphertext.end(), output); - //process_final(Mode::kDecrypt, output, len); - //return output; - raw_data output; - size_t len; - if (begin) - len = process_start(Mode::kDecrypt, ciphertext.begin(), ciphertext.end(), output); - else - len = process_chunk(Mode::kDecrypt, ciphertext.begin(), ciphertext.end(), output, 0); - if (end) - process_final(Mode::kDecrypt, output, len); - return output; + +raw_data AES256_CBC::decrypt(const raw_data& ciphertext, + const std::bitset& flags) { + mode = Mode::kDecrypt; + return process(ciphertext, flags); } +/* OVERLOADED FUNCTIONS FOR TEST PURPOSE */ + raw_data AES256_CBC::encrypt(const raw_data& key, const raw_data& iv, const raw_data& input) { init(Mode::kEncrypt, key, iv); diff --git a/src/crypto.h b/src/crypto.h index 547953c..5b6d147 100644 --- a/src/crypto.h +++ b/src/crypto.h @@ -4,6 +4,7 @@ #include #include #include +#include #include "utils.h" @@ -16,6 +17,10 @@ enum class Mode { kEncrypt, kDecrypt }; */ class Cipher { public: + enum flag { kBegin = 0x01, kEnd = 0x02}; + constexpr static int Flags_num = 2; + typedef std::bitset Flags; + /** @brief default constructor */ Cipher() = delete; @@ -35,10 +40,14 @@ public: protected: /** @brief key used for encryption / decryption */ raw_data encryption_key; + /** @brief current cipher mode */ + Mode mode; public: - /** @brief Cipher mode (encrypt or decrypt) */ - const Mode mode; + /** @brief get current Cipher mode (encrypt or decrypt) */ + const Mode& get_mode() { + return mode; + } /** @brief get Cipher encryption key * @return the encryption key @@ -48,31 +57,37 @@ public: } /** @brief process inputdata according to Cipher mode + * @details + * set kBegin flag to specify this is the first chunk to process. + * set kEnd flag to specify this is the last chunk to process. * @param inputdata buffer containing data to process - * @param begin set to true if this is the first chunk of data to process. - * @param end set to true if this is the last chunk of data to process. + * @param flags flag bitset. Default (kBegin | kEnd) * @return processed data */ virtual raw_data process( - const raw_data& inputdata, bool begin = true, bool end = true) = 0; + const raw_data& inputdata, const std::bitset& flags = Flags(kBegin | kEnd)) = 0; /** @brief encrypt the content of a buffer using Cipher key + * @details + * set kBegin flag to specify this is the first chunk to process. + * set kEnd flag to specify this is the last chunk to process. * @param plaintext the input buffer - * @param begin set to true if this is the first chunk of data to process. - * @param end set to true if this is the last chunk of data to process. + * @param flags flag bitset. Default (kBegin | kEnd) * @return encrypted data */ virtual raw_data encrypt( - const raw_data& plaintext, bool begin = true, bool end = true) = 0; + const raw_data& plaintext, const std::bitset& flags = Flags(kBegin | kEnd)) = 0; /** @brief decrypt the content of a buffer using Cipher key + * @details + * set kBegin flag to specify this is the first chunk to process. + * set kEnd flag to specify this is the last chunk to process. * @param decrypt the input buffer - * @param begin set to true if this is the first chunk of data to process. - * @param end set to true if this is the last chunk of data to process. + * @param flags flag bitset. Default (kBegin | kEnd) * @return decrypted data */ virtual raw_data decrypt( - const raw_data& ciphertext, bool begin = true, bool end = true) = 0; + const raw_data& ciphertext, const std::bitset& flags = Flags(kBegin | kEnd)) = 0; }; /** Cipher context class @@ -113,7 +128,7 @@ class AES256_CBC : public Cipher { const size_t kKeySize = 32; /** AES256 CBC initialization vector size, from AES256 specs */ const size_t kIvSize = 16; - + /** pointer to Cipher context */ std::unique_ptr ctx = nullptr; public: @@ -133,100 +148,62 @@ public: virtual ~AES256_CBC() = default; /** @brief process inputdata according to Cipher mode + * @details + * if kBegin flag is set, this function will handle the initial IV. + * if kEnd flag is set, this function will handle the final padding. * @param inputdata buffer containing data to process - * @param begin set to true if this is the first chunk of data to process. - * (deals with IV) - * @param end set to true if this is the last chunk of data to process. - * (deals with padding) + * @param flags flag bitset. Default (kBegin | kEnd) * @return processed data */ raw_data process( - const raw_data& inputdata, bool begin = true, bool end = true) override; + const raw_data& inputdata, const std::bitset& flags = Flags(kBegin | kEnd)) override; /** @brief encrypt the content of a buffer using Cipher key + * @details + * if kBegin flag is set, this function will handle the initial IV. + * if kEnd flag is set, this function will handle the final padding. * @param plaintext the input buffer - * @param begin set to true if this is the first chunk of data to process. - * (deals with IV) - * @param end set to true if this is the last chunk of data to process. - * (deals with padding) + * @param flags flag bitset. Default (kBegin | kEnd) * @return encrypted data */ raw_data encrypt( - const raw_data& plaintext, bool begin = true, bool end = true) override; + const raw_data& plaintext, const std::bitset& flags = Flags(kBegin | kEnd)) override; /** @brief decrypt the content of a buffer using Cipher key + * @details + * if kBegin flag is set, this function will handle the initial IV. + * if kEnd flag is set, this function will handle the final padding. * @param decrypt the input buffer - * @param begin set to true if this is the first chunk of data to process. - * (deals with IV) - * @param end set to true if this is the last chunk of data to process. - * (deals with padding) + * @param flags flag bitset. Default (kBegin | kEnd) * @return decrypted data */ raw_data decrypt( - const raw_data& ciphertext, bool begin = true, bool end = true) override; - - /** @brief Read the whole input and return a buffer with the encrypted data - * @param key the encryption key - * @param iv the initialization vector that will be copied to the beginning - * of the output - * @param input the input buffer - * @return the encrypted data buffer - */ - raw_data encrypt(const raw_data& key, - const raw_data& iv, - const raw_data& input); - -/* - int encrypt(const raw_data& key, - const raw_data& iv, - std::istream input, const size_t input_len, - std::ostream output, size_t& output_len); - -*/ - /** @brief Read the whole input and return a buffer with the decrypted data. - * Read the iv from the beginning of ciphertext - * @param key the encryption key - * @param ciphertext the input buffer - * @return the encrypted data buffer - */ - raw_data decrypt( - const raw_data& key, - const raw_data& ciphertext); - - /** @brief Read the whole input and return a buffer with the decrypted data. - * ciphertext does not contain the iv. - * @param key the encryption key - * @param iv the initialization vector - * @param ciphertext the input buffer - * @return the encrypted data buffer - */ - raw_data decrypt( - const raw_data& key, - const raw_data& iv, - const raw_data& ciphertext); - - /** @brief Read some input from begin to end and return a buffer with the - * decrypted data. Output data is finalized. - * @param key the encryption key - * @param iv the initialization vector - * @param begin beginning of the input buffer - * @param end end of the input buffer - * @return the encrypted data buffer - */ - raw_data decrypt( - const raw_data& key, - const raw_data& iv, - raw_data::const_iterator begin, - raw_data::const_iterator end); + const raw_data& ciphertext, const std::bitset& flags = Flags(kBegin | kEnd)) override; private: /** @brief Initialize the Cipher context object needed for enc/dec operations + * @details this operation will reset the cipher, any pending operation will + * be lost. * @param mode encrypt or decrypt * @param key the enc/dec key * @param the initialization vector */ - void init(Mode mode, const raw_data& key, - const raw_data& iv); + void init(Mode mode, const raw_data& key, const raw_data& iv); + + /** @brief process the first chunk of input, it handles the IV. + * @param mode encrypt or decrypt + * @param begin beginning of the input chunk + * @param end end of the input chunk + * @param output output buffer, it will be overwritten from the beginning + * @param resize_in resize the output buffer before appending data to it in + * order to make sure there is enough room. + * @param resize_out shrink the output buffer at the end + */ + size_t process_start(Mode mode, + raw_data::const_iterator begin, + raw_data::const_iterator end, + raw_data& output, + bool resize_in = true, bool resize_out = true); /** @brief process a chunk of input * @param mode encrypt or decrypt @@ -245,11 +222,6 @@ private: raw_data& output, size_t output_offset, bool resize_in = true, bool resize_out = true); - size_t process_start(Mode mode, - raw_data::const_iterator begin, - raw_data::const_iterator end, - raw_data& output, - bool resize_in = true, bool resize_out = true); /** @brief Finalize a encryption/decryption process. Write remaining bytes * to the output @@ -267,6 +239,7 @@ private: bool resize_in = true, bool resize_out = true); /** @brief process a input and finalize it. + * @node this function does not handle the IV * @param mode encrypt or decrypt * @param begin beginning of the input chunk * @param end end of the input chunk @@ -282,6 +255,50 @@ private: raw_data::const_iterator end, raw_data& output, size_t output_offset, bool resize_in = true, bool resize_out = true); + +public: + /* OVERLOADED FUNCTIONS FOR TEST PURPOSE */ + + /** @brief Read the whole input and return a buffer with the encrypted data + * @param key the encryption key + * @param iv the initialization vector that will be copied to the beginning + * of the output + * @param input the input buffer + * @return the encrypted data buffer + */ + raw_data encrypt(const raw_data& key, const raw_data& iv, + const raw_data& input); + + /** @brief Read the whole input and return a buffer with the decrypted data. + * Read the iv from the beginning of ciphertext + * @param key the encryption key + * @param ciphertext the input buffer + * @return the encrypted data buffer + */ + raw_data decrypt(const raw_data& key, const raw_data& ciphertext); + + /** @brief Read the whole input and return a buffer with the decrypted data. + * ciphertext does not contain the iv. + * @param key the encryption key + * @param iv the initialization vector + * @param ciphertext the input buffer + * @return the encrypted data buffer + */ + raw_data decrypt(const raw_data& key, const raw_data& iv, + const raw_data& ciphertext); + + /** @brief Read some input from begin to end and return a buffer with the + * decrypted data. Output data is finalized. + * @param key the encryption key + * @param iv the initialization vector + * @param begin beginning of the input buffer + * @param end end of the input buffer + * @return the encrypted data buffer + */ + raw_data decrypt(const raw_data& key, const raw_data& iv, + raw_data::const_iterator begin, + raw_data::const_iterator end); + }; }; // namespace horcrux diff --git a/src/horcrux.cpp b/src/horcrux.cpp index ad53ff6..905924c 100644 --- a/src/horcrux.cpp +++ b/src/horcrux.cpp @@ -5,25 +5,28 @@ namespace horcrux { -void Horcrux::read_chunks(size_t size) { +void Horcrux::process_chunks(size_t size) { raw_data buf = input->read(size); - output->write_chunk(cipher->process(buf, true, false)); + //std::cout << "read " << buf.size() << " " << size << std::endl; + output->write_chunk(cipher->process(buf, Cipher::Flags(Cipher::flag::kBegin))); while ((buf = input->read(size)).size() == size) { - output->write_chunk(cipher->process(buf, false, false)); + //std::cout << "read " << buf.size() << std::endl; + output->write_chunk(cipher->process(buf, Cipher::Flags())); } - output->write_chunk(cipher->process(buf, false, true)); + output->write_chunk(cipher->process(buf, Cipher::Flags(Cipher::flag::kEnd))); } void Horcrux::init() { if (options.count <= 0) { - throw std::invalid_argument("Invalid horcrux count"); + throw std::invalid_argument("init: Invalid horcrux count"); } switch (options.mode) { case Mode::kEncrypt: cipher = std::make_unique(); input = std::make_unique(options.input[0]); output = std::make_unique(options.output, options.count, - std::filesystem::path(options.input[0]).filename().string()); + std::filesystem::path(options.input[0]).filename().string(), + static_cast(input.get())->get_file_size() / options.count); options.base64_key = to_base64(cipher->get_encryption_key()); break; case Mode::kDecrypt: @@ -39,12 +42,23 @@ void Horcrux::init() { } } -void Horcrux::run() { - if (options.mode == Mode::kEncrypt) +void Horcrux::run(size_t buffer_size) { + if (buffer_size == 0){ + // read the whole input and process it. It may use a lot of memory. + output->write(cipher->process(input->read())); + } else { + // read and process the input buffer_size bytes at a time + process_chunks(buffer_size); + } + /* + if (options.mode == Mode::kEncrypt){ //output->write(cipher->encrypt(input->read())); - read_chunks(512); - else - output->write(cipher->decrypt(input->read())); + read_chunks(buffer_size); + } else { + //output->write(cipher->decrypt(input->read())); + read_chunks(buffer_size); + } + */ } static void print_usage(const std::string& program) { diff --git a/src/horcrux.h b/src/horcrux.h index 89e2d6d..0316fc5 100644 --- a/src/horcrux.h +++ b/src/horcrux.h @@ -38,7 +38,7 @@ class Horcrux { * Used cipher, input and output are generic and the actual implementation * can be selected in this function */ void init(); - void read_chunks(size_t size); + void process_chunks(size_t size); public: /** remove default constructor */ @@ -53,7 +53,7 @@ public: /** @brief execute the horcrux process * @details read from Input, process with Cipher and write to output */ - void run(); + void run(size_t size = 0); }; /** @brief utility function to parse the command line arguments * @param arguments arguments as read from command line diff --git a/src/io.cpp b/src/io.cpp index d95c55f..c707039 100644 --- a/src/io.cpp +++ b/src/io.cpp @@ -25,9 +25,11 @@ raw_data FsPlainInput::read() { } raw_data FsPlainInput::read(size_t size) { raw_data result(size); - file_stream.open(file_path, std::ios::binary); + if (!file_stream.is_open()) + file_stream.open(file_path, std::ios::binary); file_stream.read(reinterpret_cast(result.data()), size); result.resize(file_stream.gcount()); + //std::cout << size << " read " << file_stream.gcount() << std::endl; return result; } @@ -77,6 +79,8 @@ raw_data FsCryptoInput::read(size_t size) { ++current_file; } } + result.resize(data_read); + //std::cout << "data_read: " << data_read << std::endl; return result; } @@ -94,7 +98,8 @@ size_t FsPlainOutput::write(const raw_data& to_write) { return write_chunk(to_write); } size_t FsPlainOutput::write_chunk(const raw_data& to_write) { - file_stream.open(file_path, std::ios::binary); + if(!file_stream.is_open()) + file_stream.open(file_path, std::ios::binary); file_stream.write(reinterpret_cast(to_write.data()), to_write.size()); return to_write.size(); @@ -138,6 +143,7 @@ size_t FsCryptoOutput::write_chunk(const raw_data& to_write) { if (size <= 0) { throw std::invalid_argument("Invalid horcrux size"); } + //std::cout << "writing " << to_write.size() << std::endl; size_t data_written{0}; while (data_written < to_write.size()){ if(!file_stream.is_open()){ diff --git a/src/io.h b/src/io.h index 49054fb..5c7b544 100644 --- a/src/io.h +++ b/src/io.h @@ -52,6 +52,8 @@ public: * @param path input file path, must be a regular file */ explicit FsPlainInput(const std::string& path); + const size_t& get_file_size() { return file_size; }; + /** @brief Read from the input * @return new buffer containing all read data */ raw_data read() override; diff --git a/src/main.cpp b/src/main.cpp index d0287fb..602987e 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -2,6 +2,9 @@ #include #include "horcrux.h" +/* arbitrary buffer size */ +constexpr size_t buffer_size = 1024; + int main(const int argc, const char *argv[]) { int ret = -1; std::vector arguments; @@ -11,7 +14,7 @@ int main(const int argc, const char *argv[]) { try { horcrux::Horcrux h(horcrux::parse_arguments(std::move(arguments))); - h.run(); + h.run(buffer_size); if (h.mode() == horcrux::Mode::kEncrypt) { // print encryption key to stdout std::cout << h.key() << std::endl; diff --git a/test/integration-test.cpp b/test/integration-test.cpp index 18517b3..2c69630 100644 --- a/test/integration-test.cpp +++ b/test/integration-test.cpp @@ -62,35 +62,72 @@ TEST(IntegrationTest, HorcruxOptions){ options.base64_key = to_base64(generate_random(32)); EXPECT_NO_THROW(Horcrux {options}); } -TEST(IntegrationTest, HorcruxEncryptiDecrypt) { - HorcruxOptions options { - .mode = Mode::kEncrypt, - .count = 3, - .input = {image}, - .output = {folder} - }; - Horcrux enc{options}; +TEST(IntegrationTest, HorcruxEncryptDecrypt) { + { + HorcruxOptions options { + .mode = Mode::kEncrypt, + .count = 3, + .input = {image}, + .output = {folder} + }; + Horcrux enc{options}; + enc.run(); + auto key = enc.key(); - enc.run(); - auto key = enc.key(); + HorcruxOptions options2 { + .mode = Mode::kDecrypt, + .count = 3, + .input = {get_encrypted_files(folder, image)}, + .output = {noexist}, + .base64_key = key + }; - HorcruxOptions options2 { - .mode = Mode::kDecrypt, - .count = 3, - .input = {get_encrypted_files(folder, image)}, - .output = {noexist}, - .base64_key = key - }; - - Horcrux dec{options2}; - dec.run(); + Horcrux dec{options2}; + dec.run(); + std::for_each(options2.input.begin(), options2.input.end(), [](auto& p){ + std::filesystem::remove(p);}); + } // Close scope so that Output objects get flushed //Compare input and ouput files + EXPECT_EQ(std::filesystem::file_size(image), std::filesystem::file_size(noexist)); EXPECT_EQ(generic_read_file(image), generic_read_file(noexist)); std::filesystem::remove(noexist); - std::for_each(options2.input.begin(), options2.input.end(), [](auto& p){ - std::filesystem::remove(p);}); } + +TEST(IntegrationTest, HorcruxEncryptDecryptBuffer) { + { + HorcruxOptions options { + .mode = Mode::kEncrypt, + .count = 3, + .input = {image}, + .output = {folder} + }; + Horcrux enc{options}; + // read and process 1024 bytes at a time + enc.run(1024); + auto key = enc.key(); + + HorcruxOptions options2 { + .mode = Mode::kDecrypt, + .count = 3, + .input = {get_encrypted_files(folder, image)}, + .output = {noexist}, + .base64_key = key + }; + + Horcrux dec{options2}; + // read and process 2048 bytes at a time + dec.run(2048); + std::for_each(options2.input.begin(), options2.input.end(), [](auto& p){ + std::filesystem::remove(p);}); + } // Close scope so that Output objects get flushed + + //Compare input and ouput files + EXPECT_EQ(std::filesystem::file_size(image), std::filesystem::file_size(noexist)); + EXPECT_EQ(generic_read_file(image), generic_read_file(noexist)); + std::filesystem::remove(noexist); +} + TEST(IntegrationTest, arguments){ EXPECT_THROW( parse_arguments({"horcrux", "create", "-n", "4", "input_file"}),