/*
* Copyright (c) 2024 Vaughn Nugent
*
* Package: noscrypt
* File: impl/bcrypt.c
*
* This library is free software; you can redistribute it and/or
* modify it under the terms of the GNU Lesser General Public License
* as published by the Free Software Foundation; either version 2.1
* of the License, or  (at your option) any later version.
*
* This library is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
* Lesser General Public License for more details.
*
* You should have received a copy of the GNU Lesser General Public License
* along with noscrypt. If not, see http://www.gnu.org/licenses/.
*/


/*
*	This file provides as many fallback implementations on Windows plaforms
*	as possible using the bcrypt library. This file should be included behind
*	other libarry implementations, as it is a fallback.
*/

#ifdef _NC_IS_WINDOWS

#define WIN32_LEAN_AND_MEAN
#include <Windows.h>
#include <bcrypt.h>

#include "../../platform.h"
#include "../nc-util.h"

#define IF_BC_FAIL(x) if(!BCRYPT_SUCCESS(x))

struct _bcrypt_ctx
{
	BCRYPT_ALG_HANDLE hAlg;
	BCRYPT_HASH_HANDLE hHash;
};

_IMPLSTB NTSTATUS _bcInitSha256(struct _bcrypt_ctx* ctx, DWORD flags)
{
	NTSTATUS result;

	result = BCryptOpenAlgorithmProvider(
		&ctx->hAlg, 
		BCRYPT_SHA256_ALGORITHM, 
		NULL, 
		flags
	);

	/*
	*  If operation failed, ensure the algorithm handle is null
	* to make free code easier to cleanup
	*/
	if (!BCRYPT_SUCCESS(result))
	{
		ctx->hAlg = NULL;
	}

	return result;
}

_IMPLSTB NTSTATUS _bcCreateHmac(struct _bcrypt_ctx* ctx, const cspan_t* key)
{
	/*
	 * NOTE: 
	 * I am not explicitly managing the hash object buffer. By setting 
	 * the hash object to NULL, and length to 0, the buffer will be
	 * managed by the bcrypt library.
	 * 
	 * See: https://learn.microsoft.com/en-us/windows/win32/api/bcrypt/nf-bcrypt-bcryptcreatehash
	 */

	return BCryptCreateHash(
		ctx->hAlg, 
		&ctx->hHash, 
		NULL, 
		0, 
		(uint8_t*)key->data, 
		key->size, 
		BCRYPT_HASH_REUSABLE_FLAG	/* Enable reusable for expand function */
	);
}

_IMPLSTB NTSTATUS _bcCreate(struct _bcrypt_ctx* ctx)
{
	cspan_t key;
	
	/* Zero out key span for 0 size and NULL data ptr */
	SecureZeroMemory(&key, sizeof(cspan_t));

	return _bcCreateHmac(ctx, &key);
}

_IMPLSTB NTSTATUS _bcHashDataRaw(const struct _bcrypt_ctx* ctx, const uint8_t* data, uint64_t len)
{
	return BCryptHashData(ctx->hHash, (uint8_t*)data, len, 0);
}

_IMPLSTB NTSTATUS _bcHashData(const struct _bcrypt_ctx* ctx, const cspan_t* data)
{
	return _bcHashDataRaw(ctx, data->data, data->size);
}

_IMPLSTB NTSTATUS _bcFinishHash(const struct _bcrypt_ctx* ctx, sha256_t digestOut32)
{
	return BCryptFinishHash(ctx->hHash, digestOut32, sizeof(sha256_t), 0);
}

_IMPLSTB void _bcDestroyCtx(struct _bcrypt_ctx* ctx)
{
	/* Free the hash memory if it was allocated */
	if(ctx->hHash) BCryptDestroyHash(ctx->hHash);
	
	/* Close the algorithm provider */
	if (ctx->hAlg) BCryptCloseAlgorithmProvider(ctx->hAlg, 0);

	ctx->hAlg = NULL;
	ctx->hHash = NULL;
}

#ifndef _IMPL_SECURE_ZERO_MEMSET
	/*
	* On Windows, we can use SecureZeroMemory
	* as platform zeroing function.
	*
	* NOTE:
	* SecureZeroMemory2 uses volitle function argument
	* pointers, which is a contested mehtod of compiler
	* optimization prevention. GNU seems to oppose this method
	*
	* https://learn.microsoft.com/en-us/windows/win32/memory/winbase-securezeromemory2
	*/
	#define _IMPL_SECURE_ZERO_MEMSET SecureZeroMemory
#endif /* !_IMPL_SECURE_ZERO_MEMSET */

/*
* Provide win32 fallback for sha256 digest if needed
*/

#ifndef _IMPL_CRYPTO_SHA256_DIGEST
	
	/* Export function fallack */
	#define _IMPL_CRYPTO_SHA256_DIGEST			_bcrypt_sha256_digest	

	_IMPLSTB cstatus_t _bcrypt_sha256_digest(const cspan_t* data, sha256_t digestOut32)
	{
		cstatus_t result;
		struct _bcrypt_ctx ctx;

		result = CSTATUS_FAIL;	/* Start in fail state */

		IF_BC_FAIL(_bcInitSha256(&ctx, 0)) goto Exit;

		IF_BC_FAIL(_bcCreate(&ctx)) goto Exit;

		IF_BC_FAIL(_bcHashData(&ctx, data)) goto Exit;
	
		IF_BC_FAIL(_bcFinishHash(&ctx, digestOut32)) goto Exit;

		result = CSTATUS_OK;	/* Hash operation completed, so set success */

	Exit:

		_bcDestroyCtx(&ctx);

		return result;
	}

#endif /* !_IMPL_CRYPTO_SHA256_DIGEST */

#ifndef _IMPL_CRYPTO_SHA256_HMAC

	/* Export function */
	#define _IMPL_CRYPTO_SHA256_HMAC				_bcrypt_hmac_sha256	
	
	_IMPLSTB cstatus_t _bcrypt_hmac_sha256(const cspan_t* key, const cspan_t* data, sha256_t hmacOut32)
	{
		cstatus_t result;
		struct _bcrypt_ctx ctx;

		result = CSTATUS_FAIL;		/* Start in fail state */

		/* Init context with hmac flag set */
		IF_BC_FAIL(_bcInitSha256(&ctx, BCRYPT_ALG_HANDLE_HMAC_FLAG)) goto Exit;

		IF_BC_FAIL(_bcCreateHmac(&ctx, key)) goto Exit;

		IF_BC_FAIL(_bcHashData(&ctx, data)) goto Exit;

		IF_BC_FAIL(_bcFinishHash(&ctx, hmacOut32)) goto Exit;

		result = CSTATUS_OK;	/* HMAC operation completed, so set success */

	Exit:
		
		_bcDestroyCtx(&ctx);

		return result;
	}

#endif /* !_IMPL_CRYPTO_SHA256_HMAC */

/*
* Provide a fallback HKDF expand function using the
* HMAC function as a base.
*/

#ifndef _IMPL_CRYPTO_SHA256_HKDF_EXPAND

	#define _IMPL_CRYPTO_SHA256_HKDF_EXPAND		_fallbackHkdfExpand

	/* Include string for memmove */
	#include <string.h>

	static void ncWriteSpanS(span_t* span, uint64_t offset, const uint8_t* data, uint64_t size)
	{
		DEBUG_ASSERT2(span != NULL, "Expected span to be non-null")
		DEBUG_ASSERT2(data != NULL, "Expected data to be non-null")
		DEBUG_ASSERT2(offset + size <= span->size, "Expected offset + size to be less than span size")

		/* Copy data to span */
		memmove(span->data + offset, data, size);
	}

	STATIC_ASSERT(HKDF_IN_BUF_SIZE > SHA256_DIGEST_SIZE, "HDK Buffer must be at least the size of the underlying hashing alg output")

	/*
	* The following functions implements the HKDF expand function using an existing 
	* HMAC function. This is a fallback implementation for Windows platforms at the moment.
	* 
	* This follows the guidence from RFC 5869: https://tools.ietf.org/html/rfc5869
	*/

	#define _BC_MIN(a, b) (a < b ? a : b)

	_IMPLSTB cstatus_t _fallbackHkdfExpand(const cspan_t* prk, const cspan_t* info, span_t* okm)
	{
		cstatus_t result;
		struct _bcrypt_ctx ctx;

		uint8_t counter;
		uint64_t tLen, okmOffset;
		uint8_t t[HKDF_IN_BUF_SIZE];

		_IMPL_SECURE_ZERO_MEMSET(t, sizeof(t));
		
		tLen = 0;				/* T(0) is an empty string(zero length) */
		okmOffset = 0;
		result = CSTATUS_FAIL;	/* Start in fail state */

		/* Init context with hmac flag set, it will be reused */
		IF_BC_FAIL(_bcInitSha256(&ctx, BCRYPT_ALG_HANDLE_HMAC_FLAG)) goto Exit;

		/* Set hmac key to the prk, alg is set to reusable */
		IF_BC_FAIL(_bcCreateHmac(&ctx, prk)) goto Exit;

		/* Compute T(N) = HMAC(prk, T(n-1) | info | n) */
		for (counter = 1; okmOffset < okm->size; counter++)
		{
			IF_BC_FAIL(_bcHashDataRaw(&ctx, t, tLen)) goto Exit;
			IF_BC_FAIL(_bcHashData(&ctx, info)) goto Exit;
			IF_BC_FAIL(_bcHashDataRaw(&ctx, &counter, sizeof(counter))) goto Exit;

			/* Write current hash state to t buffer */
			IF_BC_FAIL(_bcFinishHash(&ctx, t)) goto Exit;

			/* Set the length of the current hash state */
			tLen = _BC_MIN(okm->size - okmOffset, SHA256_DIGEST_SIZE);

			DEBUG_ASSERT(tLen <= sizeof(t));
			DEBUG_ASSERT((tLen + okmOffset) < okm->size);

			/* write the T buffer back to okm */
			ncWriteSpanS(okm, okmOffset, t, tLen);

			/* shift base okm pointer by T */
			okmOffset += tLen;
		}		

		result = CSTATUS_OK;	/* HMAC operation completed, so set success */

	Exit:

		_bcDestroyCtx(&ctx);

		return result;
	}

#endif /* !_IMPL_CRYPTO_SHA256_HKDF_EXPAND */

#endif /* _NC_IS_WINDOWS */