Repository: duckdb/duckdb
File: third_party/mbedtls/mbedtls_wrapper.cpp (C++)

Code before:
	}
}

void MbedTlsWrapper::AESStateMBEDTLS::GenerateRandomData(duckdb::data_ptr_t data, duckdb::idx_t len) {
	GenerateRandomDataInsecure(data, len);
}

void MbedTlsWrapper::AESStateMBEDTLS::InitializeInternal(duckdb::const_data_ptr_t iv, duckdb::idx_t iv_len, duckdb::const_data_ptr_t aad, duckdb::idx_t aad_len){
	if (mbedtls_cipher_set_iv(context.get(), iv, iv_len)) {
		throw runtime_error("Failed to set IV for encryption");
	}

	if (aad_len > 0) {
		if (mbedtls_cipher_update_ad(context.get(), aad, aad_len)) {
			throw std::runtime_error("Failed to set AAD");
		}
	}
}

void MbedTlsWrapper::AESStateMBEDTLS::InitializeEncryption(duckdb::const_data_ptr_t iv, duckdb::idx_t iv_len, duckdb::const_data_ptr_t key, duckdb::idx_t key_len_p, duckdb::const_data_ptr_t aad, duckdb::idx_t aad_len) {
	mode = duckdb::EncryptionTypes::ENCRYPT;

	if (mbedtls_cipher_setkey(context.get(), key, key_len_p * 8, MBEDTLS_ENCRYPT) != 0) {
		runtime_error("Failed to set AES key for encryption");
	}

	InitializeInternal(iv, iv_len, aad, aad_len);
}

void MbedTlsWrapper::AESStateMBEDTLS::InitializeDecryption(duckdb::const_data_ptr_t iv, duckdb::idx_t iv_len, duckdb::const_data_ptr_t key, duckdb::idx_t key_len_p, duckdb::const_data_ptr_t aad, duckdb::idx_t aad_len) {
	mode = duckdb::EncryptionTypes::DECRYPT;

	if (key_len_p != key_len) {
		throw duckdb::InternalException("Invalid encryption key length, expected %llu, got %llu", key_len, key_len_p);
	}
	if (mbedtls_cipher_setkey(context.get(), key, key_len * 8, MBEDTLS_DECRYPT)) {
		throw runtime_error("Failed to set AES key for encryption");
	}

	InitializeInternal(iv, iv_len, aad, aad_len);
}

size_t MbedTlsWrapper::AESStateMBEDTLS::Process(duckdb::const_data_ptr_t in, duckdb::idx_t in_len, duckdb::data_ptr_t out,
                                                   duckdb::idx_t out_len) {

	// GCM works in-place, CTR and CBC don't
	auto use_out_copy = in == out && cipher != duckdb::EncryptionTypes::CipherType::GCM;

	auto out_ptr = out;
	std::unique_ptr<duckdb::data_t[]> out_copy;
	if (use_out_copy) {

Review comment (security):
I'm slightly worried that enabling this could open the door to security issues in the future because it will make it really easy to accidentally invoke this even though its not secure. Could we maybe look into hardening this against refactors accidentally calling this without the user opting into insure mbedtls usage?

Code after:
		mbedtls_cipher_free(context.get());
	}
}

void MbedTlsWrapper::AESStateMBEDTLS::GenerateRandomDataInsecure(duckdb::data_ptr_t data, duckdb::idx_t len) {
	if (!force_mbedtls) {
		// To use this insecure MbedTLS random number generator
		// we double check if force_mbedtls_unsafe is set
		// such that we do not accidentaly opt-in
		throw duckdb::InternalException("Insecure random generation called without setting 'force_mbedtls_unsafe' = true");
	}

	duckdb::RandomEngine random_engine;

	while (len != 0) {
		const auto random_integer = random_engine.NextRandomInteger();
		const auto next = duckdb::MinValue<duckdb::idx_t>(len, sizeof(random_integer));
		memcpy(data, duckdb::const_data_ptr_cast(&random_integer), next);
		data += next;
		len -= next;
	}
}

void MbedTlsWrapper::AESStateMBEDTLS::GenerateRandomData(duckdb::data_ptr_t data, duckdb::idx_t len) {
	// generate insecure random data
	GenerateRandomDataInsecure(data, len);
}

void MbedTlsWrapper::AESStateMBEDTLS::InitializeInternal(duckdb::const_data_ptr_t iv, duckdb::idx_t iv_len, duckdb::const_data_ptr_t aad, duckdb::idx_t aad_len){
	if (mbedtls_cipher_set_iv(context.get(), iv, iv_len)) {
		throw runtime_error("Failed to set IV for encryption");
	}

	if (aad_len > 0) {
		if (mbedtls_cipher_update_ad(context.get(), aad, aad_len)) {
			throw std::runtime_error("Failed to set AAD");
		}
	}
}

void MbedTlsWrapper::AESStateMBEDTLS::InitializeEncryption(duckdb::const_data_ptr_t iv, duckdb::idx_t iv_len, duckdb::const_data_ptr_t key, duckdb::idx_t key_len_p, duckdb::const_data_ptr_t aad, duckdb::idx_t aad_len) {
	mode = duckdb::EncryptionTypes::ENCRYPT;

	if (mbedtls_cipher_setkey(context.get(), key, key_len_p * 8, MBEDTLS_ENCRYPT) != 0) {
		runtime_error("Failed to set AES key for encryption");
	}

	InitializeInternal(iv, iv_len, aad, aad_len);
}

void MbedTlsWrapper::AESStateMBEDTLS::InitializeDecryption(duckdb::const_data_ptr_t iv, duckdb::idx_t iv_len, duckdb::const_data_ptr_t key, duckdb::idx_t key_len_p, duckdb::const_data_ptr_t aad, duckdb::idx_t aad_len) {