Skip to content

Commit

Permalink
Optimized AES-CBC decryption
Browse files Browse the repository at this point in the history 
Since CBC decryption can be parallelize, instead of decrypting a
single file for each thread, with this update multiple threads
will now instead decrypt a single file, this is much more effecient
since all threads will work until all files are decrypted in the queue.
  • Loading branch information
@mrdcvlsc
mrdcvlsc committed Jul 17, 2023
1 parent 359cb0e commit ede01ef
Showing 1 changed file with 107 additions and 114 deletions.
221 changes: 107 additions & 114 deletions main.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,13 @@ std::mutex vector_mtx;
std::mutex output_mtx;

int main(int argc, char *args[]) {
const size_t processor_count = std::thread::hardware_concurrency();
size_t calculate_necessary_processors = std::thread::hardware_concurrency();

while (BUFFER_BYTESIZE % calculate_necessary_processors) {
calculate_necessary_processors--;
}

const size_t processor_count = calculate_necessary_processors;

std::cout << "\n";

Expand Down Expand Up @@ -432,160 +438,147 @@ int main(int argc, char *args[]) {

std::cout << "Decryption : AES block cipher \n";

auto decrypt_lambda = [&]() {
output_mtx.lock();
std::cout << "Decryption threads running : tid = " << std::this_thread::get_id() << "\n";
output_mtx.unlock();
// ========================= Single File - Parallel Decryption =========================

bool run_thread = true;
char *read_buffer = new char[BUFFER_BYTESIZE];
char filesig[bconst::FileSignature.size()];

char *read_buffer = new char[BUFFER_BYTESIZE];
char filesig[bconst::FileSignature.size()];
while (!file_queue.empty()) {
std::string target_file = file_queue.back();
file_queue.pop_back();

while (run_thread) {
std::string target_file;
std::string outfname(target_file);
outfname = outfname.substr(0, outfname.size() - bconst::extension.size());

vector_mtx.lock();
run_thread = !file_queue.empty();

if (run_thread) {
target_file = file_queue.back();
file_queue.pop_back();
} else {
vector_mtx.unlock();
break;
}

vector_mtx.unlock();
std::ifstream curr_file(target_file, std::ios::binary);

std::string outfname(target_file);
outfname = outfname.substr(0, outfname.size() - bconst::extension.size());
if (!curr_file.good()) {
output_mtx.lock();
std::cout << "The file : " << target_file
<< "\n"
"was not decrypted...\n"
"it might be read protected, corrupted or non-existent...\n";
output_mtx.unlock();
} else {
curr_file.read(filesig, bconst::FILESIGNATURE.size());

std::ifstream curr_file(target_file, std::ios::binary);

if (!curr_file.good()) {
if (std::memcmp(filesig, bconst::FILESIGNATURE.data(), bconst::FILESIGNATURE.size())) {
output_mtx.lock();
std::cout << "The file : " << target_file
<< "\n"
"was not decrypted...\n"
"it might be read protected, corrupted or non-existent...\n";
output_mtx.unlock();
} else {

output_mtx.lock();
std::cout << "decrypting : " << target_file << "...\n";
std::cerr << "The file '" << target_file << "' is not encrypted, no need to decrypt it!\n";
output_mtx.unlock();
continue;
}

curr_file.read(filesig, bconst::FILESIGNATURE.size());
std::ofstream output_file(outfname, std::ios::binary | std::ios::trunc);
output_file.close();
output_file.open(outfname, std::ios::binary | std::ios::app);

if (std::memcmp(filesig, bconst::FILESIGNATURE.data(), bconst::FILESIGNATURE.size())) {
output_mtx.lock();
std::cerr << "The file '" << target_file << "' is not encrypted, no need to decrypt it!\n";
output_mtx.unlock();
continue;
}
unsigned char iv[AES_BLOCKSIZE];
curr_file.read(reinterpret_cast<char *>(iv), AES_BLOCKSIZE);

std::ofstream output_file(outfname, std::ios::binary | std::ios::trunc);
output_file.close();
output_file.open(outfname, std::ios::binary | std::ios::app);
// diveded encryption

unsigned char iv[AES_BLOCKSIZE];
curr_file.read(reinterpret_cast<char *>(iv), AES_BLOCKSIZE);
constexpr size_t BUFFER_SIZE_NOLAST = BUFFER_BYTESIZE - AES_BLOCKSIZE;
char last_block[AES_BLOCKSIZE] = {};

// new decryption
bool lastblock_remains = false;

constexpr size_t BUFFER_SIZE_NOLAST = BUFFER_BYTESIZE - AES_BLOCKSIZE;
char last_block[AES_BLOCKSIZE] = {};
while (!curr_file.eof()) {

bool lastblock_remains = false;
curr_file.read(read_buffer, BUFFER_BYTESIZE);
size_t read_buffer_size = curr_file.gcount();

while (!curr_file.eof()) {
curr_file.read(read_buffer, BUFFER_BYTESIZE);
size_t read_buffer_size = curr_file.gcount();
if (lastblock_remains) {
if (read_buffer_size) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(last_block), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);

output_file.write(reinterpret_cast<char *>(last_block), AES_BLOCKSIZE);
lastblock_remains = false;
} else {
break;
}
}

if (lastblock_remains) {
if (read_buffer_size) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(last_block), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);
if (read_buffer_size == BUFFER_BYTESIZE) {

output_file.write(reinterpret_cast<char *>(last_block), AES_BLOCKSIZE);
lastblock_remains = false;
} else {
break;
}
}
size_t tread_section_size = BUFFER_BYTESIZE / processor_count;
std::vector<std::thread> threads;

if (read_buffer_size == BUFFER_BYTESIZE) {
size_t index;
for (index = 0; index < BUFFER_SIZE_NOLAST; index += AES_BLOCKSIZE) {
auto decrypt_lambda = [&](size_t start_index) {
size_t section_limit = start_index + tread_section_size;
for (size_t i = start_index; i < section_limit; i += AES_BLOCKSIZE) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(read_buffer + index), iv,
reinterpret_cast<unsigned char *>(read_buffer + i), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);
}
};

output_file.write(reinterpret_cast<char *>(read_buffer), BUFFER_SIZE_NOLAST);
std::memcpy(last_block, read_buffer + index, AES_BLOCKSIZE);
} else if (read_buffer_size % AES_BLOCKSIZE == 0) {
size_t index;
for (index = 0; index < read_buffer_size - AES_BLOCKSIZE; index += AES_BLOCKSIZE) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(read_buffer + index), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);
}
for (size_t i = 0; i < processor_count - 1; ++i) {
vector_mtx.lock();
threads.push_back(std::thread(decrypt_lambda, i * tread_section_size));
vector_mtx.unlock();
}

output_file.write(reinterpret_cast<char *>(read_buffer), read_buffer_size - AES_BLOCKSIZE);
std::memcpy(last_block, read_buffer + index, AES_BLOCKSIZE);
size_t index;
for (index = (processor_count - 1) * tread_section_size; index < BUFFER_SIZE_NOLAST;
index += AES_BLOCKSIZE) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(read_buffer + index), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);
}

lastblock_remains = true;
}
for (size_t i = 0; i < processor_count - 1; ++i) {
threads[i].join();
}

if (lastblock_remains) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(last_block), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);
output_file.write(reinterpret_cast<char *>(read_buffer), BUFFER_SIZE_NOLAST);
std::memcpy(last_block, read_buffer + index, AES_BLOCKSIZE);

} else if (read_buffer_size % AES_BLOCKSIZE == 0) {

if (last_block[AES_BLOCKSIZE - 1] < 0x10) {
output_file.write(reinterpret_cast<char *>(last_block), AES_BLOCKSIZE - last_block[AES_BLOCKSIZE - 1]);
size_t index;
for (index = 0; index < read_buffer_size - AES_BLOCKSIZE; index += AES_BLOCKSIZE) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(read_buffer + index), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);
}

lastblock_remains = false;
output_file.write(reinterpret_cast<char *>(read_buffer), read_buffer_size - AES_BLOCKSIZE);
std::memcpy(last_block, read_buffer + index, AES_BLOCKSIZE);
}

// new decryption

cnt++;

std::memset((unsigned char *) iv, 0x00, AES_BLOCKSIZE);
checkif_replace(args[COMMAND], target_file);
lastblock_remains = true;
}
}

std::memset((char *) read_buffer, 0x00, BUFFER_BYTESIZE);
std::memset((char *) filesig, 0x00, bconst::FILESIGNATURE.size());
if (lastblock_remains) {
Mode::CBC<AES_BLOCKSIZE>::decrypt(
reinterpret_cast<unsigned char *>(last_block), iv,
[&aes_cipher](unsigned char *block) { aes_cipher.decrypt_block(block); }
);

delete[] read_buffer;
};

std::vector<std::thread> threads;
if (last_block[AES_BLOCKSIZE - 1] < 0x10) {
output_file.write(reinterpret_cast<char *>(last_block), AES_BLOCKSIZE - last_block[AES_BLOCKSIZE - 1]);
}

for (size_t i = 0; i < processor_count - 1; ++i) {
vector_mtx.lock();
threads.push_back(std::thread(decrypt_lambda));
vector_mtx.unlock();
}
lastblock_remains = false;
}

decrypt_lambda();
cnt++;

for (size_t i = 0; i < processor_count - 1; ++i) {
threads[i].join();
std::memset((unsigned char *) iv, 0x00, AES_BLOCKSIZE);
checkif_replace(args[COMMAND], target_file);
}
}

// =====================================================================================
delete[] read_buffer;

timing_end("De", start, cnt);
} else if (!strcmp(args[COMMAND], GENERATE_FLAG)) {
if (argc == 2) {
Expand Down

0 comments on commit ede01ef

Please sign in to comment.