Reputation: 13
I have written a code to encrypt according to AES 256 algorithm and it works to give me an output file but I realised that it hangs after running the EncryptDecryptString function and does not perform anything else after that. Why is this so? How can I further improve this code to perform encryption and decryption?
#include <windows.h>
#include <stdio.h>
#include <bcrypt.h>
#include <iostream>
#include <fstream>
#include <string>
#pragma comment(lib, "bcrypt.lib")
#define NT_SUCCESS(Status) (((NTSTATUS)(Status)) >= 0)
#define STATUS_UNSUCCESSFUL ((NTSTATUS)0xC0000001L)
// Function to decrypt using AES256 algorithm
void ENcryptDecryptString(const std::string& inputString, const std::string& outputFile, bool encrypt){
// Constants for encryption
#define KEY_SIZE 32
#define IV_SIZE 16
const unsigned char ENCRYPTION_KEY[KEY_SIZE] = {
0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0xoF,
0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
0x18, 0x19, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0x1F
};
const unsigned char IV[IV_SIZE] = {
0x20, 0x21, 0x22, 0x23, 0x24, 0x25, 0x26, 0x27,
0x28, 0x29, 0x2A, 0x2B, 0x2C, 0x2D, 0x2E, 0x2F
};
BCRYPT_ALG_HANDLE hAesAlg = NULL;
BCRYPT_KEY_HANDLE hKey = NULL;
NTSTATUS status = STATUS_UNSUCCESSFUL;
DWORD cbCipherText = 0, cbPlainText = 0, cbData = 0, cbKeyObject = 0, cbBlockLen = 0, cbBlob = 0;
PBYTE pbCipherText = NULL, pbPlainText = NULL, pbKeyObject = NULL, pbIV = NULL, pbBlob = NULL;
// Open an algorithm handle
if (!NT_SUCCESS(status = BCryptOpenAlgorithmProvider(&hAesAlg, BCRYPT_AES_ALGORITHM, NULL, 0))){
wprintf(L"**** Error 0x%x returned by BCryptOpenAlgorithmProvider\n", status);
goto Cleanup;
}
// Calculate the size of the buffer to hold the KeyObject
if (!NT_SUCCESS(status = BCryptGetProperty(hAesAlg, BCRYPT_OBJECT_LENGTH, (PBYTE)&cbKeyObject, sizeof(DWORD), &cbData, 0))){
wprintf(L"**** Error 0x%x returned by BCryptGetProperty\n", status);
goto Cleanup;
}
// Allocate the key object on the heap
pbKeyObject = (PBYTE)HeapAlloc(GetProcessHeap(), 0, cbKeyObject);
if (pbKeyObject == NULL){
wprintf(L"**** memory allocation failed\n");
goto Cleanup;
}
// Calculate the block length for the IV
if (!NT_SUCCESS(status = BCryptGetProperty(hAesAlg, BCRYPT_BLOCK_LENGTH, (PBYTE)&cbBlockLen, sizeof(DWORD), &cbData, 0))){
wprintf(L"**** Error 0x%x returned by BCryptGetProperty_IV\n", status);
goto Cleanup;
}
// Allocate a buffer for the IV
pbIV = (PBYTE)HeapAlloc(GetProcessHeap(), 0, cbBlockLen);
if (pbIV == NULL){
wprintf(L"**** IV memory allocation failed\n");
goto Cleanup;
}
// Use the provided IV
memcpy(pbIV, IV, IV_SIZE);
// Set the chaining mode to CBC
if (!NT_SUCCESS(status = BCryptSetProperty(hAesAlg, BCRYPT_CHAINING_MODE, (PBYTE)BCRYPT_CHAIN_MODE_CBC, sizeof(BCRYPT_CHAIN_MODE_CBC), 0))){
wprintf(L"**** Error 0x%x returned by BCryptGetProperty_chaining_mode\n", status);
goto Cleanup;
}
// Generate the key from supplied input key bytes
if (!NT_SUCCESS(status = BCryptGenerateSymmetricKey(hAesAlg, &hKey, pbKeyObject, cbKeyObject, (PBYTE)ENCRYPTION_KEY, KEY_SIZE, 0))){
wprintf(L"**** Error 0x%x returned by BCryptGenerateSymmetricKey\n", status);
goto Cleanup;
}
// Allocate memory for the plaintext or ciphertext buffercbData = inputString.length();
if (encrypt){
pbPlainText = (PBYTE)inputString.c_str();
cbPlainText = cbData;
// Get the size of the encrypted outputFile
if (!NT_SUCCESS(status = BCryptEncrypt(hKey, pbPlainText, cbPlainText, NULL, pbIV, cbBlockLen, NULL, 0, &cbCipherText, BCRYPT_BLOCK_PADDING))){
wprintf(L"**** Error 0x%x returned by BCryptEncrypt\n", status);
goto Cleanup;
}
pbCipherText = (PBYTE)HeapAlloc(GetProcessHeap(), 0, cbCipherText);
if (pbCipherText == NULL){
wprintf(L"**** memory allocation failed_pbCipherText\n");
goto Cleanup;
}
// Perform encryption
if (!NT_SUCCESS(status = BCryptEncrypt(hKey, pbPlainText, cbPlainText, NULL, pbIV, cbBlockLen, pbCipherText, cbCipherText, &cbData, BCRYPT_BLOCK_PADDING))){
wprintf(L"**** Error 0x%x returned by BCryptEncrypt_encryption\n", status);
goto Cleanup;
}
// Write the encrypted data to the output filebuf
std::ofstream ofs(outputFile, std::ios::binary);
if (!ifs){
std::cerr << "Error: Failed to open input file." << std::endl;
goto Cleanup;
}
ofs.write(reinterpret_cast<char*>(pbCipherText), cbCipherText);
} else {
// Decrypting
std::ifstream ifs(inputString, std::ios::binary);
if(!ifs){
std::cerr << "Error: Failed to open input file." << std::endl;
goto Cleanup;
}
// Get the size of the file
ifs.seekg(0, std::ios::end);
std::streampos fileSize = ifs.tellg();
ifs.seekg(0, std::ios::beg);
// Read the file content into the buffer
pbCipherText = (PBYTE)HeapAlloc(GetProcessHeap(), 0, fileSize);
if (pbCipherText == NULL){
wprintf(L"**** memory allocation failed_pbCipherText\n");
goto Cleanup;
}
ifs.read(reinterpret_cast<char*>(pbCipherText), fileSize);
cbCipherText = static_cast<DWORD>(fileSize);
// Decrypt the data
if (!NT_SUCCESS(status = BCryptDecrypt(hKey, pbCipherText, cbCipherText, NULL, pbIV, cbBlockLen, NULL, 0, &cbPlainText, BCRYPT_BLOCK_PADDING))){
wprintf(L"**** Error 0x%x returned by BCryptDecrypt\n", status);
goto Cleanup;
}
pbPlainText = (PBYTE)inputString.c_str();
if (pbPlainText == NULL){
wprintf(L"**** memory allocation failed_pbPlainText\n");
goto Cleanup;
}
// Perform decryption
if (!NT_SUCCESS(status = BCryptDecrypt(hKey, pbCipherText, cbCipherText, NULL, pbIV, cbBlockLen, pbPlainText, cbPlainText, &cbData, BCRYPT_BLOCK_PADDING))){
wprintf(L"**** Error 0x%x returned by BCryptDecrypt_decryption\n", status);
goto Cleanup;
}
// Determine the padding length
BYTE paddingLength = pbPlainText[cbPlainText - 1];
// Remove the padding from the plaintext
cbPlainText -= paddingLength;
// Write the decrypted data to the output file
std::ofstream ofs(outputFile, std::ios::binary);
if (!ofs){
std::cerr << "Error: Failed to open output file." << std::endl;
goto Cleanup;
}
ofs.write(reinterpret_cast<char*>(pbPlainText), cbPlainText);
}
Cleanup:
if (pbCipherText){
HeapFree(GetProcessHeap(), 0, pbCipherText);
}
if (pbPlainText){
HeapFree(GetProcessHeap(), 0, pbPlainText);
}
if (pbKeyObject){
HeapFree(GetProcessHeap(), 0, pbKeyObject);
}
if (pbIV){
HeapFree(GetProcessHeap(), 0, pbIV);
}
if (hKey){
BCryptDestroyKey(hKey);
}
if (hAesAlg){
BCryptCloseAlgorithmProvider(hAesAlg, 0);
}
}
int main(int argc, char* argv[])
{
if (argc != 3){
std::cerr << "Usage: " << argv[0] << " <plain_string> <output_file>" << std::endl;
return 1;
}
std::string inputString = argv[1];
std::string outputFile = argv[2];
EncryptDecryptString(inputString, outputFile, true);
std::cerr << "Encryption successful." << std::endl;
return 0;
}
Upvotes: 1
Views: 114
Reputation: 35440
Here is at least one error:
You do this:
pbPlainText = (PBYTE)inputString.c_str();
Then in Cleanup
, you do this:
if (pbPlainText){
HeapFree(GetProcessHeap(), 0, pbPlainText);
}
Clearly, pbPlainText
was not allocated with HeapAlloc
, yet you are calling HeapFree
using pbPlainText
.
The fix is to remove the HeapFree
for pbPlainText
.
Also, this could never occur:
pbPlainText = (PBYTE)inputString.c_str();
if (pbPlainText == NULL){
The value of c_str()
is always going to be a valid pointer. What it points to may be empty, i.e. pbPlainText.empty()
returns true
, but it will never be NULL.
Also, the better way to write that entire cleanup section of the code is to use RAII, so that not only does it get rid of the goto
game, you will be guaranteed that memory is released regardless of when the function returns. This is highly important if for some reason, an exception is thrown in the code, and you need to ensure that memory is released.
Upvotes: 0