refactor: Finish support and testing for mbedtls

8 files changed, 347 insertions, 168 deletions

diff --git a/src/crypto/hkdf.c b/src/crypto/hkdf.c
new file mode 100644
index 0000000..0d91d14
--- /dev/null
+++ b/src/crypto/hkdf.c
@@ -0,0 +1,131 @@
+/*
+* Copyright (c) 2024 Vaughn Nugent
+*
+* Package: noscrypt
+* File: hkdf.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/.
+*/
+
+
+#include "hkdf.h"
+
+/* Include string for memmove */
+#include <string.h>
+
+#define HKDF_MIN(a, b) (a < b ? a : b)
+
+STATIC_ASSERT(HKDF_IN_BUF_SIZE > SHA256_DIGEST_SIZE, "HDK Buffer must be at least the size of the underlying hashing alg output")
+
+static _nc_fn_inline void ncWriteSpanS(span_t* span, uint32_t offset, const uint8_t* data, uint32_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 _nc_fn_inline void debugValidateHandler(const struct nc_hkdf_fn_cb_struct* handler)
+{
+	DEBUG_ASSERT(handler != NULL)
+	DEBUG_ASSERT(handler->update != NULL)
+	DEBUG_ASSERT(handler->finish != NULL)
+}
+
+/*
+* The following functions implements the HKDF expand function using an existing
+* HMAC function.
+*
+* This follows the guidence from RFC 5869: https://tools.ietf.org/html/rfc5869
+*/
+
+cstatus_t hkdfExpandProcess(
+	const struct nc_hkdf_fn_cb_struct* handler,
+	void* ctx,
+	const cspan_t* info,
+	span_t* okm
+)
+{
+	cstatus_t result;
+
+	uint8_t counter;
+	uint32_t tLen, okmOffset;
+	uint8_t t[HKDF_IN_BUF_SIZE];
+	cspan_t tSpan, counterSpan;
+
+	debugValidateHandler(handler);
+
+	ncCryptoSecureZero(t, sizeof(t));
+
+	tLen = 0;				/* T(0) is an empty string(zero length) */
+	okmOffset = 0;
+	counter = 1;			/* counter is offset by 1 for init */
+	result = CSTATUS_FAIL;	/* Start in fail state */
+
+	/* counter as a span */
+	ncSpanInitC(&counterSpan, &counter, sizeof(counter));
+
+	/* Compute T(N) = HMAC(prk, T(n-1) | info | n) */
+	while (okmOffset < okm->size)
+	{
+		ncSpanInitC(&tSpan, t, tLen);
+
+		if (handler->update(ctx, &tSpan) != CSTATUS_OK)
+		{
+			goto Exit;
+		}
+
+		if (handler->update(ctx, info) != CSTATUS_OK)
+		{
+			goto Exit;
+		}
+
+		if (handler->update(ctx, &counterSpan) != CSTATUS_OK)
+		{
+			goto Exit;
+		}
+
+		/* 
+		* Write current hash state to t buffer. It is known
+		* that the t buffer must be at least the size of the
+		* underlying hash function output.
+		*/
+		if (handler->finish(ctx, t) != CSTATUS_OK)
+		{
+			goto Exit;
+		}
+
+		/* tlen becomes the hash size or remaining okm size */
+		tLen = HKDF_MIN(okm->size - okmOffset, SHA256_DIGEST_SIZE);
+
+		DEBUG_ASSERT(tLen <= sizeof(t));
+
+		/* write the T buffer back to okm */
+		ncWriteSpanS(okm, okmOffset, t, tLen);
+
+		/* shift base okm pointer by T */
+		okmOffset += tLen;
+
+		/* increment counter */
+		counter++;
+	}
+
+	result = CSTATUS_OK;	/* HMAC operation completed, so set success */
+
+Exit:
+
+	return result;
+}
diff --git a/src/crypto/impl/bcrypt.c b/src/crypto/impl/bcrypt.c
index 970f68f..9b01cac 100644
--- a/src/crypto/impl/bcrypt.c
+++ b/src/crypto/impl/bcrypt.c
@@ -32,8 +32,10 @@
 #include <bcrypt.h>
 
 #include "nc-util.h"
+#include "hkdf.h"
 
-#define IF_BC_FAIL(x) if(!BCRYPT_SUCCESS(x))
+#define IF_BC_FAIL(x) if(!BCRYPT_SUCCESS(x)) 
+#define BC_FAIL(x) if(!BCRYPT_SUCCESS(x)) return CSTATUS_FAIL;
 
 struct _bcrypt_ctx
 {
@@ -68,8 +70,8 @@ _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
+	 * I am not explicitly managing the update object buffer. By setting 
+	 * the update 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
@@ -113,7 +115,7 @@ _IMPLSTB NTSTATUS _bcFinishHash(const struct _bcrypt_ctx* ctx, sha256_t digestOu
 
 _IMPLSTB void _bcDestroyCtx(struct _bcrypt_ctx* ctx)
 {
-	/* Free the hash memory if it was allocated */
+	/* Free the update memory if it was allocated */
 	if(ctx->hHash) BCryptDestroyHash(ctx->hHash);
 	
 	/* Close the algorithm provider */
@@ -214,76 +216,42 @@ _IMPLSTB void _bcDestroyCtx(struct _bcrypt_ctx* ctx)
 
 	#define _IMPL_CRYPTO_SHA256_HKDF_EXPAND		_bcrypt_fallback_hkdf_expand
 
-	/* Include string for memmove */
-	#include <string.h>
-
-	static void ncWriteSpanS(span_t* span, uint32_t offset, const uint8_t* data, uint32_t size)
+	cstatus_t _bcrypt_hkdf_update(void* ctx, const cspan_t* data)
 	{
-		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")
+		DEBUG_ASSERT(ctx != NULL)
 
-		/* Copy data to span */
-		memmove(span->data + offset, data, size);
+		BC_FAIL(_bcHashData((struct _bcrypt_ctx*)ctx, data))
+		return CSTATUS_OK;
 	}
 
-	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
-	*/
+	cstatus_t _bcrypt_hkdf_finish(void* ctx, sha256_t hmacOut32)
+	{
+		DEBUG_ASSERT(ctx != NULL)
 
-	#define _BC_MIN(a, b) (a < b ? a : b)
+		BC_FAIL(_bcFinishHash((struct _bcrypt_ctx*)ctx, hmacOut32))
+		return CSTATUS_OK;
+	}
 
 	_IMPLSTB cstatus_t _bcrypt_fallback_hkdf_expand(const cspan_t* prk, const cspan_t* info, span_t* okm)
 	{
 		cstatus_t result;
 		struct _bcrypt_ctx ctx;
+		struct nc_hkdf_fn_cb_struct handler;
 
-		uint8_t counter;
-		uint32_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;
+		handler.update = _bcrypt_hkdf_update;
+		handler.finish = _bcrypt_hkdf_finish;
 
-			/* Set the length of the current hash state */
-			tLen = _BC_MIN(okm->size - okmOffset, SHA256_DIGEST_SIZE);
+		/* Init bcrypt */
+		BC_FAIL(_bcInitSha256(&ctx, BCRYPT_ALG_HANDLE_HMAC_FLAG))
 
-			DEBUG_ASSERT(tLen <= sizeof(t));
+		BC_FAIL(_bcCreateHmac(&ctx, prk))
 
-			/* write the T buffer back to okm */
-			ncWriteSpanS(okm, okmOffset, t, tLen);
+		/*
+		* NOTE! Hmac reusable flag must be set to allow for multiple
+		* calls to the finish function without losing the context.
+		*/
 
-			/* shift base okm pointer by T */
-			okmOffset += tLen;
-		}		
-
-		result = CSTATUS_OK;	/* HMAC operation completed, so set success */
-
-	Exit:
+		result = hkdfExpandProcess(&handler, &ctx, info, okm);
 
 		_bcDestroyCtx(&ctx);
 
diff --git a/src/crypto/impl/mbedtls.c b/src/crypto/impl/mbedtls.c
index ae36bbd..18eb9db 100644
--- a/src/crypto/impl/mbedtls.c
+++ b/src/crypto/impl/mbedtls.c
@@ -28,6 +28,11 @@
 
 #ifdef MBEDTLS_CRYPTO_LIB
 
+/* Inline errors on linux in header files on linux */
+#ifndef inline
+#define inline __inline
+#endif
+
 #include <mbedtls/md.h>
 #include <mbedtls/hkdf.h>
 #include <mbedtls/hmac_drbg.h>
@@ -35,11 +40,10 @@
 #include <mbedtls/chacha20.h>
 #include <mbedtls/constant_time.h>
 
-#include "nc-util.h"
+#ifndef inline
+#undef inline
+#endif
 
-/*
-*		EXPORT SUPPORTED FUNCTION OVERRIDES
-*/	
 
 _IMPLSTB const mbedtls_md_info_t* _mbed_sha256_alg(void)
 {
@@ -51,10 +55,8 @@ _IMPLSTB const mbedtls_md_info_t* _mbed_sha256_alg(void)
 }
 
 #if SIZE_MAX < UINT64_MAX
-	#define _ssize_guard(x) if(x > SIZE_MAX) return CSTATUS_FAIL;
 	#define _ssize_guard_int(x) if(x > SIZE_MAX) return 1;
 #else
-	#define _ssize_guard(x)
 	#define _ssize_guard_int(x)
 #endif
 
@@ -71,7 +73,7 @@ _IMPLSTB const mbedtls_md_info_t* _mbed_sha256_alg(void)
 		uint32_t dataLen
 	)
 	{
-		_ssize_guard(dataLen)
+		_sizet_check(dataLen)
 
 		/* Counter always starts at 0 */
 		return mbedtls_chacha20_crypt(
@@ -93,7 +95,7 @@ _IMPLSTB const mbedtls_md_info_t* _mbed_sha256_alg(void)
 
 	_IMPLSTB cstatus_t _mbed_sha256_digest(const cspan_t* data, sha256_t digestOut32)
 	{
-		_ssize_guard(data->size)
+		_sizet_check(data->size)
 
 		return mbedtls_sha256(
 			data->data, 
@@ -112,7 +114,7 @@ _IMPLSTB const mbedtls_md_info_t* _mbed_sha256_alg(void)
 
 	_IMPLSTB cstatus_t _mbed_sha256_hmac(const cspan_t* key, const cspan_t* data, sha256_t hmacOut32)
 	{
-		_ssize_guard(data->size)
+		_sizet_check(data->size)
 
 		/* Keys should never be large enough for this to matter, but sanity check. */
 		DEBUG_ASSERT2(key->size < SIZE_MAX, "Expected key size to be less than SIZE_MAX")
diff --git a/src/crypto/impl/monocypher.c b/src/crypto/impl/monocypher.c
index 790f5e9..b695d08 100644
--- a/src/crypto/impl/monocypher.c
+++ b/src/crypto/impl/monocypher.c
@@ -53,10 +53,7 @@
 		uint32_t dataLen
 	)
 	{
-		if(dataLen > SIZE_MAX)
-		{
-			return CSTATUS_FAIL;
-		}
+		_sizet_check(dataLen)
 
 		/*
 		* Function returns the next counter value which is not
diff --git a/src/crypto/impl/openssl.c b/src/crypto/impl/openssl.c
index 1217e60..90028e6 100644
--- a/src/crypto/impl/openssl.c
+++ b/src/crypto/impl/openssl.c
@@ -23,24 +23,156 @@
 
 #ifdef OPENSSL_CRYPTO_LIB
 
-#include <openssl/sha.h>
-
 #include "nc-util.h"
+#include <openssl/crypto.h>
 
-/*
-*		EXPORT SUPPORTED FUNCTIONS AS MACROS
-*/
-#define _IMPL_CHACHA20_CRYPT				_mbed_chacha20_encrypt			
-#define _IMPL_CRYPTO_SHA256_HMAC			_mbed_sha256_hmac
-#define _IMPL_CRYPTO_SHA256_DIGEST			_ossl_sha256_digest	
-#define _IMPL_CRYPTO_FIXED_TIME_COMPARE		_mbed_fixed_time_compare
-#define _IMPL_CRYPTO_SHA256_HKDF_EXPAND		_mbed_sha256_hkdf_expand
-#define _IMPL_CRYPTO_SHA256_HKDF_EXTRACT	_mbed_sha256_hkdf_extract
-
-_IMPLSTB int _ossl_sha256_digest(const uint8_t* data, uint8_t* digest)
-{
+#define _OSSL_FAIL(x) if(!(x)) return CSTATUS_FAIL;
+
+#ifndef _IMPL_SECURE_ZERO_MEMSET
+
+	#define _IMPL_SECURE_ZERO_MEMSET			_ossl_secure_zero_memset
+
+	_IMPLSTB void _ossl_secure_zero_memset(void* ptr, size_t size)
+	{
+		_sizet_check(size)
+
+		OPENSSL_cleanse(ptr, size);
+	}
+#endif
+
+#ifndef _IMPL_CRYPTO_FIXED_TIME_COMPARE
+
+	#define _IMPL_CRYPTO_FIXED_TIME_COMPARE		_ossl_fixed_time_compare
+
+	_IMPLSTB uint32_t _ossl_fixed_time_compare(const uint8_t* a, const uint8_t* b, uint32_t size)
+	{
+		int result;
+
+		_sizet_check(size)
+
+		result = CRYPTO_memcmp(a, b, size);
+
+		return (uint32_t)result;
+	}
+
+#endif /* _IMPL_CRYPTO_FIXED_TIME_COMPARE */
+
+
+#ifndef _IMPL_CRYPTO_SHA256_DIGEST	
+
+	#include <openssl/sha.h>
+
+	#define _IMPL_CRYPTO_SHA256_DIGEST			_ossl_sha256_digest	
+
+	_IMPLSTB cstatus_t _ossl_sha256_digest(const cspan_t* data, sha256_t digestOut32)
+	{
+		_sizet_check(data->size)
+
+		_OSSL_FAIL(SHA256(data->data, data->size, digestOut32))
+
+		return CSTATUS_OK;
+	}
+
+#endif
+
+#ifndef _IMPL_CRYPTO_SHA256_HMAC
+
+	#include <openssl/hmac.h>
+
+	/* Export function */
+	#define _IMPL_CRYPTO_SHA256_HMAC			_ossl_hmac_sha256	
+	
+	_IMPLSTB cstatus_t _ossl_hmac_sha256(const cspan_t* key, const cspan_t* data, sha256_t hmacOut32)
+	{
+		unsigned int hmacLen;
+
+		_sizet_check(key->size)
+		_sizet_check(data->size)
+
+		hmacLen = sizeof(sha256_t);
+
+		_OSSL_FAIL(
+			HMAC(
+				EVP_sha256(),
+				key->data,
+				key->size,
+				data->data,
+				data->size,
+				hmacOut32,
+				&hmacLen
+			)
+		)
+		
+		/* digest length should match the actual digest size */
+		_OSSL_FAIL(hmacLen != sizeof(sha256_t))
+
+		return CSTATUS_OK;
+	}
+
+#endif /* !_IMPL_CRYPTO_SHA256_HMAC */
+
+#ifndef _IMPL_CRYPTO_SHA256_HKDF_EXPAND
+	
+	#include <openssl/hmac.h>
+	#include "hkdf.h"
+
+	#define _IMPL_CRYPTO_SHA256_HKDF_EXPAND		_ossl_sha256_hkdf_expand
+
+	cstatus_t _ossl_hkdf_update(void* ctx, const cspan_t* data)
+	{
+		DEBUG_ASSERT(ctx != NULL)
+
+		_OSS_FAIL(HMAC_Update((HMAC_CTX*)ctx, data->data, data->size))
+		
+		return CSTATUS_OK;
+	}
+
+	cstatus_t _ossl_hkdf_finish(void* ctx, sha256_t hmacOut32)
+	{
+		DEBUG_ASSERT(ctx != NULL)
+
+		_OSSL_FAIL(HMAC_Final((HMAC_CTX*)ctx, hmacOut32, NULL))
+		
+		return CSTATUS_OK;
+	}
+
+	_IMPLSTB cstatus_t _ossl_fallback_hkdf_expand(const cspan_t* prk, const cspan_t* info, span_t* okm)
+	{
+		HMAC_CTX* hmac;
+		cstatus_t result;
+		struct nc_hkdf_fn_cb_struct handler;
 	
-}
+		/*
+		* NOTE! Hmac reusable flag must be set to allow for multiple
+		* calls to the finish function without losing the context.
+		*/
+
+		if ((hmac = HMAC_CTX_new()) == NULL)
+		{
+			return CSTATUS_FAIL;
+		}
+
+
+		_OSSL_FAIL(
+			HMAC_Init_ex(
+				hmac,
+				prk->data,
+				pkr->size,
+				EVP_sha256(),
+				NULL
+			)
+		)
+		
+		handler.update = _ossl_hkdf_update;
+		handler.finish = _ossl_hkdf_finish;
+
+		result = hkdfExpandProcess(&handler, hmac, info, okm);
+
+		HMAC_CTX_free(hmac);
+
+		return result;
+	}
 
+#endif /* !_IMPL_CRYPTO_SHA256_HKDF_EXPAND */
 
 #endif	/*!OPENSSL_CRYPTO_LIB */
\ No newline at end of file
diff --git a/src/crypto/nc-crypto.c b/src/crypto/nc-crypto.c
index fb2c0da..587d59d 100644
--- a/src/crypto/nc-crypto.c
+++ b/src/crypto/nc-crypto.c
@@ -23,7 +23,7 @@
 
 /*
 *  Functions are not forced inline, just suggested.
-*  So unless it beomes a performance issue, I will leave
+*  So unless it becomes a performance issue, I will leave
 *  most/all impl functions inline and let the compiler 
 *  decide.
 */
@@ -43,10 +43,12 @@
 * 
 * Macros are used to allow the preprocessor to select the correct implementation
 * or raise errors if no implementation is defined.
+* 
+* Implementation functions can assume inputs have been checked/sanitized by the
+* calling function, and should return CSTATUS_OK on success, CSTATUS_FAIL on failure.
 */
 
 
-
 /*
 * Prioritize embedded builds with mbedtls
 */
@@ -69,7 +71,8 @@
 * memset 0 functions for each platform.
 */
 #ifndef _IMPL_SECURE_ZERO_MEMSET
-	#if defined(__GNUC__)	
+   /* only incude bzero if libc version greater than 2.25 */
+	#if defined(__GLIBC__) && defined(__GLIBC_MINOR__) && __GLIBC__ >= 2 && __GLIBC_MINOR__ >= 25
 		/*
 		*	When using libc, we can use explicit_bzero
 		*	as secure memset implementation.
@@ -77,6 +80,7 @@
 		*	https://sourceware.org/glibc/manual/2.39/html_mono/libc.html#Erasing-Sensitive-Data
 		*/
 		#include <string.h>
+		extern void explicit_bzero(void* block, size_t len);
 		#define _IMPL_SECURE_ZERO_MEMSET explicit_bzero
 	#endif
 #endif
@@ -161,6 +165,11 @@
 * Internal function implementations that perform
 * basic checking and call the correct implementation
 * for the desired crypto impl.
+* 
+* The following functions MUST be assumed to 
+* perform basic input validation. Since these apis are 
+* internal, debug asserts are used to ensure the
+* function has been used correctly.
 */
 
 void ncCryptoSecureZero(void* ptr, uint32_t size)
@@ -223,6 +232,8 @@ cstatus_t ncCryptoSha256HkdfExpand(const cspan_t* prk, const cspan_t* info, span
 	/*
 	* RFC 5869: 2.3
 	* "length of output keying material in octets (<= 255 * HashLen)"
+	* 
+	* important as the counter is 1 byte, so it cannot overflow
 	*/
 
 	if(okm->size > (uint32_t)(0xFFu * SHA256_DIGEST_SIZE))
diff --git a/src/crypto/nc-crypto.h b/src/crypto/nc-crypto.h
deleted file mode 100644
index 64d4ad8..0000000
--- a/src/crypto/nc-crypto.h
+++ /dev/null
@@ -1,73 +0,0 @@
-
-/*
-* Copyright (c) 2024 Vaughn Nugent
-*
-* Package: noscrypt
-* File: nc-crypto.h
-*
-* 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/.
-*/
-
-#pragma once
-
-#ifndef _NC_CRYPTO_H
-#define _NC_CRYPTO_H
-
-#include <stdint.h>
-
-#define CHACHA_NONCE_SIZE		0x0cu		/* Size of 12 is set by the cipher spec */
-#define CHACHA_KEY_SIZE			0x20u		/* Size of 32 is set by the cipher spec */
-#define SHA256_DIGEST_SIZE		0x20u		/* Size of 32 is set by the cipher spec */
-
-typedef uint8_t cstatus_t;
-#define CSTATUS_OK				((cstatus_t)0x01u)
-#define CSTATUS_FAIL			((cstatus_t)0x00u)
-
-typedef uint8_t sha256_t[SHA256_DIGEST_SIZE];
-
-/*
-* IMPORTANT:
-*	The HKDF_IN_BUF_SIZE defintion sets the internal stack buffer size to use
-*	during fallback HKDF_Expand operations.
-*
-*	128 bytes should be more than enough for most use cases, without going 
-*	overboard. Could be dialed in better for specific use cases later.
-*/
-
-#ifndef HKDF_IN_BUF_SIZE
-	#define HKDF_IN_BUF_SIZE	0x80	
-#endif
-
-
-uint32_t ncCryptoFixedTimeComp(const uint8_t* a, const uint8_t* b, uint32_t size);
-
-void ncCryptoSecureZero(void* ptr, uint32_t size);
-
-cstatus_t ncCryptoDigestSha256(const cspan_t* data, sha256_t digestOut32);
-
-cstatus_t ncCryptoHmacSha256(const cspan_t* key, const cspan_t* data, sha256_t hmacOut32);
-
-cstatus_t ncCryptoSha256HkdfExpand(const cspan_t* prk, const cspan_t* info, span_t* okm);
-
-cstatus_t ncCryptoSha256HkdfExtract(const cspan_t* salt, const cspan_t* ikm, sha256_t prk);
-
-cstatus_t ncCryptoChacha20(
-	const uint8_t key[CHACHA_KEY_SIZE],
-	const uint8_t nonce[CHACHA_NONCE_SIZE],
-	const uint8_t* input,
-	uint8_t* output,
-	uint32_t dataSize
-);
-
-#endif /* !_NC_CRYPTO_H */
diff --git a/src/noscrypt.c b/src/noscrypt.c
index 00684b8..fac3dfb 100644
--- a/src/noscrypt.c
+++ b/src/noscrypt.c
@@ -21,10 +21,10 @@
 #include "noscrypt.h"
 
 #include "nc-util.h"
-#include "crypto/nc-crypto.h"
+#include "nc-crypto.h"
 
-#include <secp256k1_ecdh.h>
-#include <secp256k1_schnorrsig.h>
+#include <secp256k1/secp256k1_ecdh.h>
+#include <secp256k1/secp256k1_schnorrsig.h>
 
 /*
 * Local macro for secure zero buffer fill
@@ -768,19 +768,30 @@ NC_EXPORT NCResult NC_CC NCEncrypt(
 	CHECK_INVALID_ARG(args->hmacKeyOut32, 3)
 	CHECK_ARG_RANGE(args->dataSize, NIP44_MIN_ENC_MESSAGE_SIZE, NIP44_MAX_ENC_MESSAGE_SIZE, 3)
 
+	switch(args->version)
+	{		
+		case NC_ENC_VERSION_NIP44:
+			break; /* Allow nip44 */
+
+		/* At the moment nip04 compatability is not supported */
+		case NC_ENC_VERSION_NIP04:
+		default:
+			return E_VERSION_NOT_SUPPORTED;
+	}
+	
 	/* Compute the shared point */
 	if ((result = _computeSharedSecret(ctx, sk, pk, &sharedSecret)) != NC_SUCCESS)
 	{
 		goto Cleanup;
 	}
-	
+
 	/* Compute the conversation key from secret and pubkic keys */
 	if ((result = _computeConversationKey(ctx, &sharedSecret, &conversationKey)) != NC_SUCCESS)
 	{
 		goto Cleanup;
 	}
 
-	result = _encryptEx(ctx, &conversationKey, args->hmacKeyOut32, args);
+	result = _encryptEx(ctx, &conversationKey, args->hmacKeyOut32, args);	
 
 Cleanup:
 	/* Clean up sensitive data */