/********************************************************/
/*  ntapi: Native API core library                      */
/*  Copyright (C) 2013--2017  Z. Gilboa                 */
/*  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"

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

static intptr_t __fastcall __ntapi_tt_fork_finalize(
	void **		hprocess,
	void **		hthread)
{
	int32_t			status;
	nt_oa			oa;
	nt_cid			cid;
	ntapi_internals *	__internals;

	__internals = __ntapi_internals();

	*hprocess = 0;
	*hthread  = 0;

	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();

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

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

	return 0;
}

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 intptr_t __fastcall __ntapi_tt_fork_child(
	void *		hresumed,
	void *		hready)
{
	int32_t			status;
	nt_thread_params	tparams;
	nt_timeout		timeout;
	nt_timeout		zerowait;
	intptr_t		state;

	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	= 0x10000;
	tparams.stack_size_reserve	= 0x20000;

	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 0;
	}

	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 0;
	}

	__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 0;
	}

	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) {
		__ntapi->zw_close(hresumed);
		__ntapi->zw_close(hready);
		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);

		__ntapi->zw_close(hresumed);
		__ntapi->zw_close(hready);
		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;
}

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

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

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

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

		if (pid == 0) {
			__ntapi_tt_fork_child(
				hresumed,hready);

			return __ntapi_tt_fork_finalize(
				hprocess,hthread);

		} else if (pid > 0) {
			if (!(__ntapi_tt_fork_parent(
					hprocess,hthread,
					hresumed,hready)))
				return pid;

		} else {
			__ntapi->zw_close(hresumed);
			__ntapi->zw_close(hready);
			return -1;
		}
	}

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

	return -1;
}