/**
 * \file ctr_drbg.c
 * \brief CTR_DRBG AES-256 implementation according to SP800-90A.
 * \author Dmytro Podgornyi (d.podgornyi@samsung.com)
 * \version 0.1
 * \date Created Sep 04, 2013
 * \par In Samsung Ukraine R&D Center (SURC) under a contract between
 * \par LLC "Samsung Electronics Ukraine Company" (Kiev, Ukraine) and
 * \par "Samsung Elecrtronics Co", Ltd (Seoul, Republic of Korea)
 * \par Copyright: (c) Samsung Electronics Co, Ltd 2012. All rights reserved.
 **/

#include "drbg/ctr_drbg.h"

#include <stdint.h>
#include <string.h>
#ifdef CTR_DRBG_DYNAMIC_BUFFERS
#include "stdlib.h"
#endif

#include "openssl/aes/aes.h"

/* Allocator */
#define DRBG_malloc malloc
#define DRBG_free free

/** @def UPDATE_BUFSIZE
 *  At least seedlen in length but aligned to outlen boundary. Used for
 *  storing key + v obtained trough block cipher.
 */
#if (CTR_DRBG_SEEDLEN_BYTES % CTR_DRBG_OUTLEN_BYTES == 0)
#define UPDATE_BUFSIZE CTR_DRBG_SEEDLEN_BYTES
#else
#define UPDATE_BUFSIZE (CTR_DRBG_SEEDLEN_BYTES + CTR_DRBG_OUTLEN_BYTES - \
			CTR_DRBG_SEEDLEN_BYTES % CTR_DRBG_OUTLEN_BYTES)
#endif

/** @def ALLOCATE_BUFFER
 *  Allocate buffer in stack or in heap depending on configuration
 */
#ifdef CTR_DRBG_DYNAMIC_BUFFERS
#define ALLOCATE_BUFFER(buf, size) \
	const size_t __local_size_##buf = size; \
	uint8_t * const buf = DRBG_malloc(size);
#else
#define ALLOCATE_BUFFER(buf, size) \
	uint8_t __local_##buf[size]; \
	uint8_t * const buf = __local_##buf;
#endif

/** @def ZEROIZE_BUFFER
 *  Fill buf with zero. The buf must be allocated with ALLOCATE_BUFFER
 */
#ifdef CTR_DRBG_DYNAMIC_BUFFERS
#define ZEROIZE_BUFFER(buf) \
	do { memset(buf, 0, __local_size_##buf); } while (0)
#else
#define ZEROIZE_BUFFER(buf) \
	do { memset(buf, 0, sizeof(__local_##buf)); } while (0)
#endif

/** @def RELEASE_BUFFER
 *  Zeroize and free buffer if needed. buf must be allocated with
 *  ALLOCATE_BUFFER
 */
#ifdef CTR_DRBG_DYNAMIC_BUFFERS
#define RELEASE_BUFFER(buf) \
	do \
	{ \
		ZEROIZE_BUFFER(buf); \
		DRBG_free(buf); \
	} \
	while (0)
#else
#define RELEASE_BUFFER(buf) ZEROIZE_BUFFER(buf)
#endif

/** @def RELEASE_STACK_BUFFER
 *  Fill buf with zero if needed. buf must be declared in stack
 */
#define RELEASE_STACK_BUFFER(buf) do { memset(buf, 0, sizeof(buf)); } while (0)

/** @def ERROR_ON
 *  Return error code if condition is met. 
 */
#ifdef ERROR_ON
#undef ERROR_ON
#endif
#define ERROR_ON(cond, err_code) \
	if (cond) \
	{ \
		err = err_code; \
		goto cleanup; \
	}


/**
 * @brief	V = (V + 1) mod 2^outlen
 * @param	[in,out] v The current value of V
 */
static void v_increase(uint8_t *v)
{
	int i;
	for (i = CTR_DRBG_OUTLEN_BYTES - 1; i >= 0; i--)
	{
		if (++v[i])
		{
			break;
		}
	}
}

#ifndef CTR_DRBG_WITHOUT_DF

/**
 * @brief	Put uint32_t as a bitstring to out buffer
 * @param	[in] val 32-bit value to put
 * @param	[out] out Output buffer
 */
static void set_uint32(uint32_t val, uint8_t *out)
{
	*out++ = (val >> 24) & 0xff;
	*out++ = (val >> 16) & 0xff;
	*out++ = (val >> 8)  & 0xff;
	*out = val & 0xff;
}

/**
 * @brief	A basic encryption operation that uses the selected block
 * 		cipher algorithm.
 * @param	[in] key The key to be used for the block cipher operation
 * @param	[in] data The data to be operated upon. Note that the length
 * 		of data must be a multiple of outlen. This is guaranteed by
 * 		Block_Cipher_df process
 * @param	[in] data_len Length of data
 * @param	[out] output_block The result to be returned from the BCC
 * 		operation, must be outlen in length
 */
static void BCC(AES_KEY *key, const uint8_t *data, uint32_t data_len,
	uint8_t *output_block)
{
	int i;
	int j;

	/* output_block is used as chaining_value in SP800-90A */
	/* input_block and output_block can overlap for AES_encrypt
	 * so use output_block as input_block too */
	memset(output_block, 0, CTR_DRBG_OUTLEN_BYTES);

	i = 0;
	while (i < data_len)
	{
		for (j = 0; j < CTR_DRBG_OUTLEN_BYTES; j++)
		{
			output_block[j] ^= data[i + j];
		}
		AES_encrypt(output_block, output_block, key);
		i += CTR_DRBG_OUTLEN_BYTES;
	}
}

/* key used in derivation function */
static const uint8_t ctr_drbg_df_key[] = {
	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
	0x08, 0x09, 0x0a, 0x0b,	0x0c, 0x0d, 0x0e, 0x0f,
	0x10, 0x11, 0x12, 0x13, 0x14, 0x15, 0x16, 0x17,
	0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f
};

/**
 * @brief	The derivation function is used by the CTR_DRBG
 * @param	[in] input_string The string to be operated on. This string
 * 		shall be a multiple of 8 bits.
 * @param	[in] input_string_len Length of input_string
 * @param	[out] out Output buffer
 * @param	[in] bits Number of requested bits
 * @return	DRBG_NO_ERROR on success, negative error code otherwise
 */
static int Block_Cipher_df(const uint8_t *input_string,
	uint32_t input_string_len,
	uint8_t *out, uint32_t bits)
{
	int i;
	int err = DRBG_NO_ERROR;
	ALLOCATE_BUFFER(IV_S_buf, input_string_len + 9 + 2 * CTR_DRBG_OUTLEN_BYTES)
	ALLOCATE_BUFFER(temp, UPDATE_BUFSIZE)
	uint8_t *X;
	uint8_t * const S = IV_S_buf + CTR_DRBG_OUTLEN_BYTES;
	const uint32_t N = (bits + 7) / 8;
	const uint32_t L = input_string_len;
	uint32_t len;
	AES_KEY aes_key;

	if (bits == 0)
	{
		goto cleanup;
	}

	ZEROIZE_BUFFER(IV_S_buf);
	ZEROIZE_BUFFER(temp);

	/* create S = (L || N || input_string || 0x80) */
	set_uint32(L, S);
	set_uint32(N, S + 4);
	memcpy(S + 8, input_string, input_string_len);
	len = 4 + 4 + input_string_len + 1;
	S[len - 1] = 0x80;
	if (len % CTR_DRBG_OUTLEN_BYTES != 0)
	{
		/* len must be multiple of outlen, S will be padded with 0 */
		len += CTR_DRBG_OUTLEN_BYTES - len % CTR_DRBG_OUTLEN_BYTES;
	}
	/* add IV_LEN to len */
	len += CTR_DRBG_OUTLEN_BYTES;

	memset(&aes_key, 0, sizeof(aes_key));
	err = AES_set_encrypt_key(ctr_drbg_df_key, CTR_DRBG_KEYLEN, &aes_key);
	if (err != 0)
	{
		err = DRBG_E_AES_KEY;
		goto cleanup;
	}

	i = 0;
	while (i < UPDATE_BUFSIZE / CTR_DRBG_OUTLEN_BYTES)
	{
		/* set IV padded with zeros */
		set_uint32((uint32_t)i, IV_S_buf);
		BCC(&aes_key, IV_S_buf, len, &temp[i * CTR_DRBG_OUTLEN_BYTES]);
		++i;
	}

	/* first keylen bits of temp represent K */
	err = AES_set_encrypt_key(temp, CTR_DRBG_KEYLEN, &aes_key);
	if (err != 0)
	{
		err = DRBG_E_AES_KEY;
		goto cleanup;
	}

	X = temp + CTR_DRBG_KEYLEN_BYTES;
	i = 0;
	while (i < N)
	{
		if (N - i <= CTR_DRBG_OUTLEN_BYTES)
		{
			/* this is the last interation */
			AES_encrypt(X, temp, &aes_key);
			/* XXX: bits % 8 == 0 in current implementation */
			memcpy(&out[i], temp, N - i);
		}
		else
		{
			AES_encrypt(X, &out[i], &aes_key);
			X = &out[i];
		}
		i += CTR_DRBG_OUTLEN_BYTES;
	}

	err = DRBG_NO_ERROR;

cleanup:
	RELEASE_BUFFER(temp);
	RELEASE_BUFFER(IV_S_buf);

	return err;
}

#endif /* CTR_DRBG_WITHOUT_DF */

/**
 * @brief	Update internal state
 * @param	[in] provided_data The data to be used, must be exactly
 * 		CTR_DRBG_SEEDLEN bits in length
 * @param	[in,out] key The current value of Key
 * @param	[in,out] v The current value of V
 * @return	DRBG_NO_ERROR on success, negative error code otherwise
 */
static int CTR_DRBG_Update(const uint8_t *provided_data, uint8_t *key,
	uint8_t *v)
{
	int i;
	int err = DRBG_NO_ERROR;
	ALLOCATE_BUFFER(temp, UPDATE_BUFSIZE)
	AES_KEY aes_key;

	ZEROIZE_BUFFER(temp);

	memset(&aes_key, 0, sizeof(aes_key));
	if (AES_set_encrypt_key(key, CTR_DRBG_KEYLEN, &aes_key) != 0)
	{
		err = DRBG_E_AES_KEY;
		goto cleanup;
	}

	i = 0;
	while (i < CTR_DRBG_SEEDLEN_BYTES)
	{
		v_increase(v);
		AES_encrypt(v, &temp[i], &aes_key);
		i += CTR_DRBG_OUTLEN_BYTES;
	}

	for (i = 0; i < CTR_DRBG_SEEDLEN_BYTES; i++)
	{
		temp[i] ^= provided_data[i];
	}

	memcpy(key, temp, CTR_DRBG_KEYLEN_BYTES);
	memcpy(v, temp + CTR_DRBG_KEYLEN_BYTES, CTR_DRBG_OUTLEN_BYTES);

	err = DRBG_NO_ERROR;

cleanup:
	RELEASE_BUFFER(temp);

	return err;
}

/**
 * @brief	Instantiation when full entropy input is available and
 * 		a derivation function is not used
 * @note	Without nonce this function repeats functionality of reseed
 * @param	[out] ctx CTR_DRBG context to initialize
 * @param	[in] entropy_input The string of bits obtained from the
 * 		source of entropy input
 * @param	[in] entropy_input_len Length of entropy_input
 * @param	[in] nonce A bitstring to provide a security cushion to block
 * 		certain attacks
 * @param	[in] nonce_len Length of nonce
 * @param	[in] personalization_string The personalization string
 * 		received from the consuming application, may be zero
 * @param	[in] personalization_string_len Length of personalization_string
 * @param	[in] security_strength The security strength for the
 * 		instantiation, optional for CTR_DRBG
 * @return	DRBG_NO_ERROR on success, negative error code otherwise
 */
int CTR_DRBG_Instantiate(CTR_DRBG_CTX *ctx,
	const uint8_t *entropy_input,
	uint32_t entropy_input_len,
	const uint8_t *nonce,
	uint32_t nonce_len,
	const uint8_t *personalization_string,
	uint32_t personalization_string_len,
	int security_strength)
{
	int err = DRBG_NO_ERROR;
	ALLOCATE_BUFFER(seed_material, CTR_DRBG_SEEDLEN_BYTES)
#ifdef CTR_DRBG_WITHOUT_DF
	int i;
#else  /* CTR_DRBG_WITHOUT_DF */
	ALLOCATE_BUFFER(seed_material_df, entropy_input_len +
			nonce_len + personalization_string_len)
	uint32_t seed_material_df_len = 0;
#endif /* CTR_DRBG_WITHOUT_DF */

    if ((entropy_input_len + nonce_len) < entropy_input_len
        || (entropy_input_len + personalization_string_len) < entropy_input_len
        || (personalization_string_len + nonce_len) < personalization_string_len
        || (entropy_input_len + nonce_len + personalization_string_len) < (entropy_input_len + nonce_len)) {
        err = DRBG_NO_ERROR;
        goto cleanup;
    }

	ERROR_ON(!ctx || !entropy_input, DRBG_E_WRONG_DATA);
	ERROR_ON(security_strength > CTR_DRBG_MAX_SECURITY_STRENGTH,
		DRBG_E_TOO_BIG);

	/* also set key and v to zero */
	memset(ctx, 0, sizeof(*ctx));
	ctx->security_strength = security_strength;

	ERROR_ON(entropy_input_len < CTR_DRBG_MIN_ENTROPY, DRBG_E_TOO_LOW);
	ERROR_ON(entropy_input_len > CTR_DRBG_MAX_ENTROPY, DRBG_E_TOO_BIG);
	ERROR_ON(personalization_string &&
		personalization_string_len > CTR_DRBG_MAX_PERSONAL,
		DRBG_E_TOO_BIG);
	ERROR_ON(nonce && nonce_len > CTR_DRBG_MAX_NONCE, DRBG_E_TOO_BIG);

	ZEROIZE_BUFFER(seed_material);

#ifdef CTR_DRBG_WITHOUT_DF

	memcpy(seed_material, entropy_input, entropy_input_len);
	if (personalization_string != NULL)
	{
		for (i = 0; i < personalization_string_len; i++)
		{
			seed_material[i] ^= personalization_string[i];
		}
	}

#else  /* CTR_DRBG_WITHOUT_DF */

	/* ZEROIZE_BUFFER(seed_material_df); */
	seed_material_df_len = entropy_input_len;
	memcpy(seed_material_df, entropy_input, entropy_input_len);

	if (nonce != NULL)
	{
		memcpy(seed_material_df + seed_material_df_len,
			nonce, nonce_len);
		seed_material_df_len += nonce_len;
	}

	if (personalization_string != NULL)
	{
		memcpy(seed_material_df + seed_material_df_len,
			personalization_string,
			personalization_string_len);
		seed_material_df_len += personalization_string_len;
	}

	err = Block_Cipher_df(seed_material_df, seed_material_df_len,
			seed_material, CTR_DRBG_SEEDLEN);
	if (err != DRBG_NO_ERROR)
	{
		goto cleanup;
	}

#endif /* CTR_DRBG_WITHOUT_DF */

	err = CTR_DRBG_Update(seed_material, ctx->key, ctx->v);
	if (err != DRBG_NO_ERROR)
	{
		goto cleanup;
	}

	ctx->reseed_counter = 1;
	err = DRBG_NO_ERROR;

cleanup:
	RELEASE_BUFFER(seed_material);
#ifndef CTR_DRBG_WITHOUT_DF
	RELEASE_BUFFER(seed_material_df);
#endif /* CTR_DRBG_WITHOUT_DF */

	return err;
}

/**
 * @brief	Reseeding When Full Entropy Input is Available, and
 * 		a Derivation Function is Not Used
 * @param	[in,out] ctx CTR_DRBG context
 * @param	[in] entropy_input The string of bits obtained from the
 * 		source of entropy input
 * @param	[in] entropy_input_len Length of entropy_input
 * @param	[in] additional_input The additional input string received
 * 		from the consuming application, may be NULL
 * @param	[in] additional_input_len Length of additional_input
 * @return	DRBG_NO_ERROR on success, negative error code otherwise
 */
int CTR_DRBG_Reseed(CTR_DRBG_CTX *ctx,
	const uint8_t *entropy_input,
	uint32_t entropy_input_len,
	const uint8_t *additional_input,
	uint32_t additional_input_len)
{
	int err = DRBG_NO_ERROR;
	ALLOCATE_BUFFER(seed_material, CTR_DRBG_SEEDLEN_BYTES)
#ifdef CTR_DRBG_WITHOUT_DF
	int i;
#else  /* CTR_DRBG_WITHOUT_DF */
	ALLOCATE_BUFFER(seed_material_df, entropy_input_len +
			additional_input_len)
	uint32_t seed_material_df_len;
#endif /* CTR_DRBG_WITHOUT_DF */
    if ((entropy_input_len + additional_input_len) < entropy_input_len) {
        err = DRBG_NO_ERROR;
        goto cleanup;
    }

	ERROR_ON(!ctx || !entropy_input, DRBG_E_WRONG_DATA);
	ERROR_ON(entropy_input_len < CTR_DRBG_MIN_ENTROPY, DRBG_E_TOO_LOW);
	ERROR_ON(entropy_input_len > CTR_DRBG_MAX_ENTROPY, DRBG_E_TOO_BIG);
	ERROR_ON(additional_input &&
		additional_input_len > CTR_DRBG_MAX_ADDITIONAL,
		DRBG_E_TOO_BIG);

	ZEROIZE_BUFFER(seed_material);

#ifdef CTR_DRBG_WITHOUT_DF

	memcpy(seed_material, entropy_input, entropy_input_len);
	if (additional_input != NULL && additional_input_len > 0)
	{
		for (i = 0; i < additional_input_len; i++)
		{
			seed_material[i] ^= additional_input[i];
		}
	}

#else  /* CTR_DRBG_WITHOUT_DF */

	/* ZEROIZE_BUFFER(seed_material_df); */
	seed_material_df_len = entropy_input_len;
	memcpy(seed_material_df, entropy_input, entropy_input_len);
	if (additional_input != NULL)
	{
		memcpy(seed_material_df + seed_material_df_len,
			additional_input, additional_input_len);
		seed_material_df_len += additional_input_len;
	}

	err = Block_Cipher_df(seed_material_df, seed_material_df_len,
				seed_material, CTR_DRBG_SEEDLEN);
	if (err != DRBG_NO_ERROR)
	{
		goto cleanup;
	}

#endif /* CTR_DRBG_WITHOUT_DF */

	err = CTR_DRBG_Update(seed_material, ctx->key, ctx->v);
	if (err != DRBG_NO_ERROR)
	{
		goto cleanup;
	}

	ctx->reseed_counter = 1;
	err = DRBG_NO_ERROR;

cleanup:
	RELEASE_BUFFER(seed_material);
#ifndef CTR_DRBG_WITHOUT_DF
	RELEASE_BUFFER(seed_material_df);
#endif /* CTR_DRBG_WITHOUT_DF */

	return err;
}

/**
 * @brief	Generating Pseudorandom Bits When a Derivation Function is
 * 		Not Used for the DRBG Implementation
 * @note	Entropy source is out of scope thereby prediction resistance
 * 		can't be implemented here
 * @param	[in,out] ctx CTR_DRBG context
 * @param	[out] out Output buffer
 * @param	[in] bits The number of pseudorandom bits to be returned
 * @param	[in] additional_input The additional input string received
 * 		from the consuming application, may be NULL
 * @param	[in] additional_input_len Length of additional_input
 * @return	DRBG_NO_ERROR on success, DRBG_E_RESEED_NEEDED if reseed
 * 		function must be called before, negative error code otherwise
 */
int CTR_DRBG_Generate(CTR_DRBG_CTX *ctx,
	uint8_t *out, uint32_t bits,
	const uint8_t *additional_input,
	uint32_t additional_input_len)
{
	const int bytes = (bits + 7) / 8;
	int i;
	uint8_t mod;
	int err = DRBG_NO_ERROR;
	uint8_t temp[CTR_DRBG_OUTLEN_BYTES] = {0};
	ALLOCATE_BUFFER(seed_material, CTR_DRBG_SEEDLEN_BYTES)
	AES_KEY aes_key;

	ERROR_ON(!ctx || !out, DRBG_E_WRONG_DATA);
	ERROR_ON(ctx->reseed_counter > CTR_DRBG_RESEED_INTERVAL,
		DRBG_E_RESEED_NEEDED);
	ERROR_ON(additional_input &&
		additional_input_len > CTR_DRBG_MAX_ADDITIONAL,
		DRBG_E_TOO_BIG);
	ERROR_ON(bits > CTR_DRBG_MAX_BITS, DRBG_E_TOO_BIG);

	if (bits == 0)
	{
		goto cleanup;
	}

	ZEROIZE_BUFFER(seed_material);

	if (additional_input != NULL && additional_input_len > 0)
	{
#ifdef CTR_DRBG_WITHOUT_DF
		memcpy(seed_material, additional_input, additional_input_len);
#else  /* CTR_DRBG_WITHOUT_DF */
		err = Block_Cipher_df(additional_input, additional_input_len,
					seed_material, CTR_DRBG_SEEDLEN);
		if (err != DRBG_NO_ERROR)
		{
			goto cleanup;
		}
#endif /* CTR_DRBG_WITHOUT_DF */

		err = CTR_DRBG_Update(seed_material, ctx->key, ctx->v);
		if (err != DRBG_NO_ERROR)
		{
			goto cleanup;
		}
	}

	memset(&aes_key, 0, sizeof(aes_key));
	err = AES_set_encrypt_key(ctx->key, CTR_DRBG_KEYLEN, &aes_key);
	if (err != 0)
	{
		err = DRBG_E_AES_KEY;
		goto cleanup;
	}

	i = 0;
	while (i < bytes)
	{
		v_increase(ctx->v);
		if (bytes - i <= CTR_DRBG_OUTLEN_BYTES)
		{
			/* this is the last iteration */
			AES_encrypt(ctx->v, temp, &aes_key);
			memcpy(&out[i], temp, bytes - i);

			mod = bits % 8;
			if (mod > 0)
			{
				mod = 1 << (8 - mod);
				out[bytes - 1] &= ~(mod - (uint8_t)1);
			}
		}
		else
		{
			AES_encrypt(ctx->v, &out[i], &aes_key);
		}
		i += CTR_DRBG_OUTLEN_BYTES;
	}

	err = CTR_DRBG_Update(seed_material, ctx->key, ctx->v);
	++ctx->reseed_counter;
	if (err != DRBG_NO_ERROR)
	{
		memset(out, 0, bytes);
		goto cleanup;
	}

	err = DRBG_NO_ERROR;

cleanup:
	RELEASE_BUFFER(seed_material);
	RELEASE_STACK_BUFFER(temp);

	return err;
}