user23533541
user23533541

Reputation: 13

Why does my code hang after encrypting to AES 256?

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

Answers (1)

PaulMcKenzie
PaulMcKenzie

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

Related Questions