Search code examples
c++encryptionmemory-leaksheap-corruption

Heap corruption occurs when a function or sub-routine is returned during encryption functions


I'm trying to encrypt an AES key with an RSA public_key and print the encrypted bytes out on the console. I'm using Windows Cryptographic APIs to encrypt the key and also using Microsoft's documented way of getting size of buffer that is required for encrypted bytes like in the code provided below:

The issue i am facing is that every time when my program is executed and lockKey function is called, on every return it give error of heap corruption even though all code in the lockKey function has been executed successfully. I've debugged the code, but couldn't find the reason because every time the exception is thrown on function return (ending parenthesis).

Complete code is provided below:

#include <windows.h>
#include <wincrypt.h>
#include <iostream>
#include <vector>
#include <string>
#include <iomanip>
#include <fstream>

#pragma comment(lib, "crypt32.lib")
using namespace std;

void lockKey(std::string, HCRYPTKEY, HCRYPTPROV);
void handleError(const char*);
std::string byteToHex(BYTE*, DWORD);

void handleError(const char* message)
{
    cerr << message << " error code: " << GetLastError() << endl;
    exit(1);
}

std::string byteToHex(BYTE* inputBytes, DWORD inputSize)
{
    std::string outputString;
    DWORD hexSize = 0;
    if (!CryptBinaryToStringA(inputBytes, inputSize, CRYPT_STRING_HEX, NULL, &hexSize))
    {
        // handle error
        handleError("Failed to get size of encrypted bytes key");
        return NULL;
    }

    outputString.resize(hexSize, '\0');
    if (!CryptBinaryToStringA(inputBytes, inputSize, CRYPT_STRING_HEX, &outputString[0], &hexSize))
    {
        // handle error
        handleError("Failed to convert encrypted bytes key to hexSize");
        return NULL;
    }
    //outputString.resize(hexSize - 1); // remove null terminator
    return outputString;
}

void lockKey(std::string publickey, HCRYPTKEY hKey, HCRYPTPROV hCryptProv)
{
    //Import RSA public key
    std::vector<BYTE> publicKeyBytes;
    DWORD publicKeySize = 0;
    if (!CryptStringToBinaryA(publickey.c_str(), publickey.length(), CRYPT_STRING_BASE64, NULL, &publicKeySize, NULL, NULL)) {
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error getting binary size");
    }
    publicKeyBytes.resize(publicKeySize);
    if (!CryptStringToBinaryA(publickey.c_str(), publickey.length(), CRYPT_STRING_BASE64, publicKeyBytes.data(), &publicKeySize, NULL, NULL)) {
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error converting to binary");
    }

    HCRYPTKEY phKey = NULL;
    if (!CryptImportKey(hCryptProv, publicKeyBytes.data(), publicKeyBytes.size(), 0, 0, &phKey)) {
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error importing key");
    }

    //Export AES key for encryption
    //Determine the size of the buffer needed for the exported key
    DWORD dwBufSizeAES = 0;
    if (!CryptExportKey(hKey, NULL, PLAINTEXTKEYBLOB, NULL, NULL, &dwBufSizeAES))
    {
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error determining key buffer size");
    }

    // Allocate a buffer for the exported key
    //BYTE* pbAesKey = new BYTE[dwBufSizeAES];
    //memset(pbAesKey, 0, dwBufSizeAES);
    std::vector<BYTE> pbAesKey(dwBufSizeAES, 0);

    // Export the key to the buffer
    if (!CryptExportKey(hKey, NULL, PLAINTEXTKEYBLOB, NULL, pbAesKey.data(), &dwBufSizeAES))
    {
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error exporting key");
    }

    //Encrypt AES key with RSA public key
    //Determine the size of the buffer needed for the encrypted key
    DWORD dwBufSize = 0;
    if (!CryptEncrypt(phKey, NULL, TRUE, 0, NULL, &dwBufSize, 0))
    {
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error determining encrypted buffer size");
    }

    //PBYTE pbEncryptedKey = (PBYTE)malloc(dwBufSize);
    std::vector<BYTE> pbEncryptedKey(dwBufSize, 0);

    // Encrypt the AES key with RSA-OAEP padding
    DWORD cbAesKey = 0;
    if (!CryptEncrypt(phKey, NULL, TRUE, 0, pbAesKey.data(), &cbAesKey, dwBufSize))
    {
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error encrypting AES key");
    }

    memcpy(pbEncryptedKey.data(), pbAesKey.data(), dwBufSize);   

    // Now pbEncryptedKey contains the encrypted AES key
    std::string hexEncryptedKey = byteToHex(pbEncryptedKey.data(), dwBufSize);
    cout << "Encrypted AES Key Bytes: \n" << hexEncryptedKey << endl;
    CryptDestroyKey(phKey);

} //error always on this line of code.

int main()
{
    DWORD blockSize;
    HCRYPTKEY hKey;
    HCRYPTPROV hCryptProv;

    // Import the RSA public key as in the previous example.
    std::string publickey = "valid public key"; //removed

    // Generate aes key
    hCryptProv;
    if (!CryptAcquireContext(&hCryptProv, NULL, NULL, PROV_RSA_AES, CRYPT_VERIFYCONTEXT)) {
        CryptReleaseContext(hCryptProv, 0);
        handleError("Failed to acquire cryptographic context");
    }
    hKey;
    if (!CryptGenKey(hCryptProv, CALG_AES_256, CRYPT_EXPORTABLE, &hKey)) {
        CryptReleaseContext(hCryptProv, 0);
        handleError("Failed to generate key");
    }
    // Get the block size for the algorithm
    blockSize;
    DWORD dwDataLen = sizeof(DWORD);
    if (!CryptGetKeyParam(hKey, KP_BLOCKLEN, (BYTE*)&blockSize, &dwDataLen, 0)) {
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Failed to get block size");
    }
        
    lockKey(publickey, hKey, hCryptProv);
    
    CryptDestroyKey(hKey);
    CryptReleaseContext(hCryptProv, 0);
    return 0;

}

The screenshots for errors are as follows:

Exception in visual studio

If continued then it give heap corruption error:

Heap corruption

I was using memory management functions at first like memset, malloc, free, delete etc. But because of heap errors. I've shifted my buffers to vectors which i understand are much better to use avoiding memory corruption. However, my error is still not solved. I've debugged the program and check available values and buffer of each variable as well. Everything is at it should. I am missing something. Any help would be appreciated.


Solution

  • I've asked the same question to ChatGPT-4 and got the accurate answer. It looks like the issue was the way i was handling the encrypted key's buffer. When I call CryptEncrypt, the pbAesKey buffer is being overwritten with the encrypted data. However, I have allocated pbEncryptedKey to store the encrypted key separately. Instead of copying the encrypted data from pbAesKey to pbEncryptedKey, I should have been encrypting the data directly into pbEncryptedKey.

    So replacing the code

    if (!CryptEncrypt(phKey, NULL, TRUE, 0, pbAesKey.data(), &cbAesKey, dwBufSize)){
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error encrypting AES key");
    }
    memcpy(pbEncryptedKey.data(), pbAesKey.data(), dwBufSize);
    

    with

    memcpy(pbEncryptedKey.data(), pbAesKey.data(), dwBufSizeAES);
    if (!CryptEncrypt(phKey, NULL, TRUE, 0, pbEncryptedKey.data(), &dwBufSizeAES, dwBufSize)){
        CryptDestroyKey(phKey);
        CryptDestroyKey(hKey);
        CryptReleaseContext(hCryptProv, 0);
        handleError("Error encrypting AES key");
    }
    

    resolved the issue.

    This way, I am first copying the AES key into pbEncryptedKey, and then using pbEncryptedKey as the buffer for the encrypted data. After that, I can safely use pbEncryptedKey.data() to convert the encrypted key to a hex string without any heap corruption issues.