/********************************************************/
/*  ntapi: Native API core library                      */
/*  Copyright (C) 2013--2021  SysDeer Technologies, LLC */
/*  Released under GPLv2 and GPLv3; see COPYING.NTAPI.  */
/********************************************************/

#include <psxtypes/psxtypes.h>
#include <ntapi/nt_atomic.h>
#include <ntapi/nt_status.h>
#include <ntapi/nt_object.h>
#include <ntapi/nt_memory.h>
#include <ntapi/nt_thread.h>
#include <ntapi/nt_process.h>
#include <ntapi/ntapi.h>
#include "ntapi_impl.h"

#define __FORK_THREAD_STACK_COMMIT  (0x10000)
#define __FORK_THREAD_STACK_RESERVE (0x20000)

static intptr_t	__fork_retry_stats  = 0;
static intptr_t	__fork_resume_stats = 0;

static int __ipc_memfn(
	struct dalist_ex *	dlist,
	void **			addr,
	size_t *		alloc_size)
{
	(void)dlist;
	(void)addr;
	(void)alloc_size;

	return DALIST_EMEMFN;
}

static int32_t __fastcall __ntapi_tt_fork_finalize(void ** hprocess)
{
	int32_t			status;
	int			page;
	nt_rtdata *		rtdata;
	ntapi_internals *	__internals;

	__internals = __ntapi_internals();
	rtdata      = __internals->rtdata;

	rtdata->cid_parent.process_id = rtdata->cid_self.process_id;
	rtdata->cid_parent.thread_id  = rtdata->cid_self.thread_id;

	rtdata->cid_self.process_id   = pe_get_current_process_id();
	rtdata->cid_self.thread_id    = pe_get_current_thread_id();

	if ((status = __ntapi->zw_duplicate_object(
			__internals->hprocess,
			__internals->hprocess,
			__internals->hprocess,
			hprocess,0,0,
			NT_DUPLICATE_SAME_ATTRIBUTES
			|NT_DUPLICATE_SAME_ACCESS)))
		return status;

	if ((status = dalist_init_ex(
			&__internals->ipc_conns,
			sizeof(nt_ipc_conn),
			NT_ALLOCATION_GRANULARITY,
			__ipc_memfn,
			DALIST_MEMFN_CUSTOM)))
		return status;

	dalist_deposit_memory_block(
		&__internals->ipc_conns,
		__internals->ntapi_img_sec_bss->ipc_buffer,
		__NT_BSS_IPC_BUFFER_SIZE);

	for (page=0; page<__internals->ipc_page; page++)
		dalist_deposit_memory_block(
			&__internals->ipc_conns,
			__internals->ipc_pages[page],
			NT_ALLOCATION_GRANULARITY);

	rtdata->hsemctl     = 0;
	rtdata->hsempid     = 0;

	rtdata->hmsqctl     = 0;
	rtdata->hmsqpid     = 0;

	rtdata->haflctl     = 0;
	rtdata->haflpid     = 0;

	rtdata->ipc_keys[0] = 0;
	rtdata->ipc_keys[1] = 0;
	rtdata->ipc_keys[2] = 0;
	rtdata->ipc_keys[3] = 0;
	rtdata->ipc_keys[4] = 0;
	rtdata->ipc_keys[5] = 0;

	return NT_STATUS_SUCCESS;
}

static int32_t __stdcall __fork_thread(void * ctx)
{
	intptr_t *	pstate;
	intptr_t	state;
	void *		hready;

	pstate = (intptr_t *)ctx;
	state  = *pstate;
	hready = (void *)state;

	at_store(
		pstate,
		0);

	return __ntapi->zw_terminate_thread(
		NT_CURRENT_THREAD_HANDLE,
		__ntapi->zw_set_event(
			hready,0));
}

static int32_t __fastcall __ntapi_tt_fork_child(
	void *		hresumed,
	void *		hready,
	void **		hthread)
{
	int32_t			status;
	nt_thread_params	tparams;
	nt_timeout		timeout;
	nt_timeout		zerowait;
	intptr_t		state;
	nt_oa			oa;
	nt_cid			cid;
	ntapi_internals *	__internals;

	oa.len		= sizeof(oa);
	oa.root_dir	= 0;
	oa.obj_name	= 0;
	oa.obj_attr	= 0;
	oa.sec_desc	= &__internals->seq_desc;
	oa.sec_qos	= &__internals->seq_qos;

	cid.process_id	= pe_get_current_process_id();
	cid.thread_id	= pe_get_current_thread_id();

	__internals = __ntapi_internals();

	status = __ntapi->zw_open_process(
		&__internals->hprocess,
		NT_PROCESS_ALL_ACCESS,
		&oa,&cid);

	if (status == NT_STATUS_SUCCESS)
		status = __ntapi->zw_open_thread(
			hthread,
			NT_THREAD_ALL_ACCESS,
			&oa,&cid);

	if (status) {
		__ntapi->zw_set_event(
			hresumed,0);

		__ntapi->zw_terminate_process(
			NT_CURRENT_PROCESS_HANDLE,
			status);
	}

	at_store(
		&state,
		(intptr_t)hready);

	__ntapi->tt_aligned_block_memset(
		&tparams,0,sizeof(tparams));

	tparams.start			= __fork_thread;
	tparams.arg			= &state;
	tparams.stack_size_commit	= __FORK_THREAD_STACK_COMMIT;
	tparams.stack_size_reserve	= __FORK_THREAD_STACK_RESERVE;

	status = __ntapi->tt_create_local_thread(
		&tparams);

	__ntapi->zw_set_event(
		hresumed,0);

	if (status)
		__ntapi->zw_terminate_process(
			NT_CURRENT_PROCESS_HANDLE,
			status);

	if (!state) {
		__ntapi->zw_close(hresumed);
		__ntapi->zw_close(hready);
		__ntapi->zw_close(tparams.hthread);
		return NT_STATUS_SUCCESS;
	}

	timeout.quad  = (-1) * 10 * 1000 * 250;

	status = __ntapi->zw_wait_for_single_object(
		hready,
		NT_SYNC_NON_ALERTABLE,
		&timeout);

	if (status == NT_STATUS_SUCCESS) {
		__ntapi->zw_close(hresumed);
		__ntapi->zw_close(hready);
		__ntapi->zw_close(tparams.hthread);
		return NT_STATUS_SUCCESS;
	}

	__ntapi->zw_terminate_thread(
		tparams.hthread,
		NT_STATUS_MORE_PROCESSING_REQUIRED);

	zerowait.quad = 0;

	status = __ntapi->zw_wait_for_single_object(
		hready,
		NT_SYNC_NON_ALERTABLE,
		&zerowait);

	if (status == NT_STATUS_SUCCESS) {
		__ntapi->zw_close(hresumed);
		__ntapi->zw_close(hready);
		__ntapi->zw_close(tparams.hthread);
		return NT_STATUS_SUCCESS;
	}

	return __ntapi->zw_terminate_process(
		NT_CURRENT_PROCESS_HANDLE,
		status);
}

static intptr_t __fastcall __ntapi_tt_fork_parent(
	void **		hprocess,
	void **		hthread,
	void *		hresumed,
	void *		hready)
{
	int32_t		status;
	nt_timeout	timeout;
	nt_timeout	zerowait;
	uint32_t	prev;

	__ntapi->zw_wait_for_single_object(
		hresumed,
		NT_SYNC_NON_ALERTABLE,
		0);

	timeout.quad = (-1) * 10 * 1000 * 500;

	status = __ntapi->zw_wait_for_single_object(
		hready,
		NT_SYNC_NON_ALERTABLE,
		&timeout);

	if (status == NT_STATUS_SUCCESS)
		return NT_STATUS_SUCCESS;

	__ntapi->zw_suspend_thread(
		*hthread,&prev);

	zerowait.quad = 0;

	status = __ntapi->zw_wait_for_single_object(
		hready,
		NT_SYNC_NON_ALERTABLE,
		&zerowait);

	if (status == NT_STATUS_SUCCESS) {
		at_locked_inc(
			&__fork_resume_stats);

		__ntapi->zw_resume_thread(
			*hthread,0);

		return NT_STATUS_SUCCESS;
	}

	at_locked_inc(
		&__fork_retry_stats);

	__ntapi->zw_terminate_process(
		*hprocess,
		status);

	__ntapi->zw_close(*hprocess);
	__ntapi->zw_close(*hthread);

	return status;
}

int32_t __fastcall __ntapi_tt_fork(
	__out	void **		hprocess,
	__out	void **		hthread,
	__out	nt_cid *	cid)
{
	int32_t			status;
	void *			hresumed;
	void *			hready;
	int			i;

	if ((status = __ntapi->tt_create_inheritable_event(
			&hresumed,
			NT_NOTIFICATION_EVENT,
			NT_EVENT_NOT_SIGNALED)))
		return status;

	status = __ntapi->tt_create_inheritable_event(
		&hready,
		NT_NOTIFICATION_EVENT,
		NT_EVENT_NOT_SIGNALED);

	if (status) {
		__ntapi->zw_close(hresumed);
		return status;
	}

	for (i=0; i<32; i++) {
		if (__ntapi->zw_create_user_process)
			status = __ntapi_tt_fork_v2(hprocess,hthread,cid);
		else
			status = __ntapi_tt_fork_v1(hprocess,hthread,cid);

		if (status) {
			__ntapi->zw_close(hresumed);
			__ntapi->zw_close(hready);
			return status;
		}

		if (cid->process_id == 0) {
			__ntapi_tt_fork_child(
				hresumed,hready,hthread);

			return __ntapi_tt_fork_finalize(
				hprocess);
		}

		status = __ntapi_tt_fork_parent(
			hprocess,hthread,
			hresumed,hready);

		if (status == NT_STATUS_SUCCESS) {
			__ntapi->zw_close(hresumed);
			__ntapi->zw_close(hready);
			return NT_STATUS_SUCCESS;
		}
	}

	__ntapi->zw_close(hresumed);
	__ntapi->zw_close(hready);

	return NT_STATUS_UNSUCCESSFUL;
}