/*********************************************************/
/*  toksvc: a framework-native token broker service      */
/*  Copyright (C) 2020  SysDeer Technologies, LLC        */
/*  Released under GPLv2 and GPLv3; see COPYING.TOKSVC.  */
/*********************************************************/

#include <stddef.h>
#include <ntapi/ntapi.h>
#include "toksvc_memfn_impl.h"
#include "toksvc_driver_impl.h"

void * toks_calloc(size_t n, size_t size)
{
	void *                      addr;
	size_t                      bsize;
	struct toks_memory_block *  block;

	if (!n || (size > (size_t)(-1) / n))
		return 0;

	size *= n;
	size += 0xf;
	size &= ~(size_t)0xf;

	addr  = 0;
	bsize = size + sizeof(*block);

	if (ntapi->zw_allocate_virtual_memory(
			NT_CURRENT_PROCESS_HANDLE,
			&addr,
			0,
			&bsize,
			NT_MEM_COMMIT,
			NT_PAGE_READWRITE))
		return 0;

	block        = (struct toks_memory_block *)addr;
	block->addr  = addr;
	block->size  = bsize;
	block->used  = size + offsetof(struct toks_memory_block,buffer);
	block->avail = bsize - block->used;

	return (char *)addr + offsetof(struct toks_memory_block,buffer);
}

void * toks_balloc(
	struct toks_memory_block * cache,
	struct toks_memory_block * block,
	size_t n, size_t size)
{
	void * addr;

	if (!n || (size > (size_t)(-1) / n))
		return 0;

	if (!cache || !cache->addr)
		return toks_calloc(n,size);

	size *= n;
	size += 0xf;
	size &= ~(size_t)0xf;

	if (size > block->avail)
		return toks_calloc(1,size);

	/* cached */
	addr          = &block->buffer[block->used / sizeof(size_t)];
	block->used  += size;
	block->avail -= size;

	return addr;
}

void toks_free(void * addr)
{
	char *                      baddr;
	size_t                      size;
	struct toks_memory_block *  block;

	baddr  = addr;
	baddr -= offsetof(struct toks_memory_block,buffer);

	block = (struct toks_memory_block *)baddr;
	addr  = block->addr;
	size  = block->size;

	ntapi->zw_free_virtual_memory(
		NT_CURRENT_PROCESS_HANDLE,
		&addr,
		&size,
		NT_MEM_RELEASE);
}