/********************************************************/

/*  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			= &stat;;

	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;

}