picotlsvs/bcrypt-test/bcrypt-test.c - third_party/github/h2o/picotls - Git at Google

 // bcrypt-test.cpp : This file contains the 'main' function. Program execution begins and ends there.
 //
 #include <wincompat.h>
 #include <bcrypt.h>
 #include <stdio.h>
 #include "picotls/ptlsbcrypt.h"
 #include "picotls/minicrypto.h"

 int KeyInit(BCRYPT_KEY_HANDLE *hKey, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz, const BYTE *proposedKey,
             DWORD proposedKeyLength, BYTE **ko, ULONG *ko_length)
 {
     DWORD cbData = 0;
     HANDLE hAlgo = NULL;

     // Open an algorithm handle.
     NTSTATUS ret = BCryptOpenAlgorithmProvider(&hAlgo, name, NULL, 0);

     if (BCRYPT_SUCCESS(ret)) {
         // Set the properties to define the chaining mode
         ret = BCryptSetProperty(hAlgo, BCRYPT_CHAINING_MODE, (PBYTE)chain_mode, (ULONG)chain_mode_sz, 0);
     }

     *ko = NULL;
     *ko_length = 0;

     if (BCRYPT_SUCCESS(ret)) {
         DWORD ko_size = 0;
         ULONG cbResult = 0;

         ret = BCryptGetProperty(hAlgo, BCRYPT_OBJECT_LENGTH, (PUCHAR)&ko_size, (ULONG)sizeof(ko_size), &cbResult, 0);

         if (BCRYPT_SUCCESS(ret)) {
             *ko = (uint8_t *)malloc(ko_size);
             if (*ko == NULL) {
                 ret = STATUS_NO_MEMORY;
             } else {
                 *ko_length = ko_size;
                 memset(*ko, 0, *ko_length);
             }
         }
     }

     if (BCRYPT_SUCCESS(ret)) {
         // Generate the key from supplied input key bytes.
         ret = BCryptGenerateSymmetricKey(hAlgo, hKey, *ko, *ko_length, (PBYTE)proposedKey, proposedKeyLength, 0);
     } else {
         if (*ko != NULL) {
             free(*ko);
             *ko = NULL;
             *ko_length = 0;
         }
     }

     if (hAlgo != NULL) {
         BCryptCloseAlgorithmProvider(hAlgo, 0);
     }

     return BCRYPT_SUCCESS(ret) ? 0 : -1;
 }

 void KeyRelease(BCRYPT_KEY_HANDLE *hKey, BYTE **ko, ULONG *ko_length)
 {
     BCryptDestroyKey(*hKey);
     *hKey = NULL;
     if (*ko) {
         free(*ko);
     }
     *ko = NULL;
     *ko_length = 0;
 }

 int EncodeOneShot(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz, BYTE *key,
                   ULONG key_length,
     BYTE* iv, ULONG iv_length,
     BYTE *data, ULONG dataLength, uint64_t seq, BYTE *authData, ULONG authDataLength, ULONG authTagLength,
     BYTE *encrypted, ULONG encryptedLengthMax, ULONG *encryptedLength)
 {

     BCRYPT_KEY_HANDLE hKey = NULL;
     BYTE *authTag = encrypted + dataLength;
     BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO bacmi;
     BYTE iv_nonce[PTLS_MAX_IV_SIZE];
     BYTE *ko = NULL;
     ULONG ko_length = 0;
     int ret = 0;

     *encryptedLength = 0;

     if (KeyInit(&hKey, name, chain_mode, chain_mode_sz, key, key_length, &ko, &ko_length) != 0) {
         return -1;
     }

     memset(authTag, 0, authTagLength);
     // Set the auth mode info for AEAD
     BCRYPT_INIT_AUTH_MODE_INFO(bacmi);
     ptls_aead__build_iv(aead, iv_nonce, iv, seq);
     bacmi.pbNonce = iv_nonce;
     bacmi.cbNonce = iv_length;
     bacmi.pbAuthData = authData;
     bacmi.cbAuthData = authDataLength;
     bacmi.pbTag = authTag;
     bacmi.cbTag = authTagLength;
     /* All other fields are set to NULL by the INIT macro. */

     /* If called with a NULL pointer for the data block, we will merely compute the block size. */
     DWORD cbCipherText = 0;
     NTSTATUS status = BCryptEncrypt(hKey, data, dataLength, &bacmi, NULL, 0, encrypted, encryptedLengthMax, &cbCipherText, 0);

     KeyRelease(&hKey, &ko, &ko_length);

     if (BCRYPT_SUCCESS(status)) {
         *encryptedLength = cbCipherText + authTagLength;
     } else {
         ret = -1;
     }

     return ret;
 }

 int DecodeOneShot(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz,
     BYTE *key, ULONG key_length, BYTE * iv, ULONG iv_length, BYTE *encrypted,
                   ULONG encryptedLength, uint64_t seq, BYTE *authData, ULONG authDataLength,
                   ULONG authTagLength, BYTE *decrypted, ULONG decryptedLengthMax, ULONG *decryptedLength)
 {

     BCRYPT_KEY_HANDLE hKey = NULL;
     BYTE *authTag = encrypted + (encryptedLength - authTagLength);
     BYTE iv_nonce[PTLS_MAX_IV_SIZE];
     BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO bacmi;
     BYTE *ko = NULL;
     ULONG ko_length = 0;
     int ret = 0;

     *decryptedLength = 0;

     if (KeyInit(&hKey, name, chain_mode, chain_mode_sz, key, key_length, &ko, &ko_length) != 0) {
         return -1;
     }

     // Set the auth mode info for AEAD
     BCRYPT_INIT_AUTH_MODE_INFO(bacmi);
     ptls_aead__build_iv(aead, iv_nonce, iv, seq);
     bacmi.pbNonce = iv_nonce;
     bacmi.cbNonce = iv_length;
     bacmi.pbAuthData = authData;
     bacmi.cbAuthData = authDataLength;
     bacmi.pbTag = authTag;
     bacmi.cbTag = authTagLength;
     /* All other fields are set to NULL by the INIT macro. */

     /* If called with a NULL pointer for the data block, we will merely compute the block size. */
     DWORD cbCipherText = 0;
     NTSTATUS status = BCryptDecrypt(hKey, encrypted, encryptedLength - authTagLength, &bacmi, NULL, 0, decrypted,
                                     decryptedLengthMax, &cbCipherText, 0);

     KeyRelease(&hKey, &ko, &ko_length);

     if (BCRYPT_SUCCESS(status)) {
         *decryptedLength = cbCipherText;
     } else {
         ret = -1;
     }

     return ret;
 }

 int test_oneshot(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz)
 {
     BYTE key[32];
     BYTE data[123];
     BYTE iv[PTLS_MAX_IV_SIZE];
     uint64_t nonce;
     BYTE authData[9];
     BYTE encrypted[256];
     ULONG encryptedLength;
     BYTE decrypted[256];
     ULONG decryptedLength;
     ULONG authTagLength = (ULONG)aead->tag_size;
     int ret = 0;

     assert(sizeof(key) >= aead->key_size);
     assert(sizeof(iv) >= aead->iv_size);
     assert(sizeof(data) + authTagLength <= sizeof(encrypted));
     assert(sizeof(decrypted) >= sizeof(encrypted));

     memset(key, 'k', sizeof(key));
     memset(data, 'd', sizeof(data));
     memset(iv, 'n', sizeof(iv));
     nonce = 0;
     memset(authData, 'a', sizeof(authData));

     ret = EncodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size, iv, (ULONG)aead->iv_size, data,
         123, nonce, authData, 9, authTagLength, encrypted, 256, &encryptedLength);

     printf("Encrypt one shot returns %d, l=%d\n", ret, encryptedLength);

     if (ret == 0) {
         ret = DecodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size,
             iv, (ULONG)aead->iv_size, encrypted, encryptedLength, nonce, authData, 9,
             authTagLength, decrypted, 256, &decryptedLength);

         printf("Decrypt one shot returns %d, l=%d\n", ret, decryptedLength);

         if (ret == 0) {
             if (decryptedLength != 123) {
                 printf("Wrong length, not %d\n", 123);
                 ret = -1;
             } else if (memcmp(data, decrypted, 123) != 0) {
                 printf("Data and decrypted don't match\n");
                 ret = -1;
             } else {
                 printf("One shot matches.\n");
             }
         }
     }

     return ret;
 }

 void delete_test_aead_context(ptls_aead_context_t *ctx)
 {
     if (ctx != NULL) {
         ctx->dispose_crypto(ctx);
         free(ctx);
     }
 }

 ptls_aead_context_t *new_test_aead_context(ptls_aead_algorithm_t *aead, int is_enc, BYTE *key, BYTE *iv)
 {
     int ret = 0;
     ptls_aead_context_t *ctx = (ptls_aead_context_t *)malloc(aead->context_size);

     if (ctx != NULL) {
         memset(ctx, 0, aead->context_size);
         *ctx = (ptls_aead_context_t){aead};
         if (aead->setup_crypto(ctx, is_enc, key, iv) != 0) {
             printf("For %s, setup returns %d\n", aead->name, ret);
             delete_test_aead_context(ctx);
             ctx = NULL;
         }
     } else {
         printf("For %s, memory error during setup\n", aead->name);
     }

     return (ctx);
 }

 int test_decrypt(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz)
 {
     BYTE key[32];
     BYTE iv[PTLS_MAX_IV_SIZE];
     BYTE data[123];
     uint64_t nonce;
     BYTE authData[9];
     BYTE encrypted[256];
     ULONG encryptedLength;
     BYTE decrypted[256];
     size_t decryptedLength;
     ULONG authTagLength = (ULONG)aead->tag_size;
     ptls_aead_context_t *ctx = NULL;
     int ret = 0;

     assert(sizeof(key) >= aead->key_size);
     assert(sizeof(iv) >= aead->iv_size);
     assert(sizeof(data) + authTagLength <= sizeof(encrypted));
     assert(sizeof(decrypted) >= sizeof(encrypted));

     memset(key, 'k', sizeof(key));
     memset(iv, 'n', sizeof(iv));
     memset(data, 'd', sizeof(data));
     nonce = 0;
     memset(authData, 'a', sizeof(authData));

     /* Create a decryption context */
     ctx = new_test_aead_context(aead, 0, key, iv);
     if (ctx == NULL) {
         ret = -1;
     }

     /* Do a simple encrypt using one shot bcrypt */
     if (ret == 0) {
         ret = EncodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size,
             iv, (ULONG)aead->iv_size, data, 123, nonce,
             authData, 9, authTagLength, encrypted, 256, &encryptedLength);
     }

     /* Try decrypt with library procedure */
     if (ret == 0) {
         decryptedLength = ctx->do_decrypt(ctx, decrypted, encrypted, encryptedLength, nonce, authData, 9);
         if (decryptedLength >= encryptedLength) {
             printf("For %s, decrypt returns %d\n", aead->name, (int)decryptedLength);
             ret = -1;
         } else if (decryptedLength != 123) {
             printf("For %s, decrypt returns %d instead of %d\n", aead->name, (int)decryptedLength, 123);
             ret = -1;
         } else if (memcmp(data, decrypted, decryptedLength) != 0) {
             printf("For %s, decrypted does not match clear text\n", aead->name);
             ret = -1;
         } else {
             printf("For %s, decrypting test passes.\n", aead->name);
         }
     }

     delete_test_aead_context(ctx);

     return ret;
 }

 int test_encrypt(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz)
 {
     BYTE key[32];
     BYTE iv[PTLS_MAX_IV_SIZE];
     BYTE data[123];
     uint64_t nonce;
     BYTE authData[9];
     BYTE encryptedRef[256];
     ULONG encryptedRefLength;
     BYTE encrypted[256];
     size_t encryptedLength;
     ULONG authTagLength = (ULONG)aead->tag_size;
     ptls_aead_context_t *ctx = NULL;
     int ret = 0;

     assert(sizeof(key) >= aead->key_size);
     assert(sizeof(iv) >= aead->iv_size);
     assert(sizeof(data) + authTagLength <= sizeof(encrypted));
     assert(sizeof(data) + authTagLength <= sizeof(encryptedRef));

     memset(key, 'k', sizeof(key));
     memset(iv, 'n', sizeof(iv));
     memset(data, 'd', sizeof(data));
     nonce = 0;
     memset(authData, 'a', sizeof(authData));

     /* Create an encryption context */
     ctx = new_test_aead_context(aead, 1, key, iv);
     if (ctx == NULL) {
         ret = -1;
     }

     /* Do a simple encrypt using one shot bcrypt */
     if (ret == 0) {
         ret = EncodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size,
             iv, (ULONG)aead->iv_size, data, 123, nonce,
             authData, 9, authTagLength, encryptedRef, 256, &encryptedRefLength);
     }

     /* Try encrypt with library procedure */
     if (ret == 0) {
         ctx->do_encrypt_init(ctx, nonce, authData, 9);
         encryptedLength = ctx->do_encrypt_update(ctx, encrypted, data, 123);
         encryptedLength += ctx->do_encrypt_final(ctx, &encrypted[encryptedLength]);

         if (encryptedLength != encryptedRefLength) {
             printf("For %s, encrypt returns %d instead of %d\n", aead->name, (int)encryptedLength, encryptedRefLength);
             ret = -1;
         } else if (memcmp(encryptedRef, encrypted, encryptedRefLength) != 0) {
             printf("For %s, encrypted does not match ref\n", aead->name);
             for (ULONG i = 0; i < encryptedRefLength; i++) {
                 if (encryptedRef[i] != encrypted[i]) {
                     printf("For %s, encrypted[%d] = 0x%02x vs encryptedRef[%d] = 0x%02x\n", aead->name, i, encrypted[i], i,
                            encryptedRef[i]);
                     break;
                 }
             }
             ret = -1;
         } else {
             printf("For %s, encrypting test passes.\n", aead->name);
         }
     }

     delete_test_aead_context(ctx);

     return ret;
 }

 int test_for_size(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz)
 {
     BYTE key[32];
     BYTE iv[PTLS_MAX_IV_SIZE];
     uint64_t nonce;
     BYTE authData[9];
     BYTE *data = NULL;
     BYTE *encrypted = NULL;
     BYTE *decrypted = NULL;
     size_t encryptedLength;
     size_t decryptedLength;
     ULONG authTagLength = (ULONG)aead->tag_size;
     ptls_aead_context_t *ctx_e = NULL;
     ptls_aead_context_t *ctx_d = NULL;
     ULONG packet_size[] = {1500, 128, 3, 0};
     ULONG nb_packet_size = (ULONG)(sizeof(packet_size) / sizeof(ULONG));
     int ret = 0;

     assert(sizeof(key) >= aead->key_size);
     assert(sizeof(iv) >= aead->iv_size);

     memset(key, 'k', sizeof(key));
     memset(key, 'n', sizeof(iv));
     nonce = 0;
     memset(authData, 'a', sizeof(authData));

     /* Create the encryption contexts */
     ctx_e = new_test_aead_context(aead, 1, key, iv);
     ctx_d = new_test_aead_context(aead, 0, key, iv);

     if (ctx_e == NULL || ctx_d == NULL) {
         ret = -1;
     }

     /* Test a variety of packet sizes */
     for (ULONG i = 0; ret == 0 && i < nb_packet_size; i++) {
         ULONG data_size = (packet_size[i] > 0) ? packet_size[i] : 128;
         ULONG encrypted_size = packet_size[i] + authTagLength;

         data = (BYTE *)malloc(data_size);
         encrypted = (BYTE *)malloc(encrypted_size);
         decrypted = (BYTE *)malloc(data_size);

         if (data == NULL || encrypted == NULL || decrypted == NULL) {
             printf("For %s: cannot allocate memory for packet size[%d] = %d\n", aead->name, i, packet_size[i]);
         } else {
             memset(data, 'd', data_size);

             ctx_e->do_encrypt_init(ctx_e, nonce, authData, 9);
             encryptedLength = ctx_e->do_encrypt_update(ctx_e, encrypted, data, packet_size[i]);
             encryptedLength += ctx_e->do_encrypt_final(ctx_e, &encrypted[encryptedLength]);
             decryptedLength = ctx_d->do_decrypt(ctx_d, decrypted, encrypted, encryptedLength, nonce, authData, 9);

             if (decryptedLength >= encryptedLength) {
                 printf("For %s, decrypt returns %d\n", aead->name, (int)decryptedLength);
                 ret = -1;
             } else if (decryptedLength != packet_size[i]) {
                 printf("For %s, decrypt returns %d instead of %d\n", aead->name, (int)decryptedLength, packet_size[i]);
                 ret = -1;
             } else if (memcmp(data, decrypted, decryptedLength) != 0) {
                 printf("For %s, decrypted does not match clear text\n", aead->name);
                 ret = -1;
             } else {
                 printf("For %s, test packet size[%d] = %d passes.\n", aead->name, i, packet_size[i]);
             }
         }

         if (data != NULL) {
             free(data);
             data = NULL;
         }

         if (encrypted != NULL) {
             free(encrypted);
             encrypted = NULL;
         }

         if (decrypted != NULL) {
             free(decrypted);
             decrypted = NULL;
         }
     }

     delete_test_aead_context(ctx_e);
     delete_test_aead_context(ctx_d);

     return ret;
 }

 int test_one_aead(ptls_aead_algorithm_t *aead, wchar_t *name, wchar_t *chain_mode, size_t chain_mode_sz)
 {
     int ret = test_oneshot(aead, name, chain_mode, chain_mode_sz);

     printf("For %s, test one shot returns %d\n", aead->name, ret);

     if (ret == 0) {
         ret = test_decrypt(aead, name, chain_mode, chain_mode_sz);

         printf("For %s, test decrypt returns %d\n", aead->name, ret);
     }

     if (ret == 0) {
         ret = test_encrypt(aead, name, chain_mode, chain_mode_sz);

         printf("For %s, test encrypt returns %d\n", aead->name, ret);
     }

     if (ret == 0) {
         ret = test_for_size(aead, name, chain_mode, chain_mode_sz);

         printf("For %s, test packet sizes returns %d\n", aead->name, ret);
     }

     return ret;
 }

 /* Test of cipher functions.
  * The test verifies that a message encode with a bcrypt function can be
  * decoded with a minicrypto function, and vice versa.
  */

 int test_cipher_one_way(char const *name1, char const *name2, ptls_cipher_algorithm_t *b1, ptls_cipher_algorithm_t *b2,
                         unsigned int nb_blocks)
 {
     BYTE key[32];
     BYTE nonce[16];
     BYTE data[128];
     BYTE encrypted[128];
     BYTE decrypted[128];
     size_t data_size = b1->block_size * nb_blocks;
     ptls_cipher_context_t *ctx1 = NULL;
     ptls_cipher_context_t *ctx2 = NULL;
     int ret = 0;

     assert(sizeof(key) >= b1->key_size);
     assert(sizeof(data) >= data_size);
     assert(sizeof(nonce) >= b1->iv_size);

     memset(key, 'k', sizeof(key));
     memset(data, 'd', data_size);

     ctx1 = ptls_cipher_new(b1, 1, key);
     ctx2 = ptls_cipher_new(b2, 0, key);

     if (ctx1 == NULL || ctx2 == NULL) {
         ret = -1;
     } else {
         memset(nonce, 0, sizeof(nonce));
         if (ctx1->do_init != NULL) {
             ctx1->do_init(ctx1, nonce);
         }

         if (ctx2->do_init != NULL) {
             ctx2->do_init(ctx2, nonce);
         }

         ctx1->do_transform(ctx1, encrypted, data, data_size);
         ctx2->do_transform(ctx2, decrypted, encrypted, data_size);

         if (memcmp(data, decrypted, data_size) != 0) {
             printf("For %s -> %s, decrypted does not match clear text\n", name1, name2);
             ret = -1;
         } else {
             printf("For %s -> %s, test passes.\n", name1, name2);
         }
     }

     if (ctx1 != NULL) {
         ptls_cipher_free(ctx1);
     }

     if (ctx2 != NULL) {
         ptls_cipher_free(ctx2);
     }

     return ret;
 }

 int test_cipher_pair(char const *name1, char const *name2, ptls_cipher_algorithm_t *b1, ptls_cipher_algorithm_t *b2,
                      unsigned int nb_blocks)
 {
     int ret = test_cipher_one_way(name1, name2, b1, b2, nb_blocks);

     if (ret == 0) {
         ret = test_cipher_one_way(name2, name1, b2, b1, nb_blocks);
     }

     return ret;
 }

 /* Test of the hash functions
  */

 int test_hash_calc(char const *name1, char const *name2, ptls_hash_algorithm_t *h1, ptls_hash_algorithm_t *h2)
 {
     BYTE data[123];
     BYTE tag1[128];
     BYTE tag2[128];
     ptls_hash_context_t *ctx1 = NULL;
     ptls_hash_context_t *ctx2 = NULL;
     int ret = 0;

     assert(sizeof(tag1) >= h1->digest_size);
     assert(sizeof(tag2) >= h2->digest_size);
     assert(h1->digest_size == h2->digest_size);

     memset(data, 'd', sizeof(data));
     memset(tag1, '1', sizeof(tag1));
     memset(tag2, '2', sizeof(tag2));

     if (h1->digest_size != h2->digest_size) {
         ret = -1;
     }
     if (ret == 0) {
         ret = ptls_calc_hash(h1, tag1, data, sizeof(data));
     }

     if (ret == 0) {
         ret = ptls_calc_hash(h2, tag2, data, sizeof(data));
     }

     if (ret == 0){
         if (memcmp(tag1, tag2, h1->digest_size) != 0) {
             printf("For %s -> %s, hash1 does not match hash2\n", name1, name2);
             ret = -1;
         } else {
             printf("For %s -> %s, hash test passes.\n", name1, name2);
         }
     }

     return ret;
 }

 /* Minimal test program for the bcrypt functions.
  * Need to add tests for the SHA256 and SHA384 implementations.
  */

 int main()
 {
     int ret = 0;

     ret |= test_cipher_pair("bcrypt aes128ecb", "minicrypto aes128ecb", &ptls_bcrypt_aes128ecb, &ptls_minicrypto_aes128ecb, 1);
     ret |= test_cipher_pair("bcrypt aes256ecb", "minicrypto aes256ecb", &ptls_bcrypt_aes256ecb, &ptls_minicrypto_aes256ecb, 1);

     ret |= test_cipher_pair("bcrypt aes128ctr", "minicrypto aes128ctr", &ptls_bcrypt_aes128ctr, &ptls_minicrypto_aes128ctr, 4);
     ret |= test_cipher_pair("bcrypt aes256ctr", "minicrypto aes256ctr", &ptls_bcrypt_aes256ctr, &ptls_minicrypto_aes256ctr, 4);

     ret |= test_one_aead(&ptls_bcrypt_aes128gcm, BCRYPT_AES_ALGORITHM, BCRYPT_CHAIN_MODE_GCM, sizeof(BCRYPT_CHAIN_MODE_GCM));

     ret |= test_one_aead(&ptls_bcrypt_aes256gcm, BCRYPT_AES_ALGORITHM, BCRYPT_CHAIN_MODE_GCM, sizeof(BCRYPT_CHAIN_MODE_GCM));

     ret |= test_hash_calc("bcrypt sha256", "minicrypto sha256", &ptls_bcrypt_sha256, &ptls_minicrypto_sha256);
     ret |= test_hash_calc("bcrypt sha384", "minicrypto sha384", &ptls_bcrypt_sha384, &ptls_minicrypto_sha384);

     exit(ret);
 }
	// bcrypt-test.cpp : This file contains the 'main' function. Program execution begins and ends there.
	//
	#include <wincompat.h>
	#include <bcrypt.h>
	#include <stdio.h>
	#include "picotls/ptlsbcrypt.h"
	#include "picotls/minicrypto.h"

	int KeyInit(BCRYPT_KEY_HANDLE hKey, wchar_t name, wchar_t chain_mode, size_t chain_mode_sz, const BYTE proposedKey,
	DWORD proposedKeyLength, BYTE *ko, ULONG ko_length)
	{
	DWORD cbData = 0;
	HANDLE hAlgo = NULL;

	// Open an algorithm handle.
	NTSTATUS ret = BCryptOpenAlgorithmProvider(&hAlgo, name, NULL, 0);

	if (BCRYPT_SUCCESS(ret)) {
	// Set the properties to define the chaining mode
	ret = BCryptSetProperty(hAlgo, BCRYPT_CHAINING_MODE, (PBYTE)chain_mode, (ULONG)chain_mode_sz, 0);
	}

	*ko = NULL;
	*ko_length = 0;

	if (BCRYPT_SUCCESS(ret)) {
	DWORD ko_size = 0;
	ULONG cbResult = 0;

	ret = BCryptGetProperty(hAlgo, BCRYPT_OBJECT_LENGTH, (PUCHAR)&ko_size, (ULONG)sizeof(ko_size), &cbResult, 0);

	if (BCRYPT_SUCCESS(ret)) {
	ko = (uint8_t )malloc(ko_size);
	if (*ko == NULL) {
	ret = STATUS_NO_MEMORY;
	} else {
	*ko_length = ko_size;
	memset(ko, 0, ko_length);
	}
	}
	}

	if (BCRYPT_SUCCESS(ret)) {
	// Generate the key from supplied input key bytes.
	ret = BCryptGenerateSymmetricKey(hAlgo, hKey, ko, ko_length, (PBYTE)proposedKey, proposedKeyLength, 0);
	} else {
	if (*ko != NULL) {
	free(*ko);
	*ko = NULL;
	*ko_length = 0;
	}
	}

	if (hAlgo != NULL) {
	BCryptCloseAlgorithmProvider(hAlgo, 0);
	}

	return BCRYPT_SUCCESS(ret) ? 0 : -1;
	}

	void KeyRelease(BCRYPT_KEY_HANDLE hKey, BYTE ko, ULONG ko_length)
	{
	BCryptDestroyKey(*hKey);
	*hKey = NULL;
	if (*ko) {
	free(*ko);
	}
	*ko = NULL;
	*ko_length = 0;
	}

	int EncodeOneShot(ptls_aead_algorithm_t aead, wchar_t name, wchar_t chain_mode, size_t chain_mode_sz, BYTE key,
	ULONG key_length,
	BYTE* iv, ULONG iv_length,
	BYTE data, ULONG dataLength, uint64_t seq, BYTE authData, ULONG authDataLength, ULONG authTagLength,
	BYTE encrypted, ULONG encryptedLengthMax, ULONG encryptedLength)
	{

	BCRYPT_KEY_HANDLE hKey = NULL;
	BYTE *authTag = encrypted + dataLength;
	BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO bacmi;
	BYTE iv_nonce[PTLS_MAX_IV_SIZE];
	BYTE *ko = NULL;
	ULONG ko_length = 0;
	int ret = 0;

	*encryptedLength = 0;

	if (KeyInit(&hKey, name, chain_mode, chain_mode_sz, key, key_length, &ko, &ko_length) != 0) {
	return -1;
	}

	memset(authTag, 0, authTagLength);
	// Set the auth mode info for AEAD
	BCRYPT_INIT_AUTH_MODE_INFO(bacmi);
	ptls_aead__build_iv(aead, iv_nonce, iv, seq);
	bacmi.pbNonce = iv_nonce;
	bacmi.cbNonce = iv_length;
	bacmi.pbAuthData = authData;
	bacmi.cbAuthData = authDataLength;
	bacmi.pbTag = authTag;
	bacmi.cbTag = authTagLength;
	/* All other fields are set to NULL by the INIT macro. */

	/* If called with a NULL pointer for the data block, we will merely compute the block size. */
	DWORD cbCipherText = 0;
	NTSTATUS status = BCryptEncrypt(hKey, data, dataLength, &bacmi, NULL, 0, encrypted, encryptedLengthMax, &cbCipherText, 0);

	KeyRelease(&hKey, &ko, &ko_length);

	if (BCRYPT_SUCCESS(status)) {
	*encryptedLength = cbCipherText + authTagLength;
	} else {
	ret = -1;
	}

	return ret;
	}

	int DecodeOneShot(ptls_aead_algorithm_t aead, wchar_t name, wchar_t *chain_mode, size_t chain_mode_sz,
	BYTE key, ULONG key_length, BYTE iv, ULONG iv_length, BYTE *encrypted,
	ULONG encryptedLength, uint64_t seq, BYTE *authData, ULONG authDataLength,
	ULONG authTagLength, BYTE decrypted, ULONG decryptedLengthMax, ULONG decryptedLength)
	{

	BCRYPT_KEY_HANDLE hKey = NULL;
	BYTE *authTag = encrypted + (encryptedLength - authTagLength);
	BYTE iv_nonce[PTLS_MAX_IV_SIZE];
	BCRYPT_AUTHENTICATED_CIPHER_MODE_INFO bacmi;
	BYTE *ko = NULL;
	ULONG ko_length = 0;
	int ret = 0;

	*decryptedLength = 0;

	if (KeyInit(&hKey, name, chain_mode, chain_mode_sz, key, key_length, &ko, &ko_length) != 0) {
	return -1;
	}

	// Set the auth mode info for AEAD
	BCRYPT_INIT_AUTH_MODE_INFO(bacmi);
	ptls_aead__build_iv(aead, iv_nonce, iv, seq);
	bacmi.pbNonce = iv_nonce;
	bacmi.cbNonce = iv_length;
	bacmi.pbAuthData = authData;
	bacmi.cbAuthData = authDataLength;
	bacmi.pbTag = authTag;
	bacmi.cbTag = authTagLength;
	/* All other fields are set to NULL by the INIT macro. */

	/* If called with a NULL pointer for the data block, we will merely compute the block size. */
	DWORD cbCipherText = 0;
	NTSTATUS status = BCryptDecrypt(hKey, encrypted, encryptedLength - authTagLength, &bacmi, NULL, 0, decrypted,
	decryptedLengthMax, &cbCipherText, 0);

	KeyRelease(&hKey, &ko, &ko_length);

	if (BCRYPT_SUCCESS(status)) {
	*decryptedLength = cbCipherText;
	} else {
	ret = -1;
	}

	return ret;
	}

	int test_oneshot(ptls_aead_algorithm_t aead, wchar_t name, wchar_t *chain_mode, size_t chain_mode_sz)
	{
	BYTE key[32];
	BYTE data[123];
	BYTE iv[PTLS_MAX_IV_SIZE];
	uint64_t nonce;
	BYTE authData[9];
	BYTE encrypted[256];
	ULONG encryptedLength;
	BYTE decrypted[256];
	ULONG decryptedLength;
	ULONG authTagLength = (ULONG)aead->tag_size;
	int ret = 0;

	assert(sizeof(key) >= aead->key_size);
	assert(sizeof(iv) >= aead->iv_size);
	assert(sizeof(data) + authTagLength <= sizeof(encrypted));
	assert(sizeof(decrypted) >= sizeof(encrypted));

	memset(key, 'k', sizeof(key));
	memset(data, 'd', sizeof(data));
	memset(iv, 'n', sizeof(iv));
	nonce = 0;
	memset(authData, 'a', sizeof(authData));

	ret = EncodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size, iv, (ULONG)aead->iv_size, data,
	123, nonce, authData, 9, authTagLength, encrypted, 256, &encryptedLength);

	printf("Encrypt one shot returns %d, l=%d\n", ret, encryptedLength);

	if (ret == 0) {
	ret = DecodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size,
	iv, (ULONG)aead->iv_size, encrypted, encryptedLength, nonce, authData, 9,
	authTagLength, decrypted, 256, &decryptedLength);

	printf("Decrypt one shot returns %d, l=%d\n", ret, decryptedLength);

	if (ret == 0) {
	if (decryptedLength != 123) {
	printf("Wrong length, not %d\n", 123);
	ret = -1;
	} else if (memcmp(data, decrypted, 123) != 0) {
	printf("Data and decrypted don't match\n");
	ret = -1;
	} else {
	printf("One shot matches.\n");
	}
	}
	}

	return ret;
	}

	void delete_test_aead_context(ptls_aead_context_t *ctx)
	{
	if (ctx != NULL) {
	ctx->dispose_crypto(ctx);
	free(ctx);
	}
	}

	ptls_aead_context_t new_test_aead_context(ptls_aead_algorithm_t aead, int is_enc, BYTE key, BYTE iv)
	{
	int ret = 0;
	ptls_aead_context_t ctx = (ptls_aead_context_t )malloc(aead->context_size);

	if (ctx != NULL) {
	memset(ctx, 0, aead->context_size);
	*ctx = (ptls_aead_context_t){aead};
	if (aead->setup_crypto(ctx, is_enc, key, iv) != 0) {
	printf("For %s, setup returns %d\n", aead->name, ret);
	delete_test_aead_context(ctx);
	ctx = NULL;
	}
	} else {
	printf("For %s, memory error during setup\n", aead->name);
	}

	return (ctx);
	}

	int test_decrypt(ptls_aead_algorithm_t aead, wchar_t name, wchar_t *chain_mode, size_t chain_mode_sz)
	{
	BYTE key[32];
	BYTE iv[PTLS_MAX_IV_SIZE];
	BYTE data[123];
	uint64_t nonce;
	BYTE authData[9];
	BYTE encrypted[256];
	ULONG encryptedLength;
	BYTE decrypted[256];
	size_t decryptedLength;
	ULONG authTagLength = (ULONG)aead->tag_size;
	ptls_aead_context_t *ctx = NULL;
	int ret = 0;

	assert(sizeof(key) >= aead->key_size);
	assert(sizeof(iv) >= aead->iv_size);
	assert(sizeof(data) + authTagLength <= sizeof(encrypted));
	assert(sizeof(decrypted) >= sizeof(encrypted));

	memset(key, 'k', sizeof(key));
	memset(iv, 'n', sizeof(iv));
	memset(data, 'd', sizeof(data));
	nonce = 0;
	memset(authData, 'a', sizeof(authData));

	/* Create a decryption context */
	ctx = new_test_aead_context(aead, 0, key, iv);
	if (ctx == NULL) {
	ret = -1;
	}

	/* Do a simple encrypt using one shot bcrypt */
	if (ret == 0) {
	ret = EncodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size,
	iv, (ULONG)aead->iv_size, data, 123, nonce,
	authData, 9, authTagLength, encrypted, 256, &encryptedLength);
	}

	/* Try decrypt with library procedure */
	if (ret == 0) {
	decryptedLength = ctx->do_decrypt(ctx, decrypted, encrypted, encryptedLength, nonce, authData, 9);
	if (decryptedLength >= encryptedLength) {
	printf("For %s, decrypt returns %d\n", aead->name, (int)decryptedLength);
	ret = -1;
	} else if (decryptedLength != 123) {
	printf("For %s, decrypt returns %d instead of %d\n", aead->name, (int)decryptedLength, 123);
	ret = -1;
	} else if (memcmp(data, decrypted, decryptedLength) != 0) {
	printf("For %s, decrypted does not match clear text\n", aead->name);
	ret = -1;
	} else {
	printf("For %s, decrypting test passes.\n", aead->name);
	}
	}

	delete_test_aead_context(ctx);

	return ret;
	}

	int test_encrypt(ptls_aead_algorithm_t aead, wchar_t name, wchar_t *chain_mode, size_t chain_mode_sz)
	{
	BYTE key[32];
	BYTE iv[PTLS_MAX_IV_SIZE];
	BYTE data[123];
	uint64_t nonce;
	BYTE authData[9];
	BYTE encryptedRef[256];
	ULONG encryptedRefLength;
	BYTE encrypted[256];
	size_t encryptedLength;
	ULONG authTagLength = (ULONG)aead->tag_size;
	ptls_aead_context_t *ctx = NULL;
	int ret = 0;

	assert(sizeof(key) >= aead->key_size);
	assert(sizeof(iv) >= aead->iv_size);
	assert(sizeof(data) + authTagLength <= sizeof(encrypted));
	assert(sizeof(data) + authTagLength <= sizeof(encryptedRef));

	memset(key, 'k', sizeof(key));
	memset(iv, 'n', sizeof(iv));
	memset(data, 'd', sizeof(data));
	nonce = 0;
	memset(authData, 'a', sizeof(authData));

	/* Create an encryption context */
	ctx = new_test_aead_context(aead, 1, key, iv);
	if (ctx == NULL) {
	ret = -1;
	}

	/* Do a simple encrypt using one shot bcrypt */
	if (ret == 0) {
	ret = EncodeOneShot(aead, name, chain_mode, chain_mode_sz, key, (ULONG)aead->key_size,
	iv, (ULONG)aead->iv_size, data, 123, nonce,
	authData, 9, authTagLength, encryptedRef, 256, &encryptedRefLength);
	}

	/* Try encrypt with library procedure */
	if (ret == 0) {
	ctx->do_encrypt_init(ctx, nonce, authData, 9);
	encryptedLength = ctx->do_encrypt_update(ctx, encrypted, data, 123);
	encryptedLength += ctx->do_encrypt_final(ctx, &encrypted[encryptedLength]);

	if (encryptedLength != encryptedRefLength) {
	printf("For %s, encrypt returns %d instead of %d\n", aead->name, (int)encryptedLength, encryptedRefLength);
	ret = -1;
	} else if (memcmp(encryptedRef, encrypted, encryptedRefLength) != 0) {
	printf("For %s, encrypted does not match ref\n", aead->name);
	for (ULONG i = 0; i < encryptedRefLength; i++) {
	if (encryptedRef[i] != encrypted[i]) {
	printf("For %s, encrypted[%d] = 0x%02x vs encryptedRef[%d] = 0x%02x\n", aead->name, i, encrypted[i], i,
	encryptedRef[i]);
	break;
	}
	}
	ret = -1;
	} else {
	printf("For %s, encrypting test passes.\n", aead->name);
	}
	}

	delete_test_aead_context(ctx);

	return ret;
	}

	int test_for_size(ptls_aead_algorithm_t aead, wchar_t name, wchar_t *chain_mode, size_t chain_mode_sz)
	{
	BYTE key[32];
	BYTE iv[PTLS_MAX_IV_SIZE];
	uint64_t nonce;
	BYTE authData[9];
	BYTE *data = NULL;
	BYTE *encrypted = NULL;
	BYTE *decrypted = NULL;
	size_t encryptedLength;
	size_t decryptedLength;
	ULONG authTagLength = (ULONG)aead->tag_size;
	ptls_aead_context_t *ctx_e = NULL;
	ptls_aead_context_t *ctx_d = NULL;
	ULONG packet_size[] = {1500, 128, 3, 0};
	ULONG nb_packet_size = (ULONG)(sizeof(packet_size) / sizeof(ULONG));
	int ret = 0;

	assert(sizeof(key) >= aead->key_size);
	assert(sizeof(iv) >= aead->iv_size);

	memset(key, 'k', sizeof(key));
	memset(key, 'n', sizeof(iv));
	nonce = 0;
	memset(authData, 'a', sizeof(authData));

	/* Create the encryption contexts */
	ctx_e = new_test_aead_context(aead, 1, key, iv);
	ctx_d = new_test_aead_context(aead, 0, key, iv);

	if (ctx_e == NULL \|\| ctx_d == NULL) {
	ret = -1;
	}

	/* Test a variety of packet sizes */
	for (ULONG i = 0; ret == 0 && i < nb_packet_size; i++) {
	ULONG data_size = (packet_size[i] > 0) ? packet_size[i] : 128;
	ULONG encrypted_size = packet_size[i] + authTagLength;

	data = (BYTE *)malloc(data_size);
	encrypted = (BYTE *)malloc(encrypted_size);
	decrypted = (BYTE *)malloc(data_size);

	if (data == NULL \|\| encrypted == NULL \|\| decrypted == NULL) {
	printf("For %s: cannot allocate memory for packet size[%d] = %d\n", aead->name, i, packet_size[i]);
	} else {
	memset(data, 'd', data_size);

	ctx_e->do_encrypt_init(ctx_e, nonce, authData, 9);
	encryptedLength = ctx_e->do_encrypt_update(ctx_e, encrypted, data, packet_size[i]);
	encryptedLength += ctx_e->do_encrypt_final(ctx_e, &encrypted[encryptedLength]);
	decryptedLength = ctx_d->do_decrypt(ctx_d, decrypted, encrypted, encryptedLength, nonce, authData, 9);

	if (decryptedLength >= encryptedLength) {
	printf("For %s, decrypt returns %d\n", aead->name, (int)decryptedLength);
	ret = -1;
	} else if (decryptedLength != packet_size[i]) {
	printf("For %s, decrypt returns %d instead of %d\n", aead->name, (int)decryptedLength, packet_size[i]);
	ret = -1;
	} else if (memcmp(data, decrypted, decryptedLength) != 0) {
	printf("For %s, decrypted does not match clear text\n", aead->name);
	ret = -1;
	} else {
	printf("For %s, test packet size[%d] = %d passes.\n", aead->name, i, packet_size[i]);
	}
	}

	if (data != NULL) {
	free(data);
	data = NULL;
	}

	if (encrypted != NULL) {
	free(encrypted);
	encrypted = NULL;
	}

	if (decrypted != NULL) {
	free(decrypted);
	decrypted = NULL;
	}
	}

	delete_test_aead_context(ctx_e);
	delete_test_aead_context(ctx_d);

	return ret;
	}

	int test_one_aead(ptls_aead_algorithm_t aead, wchar_t name, wchar_t *chain_mode, size_t chain_mode_sz)
	{
	int ret = test_oneshot(aead, name, chain_mode, chain_mode_sz);

	printf("For %s, test one shot returns %d\n", aead->name, ret);

	if (ret == 0) {
	ret = test_decrypt(aead, name, chain_mode, chain_mode_sz);

	printf("For %s, test decrypt returns %d\n", aead->name, ret);
	}

	if (ret == 0) {
	ret = test_encrypt(aead, name, chain_mode, chain_mode_sz);

	printf("For %s, test encrypt returns %d\n", aead->name, ret);
	}

	if (ret == 0) {
	ret = test_for_size(aead, name, chain_mode, chain_mode_sz);

	printf("For %s, test packet sizes returns %d\n", aead->name, ret);
	}

	return ret;
	}

	/* Test of cipher functions.
	* The test verifies that a message encode with a bcrypt function can be
	* decoded with a minicrypto function, and vice versa.
	*/

	int test_cipher_one_way(char const name1, char const name2, ptls_cipher_algorithm_t b1, ptls_cipher_algorithm_t b2,
	unsigned int nb_blocks)
	{
	BYTE key[32];
	BYTE nonce[16];
	BYTE data[128];
	BYTE encrypted[128];
	BYTE decrypted[128];
	size_t data_size = b1->block_size * nb_blocks;
	ptls_cipher_context_t *ctx1 = NULL;
	ptls_cipher_context_t *ctx2 = NULL;
	int ret = 0;

	assert(sizeof(key) >= b1->key_size);
	assert(sizeof(data) >= data_size);
	assert(sizeof(nonce) >= b1->iv_size);

	memset(key, 'k', sizeof(key));
	memset(data, 'd', data_size);

	ctx1 = ptls_cipher_new(b1, 1, key);
	ctx2 = ptls_cipher_new(b2, 0, key);

	if (ctx1 == NULL \|\| ctx2 == NULL) {
	ret = -1;
	} else {
	memset(nonce, 0, sizeof(nonce));
	if (ctx1->do_init != NULL) {
	ctx1->do_init(ctx1, nonce);
	}

	if (ctx2->do_init != NULL) {
	ctx2->do_init(ctx2, nonce);
	}

	ctx1->do_transform(ctx1, encrypted, data, data_size);
	ctx2->do_transform(ctx2, decrypted, encrypted, data_size);

	if (memcmp(data, decrypted, data_size) != 0) {
	printf("For %s -> %s, decrypted does not match clear text\n", name1, name2);
	ret = -1;
	} else {
	printf("For %s -> %s, test passes.\n", name1, name2);
	}
	}

	if (ctx1 != NULL) {
	ptls_cipher_free(ctx1);
	}

	if (ctx2 != NULL) {
	ptls_cipher_free(ctx2);
	}

	return ret;
	}

	int test_cipher_pair(char const name1, char const name2, ptls_cipher_algorithm_t b1, ptls_cipher_algorithm_t b2,
	unsigned int nb_blocks)
	{
	int ret = test_cipher_one_way(name1, name2, b1, b2, nb_blocks);

	if (ret == 0) {
	ret = test_cipher_one_way(name2, name1, b2, b1, nb_blocks);
	}

	return ret;
	}

	/* Test of the hash functions
	*/

	int test_hash_calc(char const name1, char const name2, ptls_hash_algorithm_t h1, ptls_hash_algorithm_t h2)
	{
	BYTE data[123];
	BYTE tag1[128];
	BYTE tag2[128];
	ptls_hash_context_t *ctx1 = NULL;
	ptls_hash_context_t *ctx2 = NULL;
	int ret = 0;

	assert(sizeof(tag1) >= h1->digest_size);
	assert(sizeof(tag2) >= h2->digest_size);
	assert(h1->digest_size == h2->digest_size);

	memset(data, 'd', sizeof(data));
	memset(tag1, '1', sizeof(tag1));
	memset(tag2, '2', sizeof(tag2));

	if (h1->digest_size != h2->digest_size) {
	ret = -1;
	}
	if (ret == 0) {
	ret = ptls_calc_hash(h1, tag1, data, sizeof(data));
	}

	if (ret == 0) {
	ret = ptls_calc_hash(h2, tag2, data, sizeof(data));
	}

	if (ret == 0){
	if (memcmp(tag1, tag2, h1->digest_size) != 0) {
	printf("For %s -> %s, hash1 does not match hash2\n", name1, name2);
	ret = -1;
	} else {
	printf("For %s -> %s, hash test passes.\n", name1, name2);
	}
	}

	return ret;
	}

	/* Minimal test program for the bcrypt functions.
	* Need to add tests for the SHA256 and SHA384 implementations.
	*/

	int main()
	{
	int ret = 0;

	ret \|= test_cipher_pair("bcrypt aes128ecb", "minicrypto aes128ecb", &ptls_bcrypt_aes128ecb, &ptls_minicrypto_aes128ecb, 1);
	ret \|= test_cipher_pair("bcrypt aes256ecb", "minicrypto aes256ecb", &ptls_bcrypt_aes256ecb, &ptls_minicrypto_aes256ecb, 1);

	ret \|= test_cipher_pair("bcrypt aes128ctr", "minicrypto aes128ctr", &ptls_bcrypt_aes128ctr, &ptls_minicrypto_aes128ctr, 4);
	ret \|= test_cipher_pair("bcrypt aes256ctr", "minicrypto aes256ctr", &ptls_bcrypt_aes256ctr, &ptls_minicrypto_aes256ctr, 4);

	ret \|= test_one_aead(&ptls_bcrypt_aes128gcm, BCRYPT_AES_ALGORITHM, BCRYPT_CHAIN_MODE_GCM, sizeof(BCRYPT_CHAIN_MODE_GCM));

	ret \|= test_one_aead(&ptls_bcrypt_aes256gcm, BCRYPT_AES_ALGORITHM, BCRYPT_CHAIN_MODE_GCM, sizeof(BCRYPT_CHAIN_MODE_GCM));

	ret \|= test_hash_calc("bcrypt sha256", "minicrypto sha256", &ptls_bcrypt_sha256, &ptls_minicrypto_sha256);
	ret \|= test_hash_calc("bcrypt sha384", "minicrypto sha384", &ptls_bcrypt_sha384, &ptls_minicrypto_sha384);

	exit(ret);
	}