/******************************************************************************/
/* SlunkCrypt, by LoRd_MuldeR <MuldeR2@GMX.de>                                */
/* This work has been released under the CC0 1.0 Universal license!           */
/******************************************************************************/

/* Internal */
#include "compiler.h"
#include "slunkcrypt.h"

/* CRT */
#include <string.h>
#include <fcntl.h>
#include <limits.h>

/* Utils */
static INLINE size_t MIN_SIZE(const size_t a, const size_t b) { return (a > b) ? b : a; }

// ==========================================================================
// Platform compatibility
// ==========================================================================

#ifdef _WIN32
#  include <Windows.h>
#  include <io.h>
#else
#  include <unistd.h>
#  include <sched.h>
#  include <strings.h>
#endif

#if defined(__APPLE__) && defined(__MACH__)
#  include <sys/random.h>
#endif

/* detect compiler destructor support */
#undef HAVE_ATTRIB_DESTRUCTOR
#if defined(__GNUC__) || defined(__clang__)
#  define HAVE_ATTRIB_DESTRUCTOR 1
#endif

/* detect getentropy() or RtlGenRandom() support */
#undef HAVE_GETENTROPY
#undef HAVE_WIN32RTLGENRANDOM
#if defined(_WIN32)
#  define HAVE_WIN32RTLGENRANDOM 1
#elif (defined(__linux__) && !defined(__UCLIBC__)) || defined(__CYGWIN__) || (defined(__FreeBSD__) && (__FreeBSD__ >= 12)) || defined(__DragonFly__) || defined(__OpenBSD__) || (defined(__sun) && defined(__SVR4)) || (defined(__APPLE__) && defined(__MACH__))
#  define HAVE_GETENTROPY 1
#else
#  pragma message("Function getentropy() is *not* available -> using fallback!")
#endif

/* detect explicit_bzero() or SecureZeroMemory() support */
#undef EXPLICIT_BZERO
#if defined(_WIN32) && defined(SecureZeroMemory)
#  define EXPLICIT_BZERO SecureZeroMemory
#elif (defined(__linux__) && !defined(__UCLIBC__)) || defined(__CYGWIN__) || (defined(__FreeBSD__) && (__FreeBSD__ >= 11)) || defined(__DragonFly__) || defined(__OpenBSD__) || (defined(__sun) && defined(__SVR4)) || defined(__HAIKU__)
#  define EXPLICIT_BZERO explicit_bzero
#else
#  pragma message("Function explicit_bzero() is *not* available -> using fallback!")
#endif

/* detect sched_yield() or Sleep() support */
#ifdef _WIN32
#  define THREAD_YIELD() Sleep(0)
#else
#  define THREAD_YIELD() sched_yield()
#endif

// ==========================================================================
// One-time initialization
// ==========================================================================

/* atomic memory access */
#if defined(_MSC_VER) && (!defined(__GNUC__))
#  define ONCE_FLAG_T LONG
#  define ATOMIC_COMPARE_EXCHANGE(X,Y,Z) InterlockedCompareExchange((X),(Y),(Z))
#  define ATOMIC_EXCHANGE(X,Y) InterlockedExchange((X),(Y))
#else
#  define ONCE_FLAG_T int
#  define ATOMIC_COMPARE_EXCHANGE(X,Y,Z) __sync_val_compare_and_swap((X),(Z),(Y))
#  define ATOMIC_EXCHANGE(X,Y) __sync_lock_test_and_set((X),(Y))
#endif

/* execute init routine once */
static INLINE void initialize_once(volatile ONCE_FLAG_T *const once_control, void (*const init_routine)(void))
{
	ONCE_FLAG_T state;
	while ((state = ATOMIC_COMPARE_EXCHANGE(once_control, -1, 0)) != 0)
	{
		if (state > 0)
		{
			return; /*already initialized*/
		}
		THREAD_YIELD();
	}
	init_routine();
	ATOMIC_EXCHANGE(once_control, 1);
}

// ==========================================================================
// Random bytes
// ==========================================================================

#define MAX_COUNT 1048576U

/* Global state */
static volatile ONCE_FLAG_T s_random_is_initialized = 0;
#if defined(_WIN32)
typedef BOOLEAN(WINAPI *ptr_genrandom_t)(void *buffer, ULONG length);
static HMODULE s_advapi32 = NULL;
static ptr_genrandom_t s_genrandom = NULL;
#else
static const char *const DEV_RANDOM[] = { "/dev/urandom", "/dev/arandom", "/dev/random", NULL };
static int s_random_fd = -1;
#endif

/* De-initialize CSRNG */
static void exit_random_source(void)
{
#ifdef _WIN32
	if (s_genrandom)
	{
		s_genrandom = NULL;
	}
	if (s_advapi32)
	{
		FreeLibrary(s_advapi32);
		s_advapi32 = NULL;
	}
#else
	if (s_random_fd >= 0)
	{
		close(s_random_fd);
		s_random_fd = -1;
	}
#endif
}

/* Initialize CSRNG */
static void init_random_source(void)
{
#ifdef HAVE_WIN32RTLGENRANDOM
	if ((s_advapi32 = LoadLibraryW(L"advapi32.dll")))
	{
		s_genrandom = (ptr_genrandom_t) GetProcAddress(s_advapi32, "SystemFunction036");
	}
#else
#if defined(HAVE_GETENTROPY)
	uint8_t temp;
	if (getentropy(&temp, sizeof(uint8_t)) >= 0)
	{
		goto init_completed;
	}
#endif
	for (size_t i = 0U; DEV_RANDOM[i]; ++i)
	{
		if ((s_random_fd = open(DEV_RANDOM[i], O_RDONLY)) >= 0)
		{
			goto init_completed;
		}
	}
init_completed: ;
#endif
#if !defined(HAVE_ATTRIB_DESTRUCTOR)
	atexit(exit_random_source);
#endif
}

/* Generate random bytes */
size_t slunkcrypt_random_bytes(uint8_t *const buffer, const size_t length)
{
	size_t offset;
	initialize_once(&s_random_is_initialized, init_random_source);
#ifdef HAVE_WIN32RTLGENRANDOM
	if (s_genrandom)
	{
		ULONG count;
		for (offset = 0U; offset < length; offset += count)
		{
			count = (ULONG) MIN_SIZE(length - offset, MAX_COUNT);
			if (!s_genrandom(buffer + offset, count))
			{
				break; /*failed*/
			}
		}
		return offset;
	}
	return 0U;
#else
	if (s_random_fd >= 0)
	{
		ssize_t count;
		for (offset = 0; offset < length; offset += (size_t)count)
		{
			if (!((count = read(s_random_fd, buffer + offset, MIN_SIZE(length - offset, MAX_COUNT))) > 0))
			{
				break; /*failed*/
			}
		}
		return offset;
	}
#if defined(HAVE_GETENTROPY)
	else
	{
		size_t count;
		for (offset = 0U; offset < length; offset += count)
		{
			count = MIN_SIZE(length - offset, 256U); /*the maximum permitted value is 256*/
			if (getentropy(buffer + offset, count) < 0)
			{
				break; /*failed*/
			}
		}
		return offset;
	}
#else
	return 0U;
#endif
#endif
}

// ==========================================================================
// Zero memory
// ==========================================================================

void slunkcrypt_bzero(void* const buffer, const size_t length)
{
	if ((buffer) && (length > 0U))
	{
#if defined(EXPLICIT_BZERO)
		EXPLICIT_BZERO(buffer, length);
#else
		volatile uint8_t* ptr = (volatile uint8_t*) buffer;
		for (size_t i = 0U; i < length; ++i)
		{
			ptr[i] = 0U;
		}
#endif
	}
}

// ==========================================================================
// Destructor
// ==========================================================================

#if defined(HAVE_ATTRIB_DESTRUCTOR)
__attribute__((destructor)) void slunkcrypt_destructor(void)
{
	exit_random_source();
}
#endif