/********************************************************/
/*  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 <pemagine/pemagine.h>
#include <ntapi/nt_status.h>
#include <ntapi/nt_object.h>
#include <ntapi/nt_thread.h>
#include <ntapi/nt_process.h>
#include <ntapi/nt_string.h>
#include <ntapi/ntapi.h>
#include "ntapi_impl.h"

#define NT_PROCESS_SPAWN_FLAG_DEBUG_MASK                \
		(NT_PROCESS_SPAWN_FLAG_DEBUG_EXECUTION   \
		 | NT_PROCESS_SPAWN_FLAG_DEBUG_SUSPENDED)

#define __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_IMGBUF_SIZE    (0x10000)
#define __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_DEF (0x80000)
#define __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_MAX (0x800000)

static int32_t __stdcall __tt_spawn_return(
	nt_runtime_data_block *	rtblock,
	void *			himgfile,
	void *			hprocess,
	void *			hthread,
	int32_t			status)
{
	nt_runtime_data * rtdata;

	rtdata = (nt_runtime_data *)rtblock->addr;

	if (hprocess) {
		__ntapi->zw_terminate_process(
			hprocess,status);

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

	if (rtdata->hready)
		__ntapi->zw_close(
			rtdata->hready);

	if (himgfile)
		__ntapi->zw_close(himgfile);

	__ntapi->zw_free_virtual_memory(
		NT_CURRENT_PROCESS_HANDLE,
		&rtblock->addr,
		&rtblock->size,
		NT_MEM_RELEASE);

	return status;
}

int32_t __stdcall __ntapi_tt_spawn_native_process(nt_spawn_process_params * sparams)
{
	int32_t				status;
	nt_create_process_params	cparams;
	nt_tty_session_info		session;
	nt_runtime_data_block		rtblock;
	nt_runtime_data_block		crtblock;
	nt_runtime_data *		rtctx;
	nt_runtime_data *		rdata;
	nt_unicode_string *		imgname;
	nt_peb * 			peb;
	size_t				asize;
	char *				patharg;
	void *				hat;
	void *				hfile;
	void *				himgfile;
	char *				src;
	char *				dst;
	int				envc;
	char **				penv;
	char **				pref;
	char **				parg;
	char **				rargv;
	char **				renvp;
	wchar16_t **			pwarg;
	wchar16_t *			wenv;
	wchar16_t *			wch;
	wchar16_t *			wcap;
	void *				hchild[2];
	wchar16_t *			imgbuf;
	uint32_t			fsuspended;
	size_t				buflen;
	size_t				written;
	size_t				needed;
	uintptr_t			addr;
	char *				raddr;
	size_t				rsize;

	/* rtctx (convenience) */
	rtctx = sparams->rtctx;

	/* validation */
	if (!sparams->himage && !sparams->patharg)
		return NT_STATUS_OBJECT_PATH_INVALID;

	if (rtctx->argc || rtctx->argv || rtctx->envc || rtctx->envp)
		return NT_STATUS_INVALID_PARAMETER;

	if (rtctx->hready || rtctx->hsession)
		return NT_STATUS_INVALID_PARAMETER_MIX;

	if (!(peb = (nt_peb *)pe_get_peb_address()))
		return NT_STATUS_INTERNAL_ERROR;

	if (!peb->process_params)
		return NT_STATUS_INTERNAL_ERROR;

	/* hat */
	hat = (sparams->hroot && (sparams->argv[0][0] == '/'))
		? sparams->hroot
		: rtctx->hcwd
			? rtctx->hcwd
			: peb->process_params->cwd_handle;

	/* patharg */
	patharg = sparams->patharg
		? (sparams->patharg[0] == '/')
			? (sparams->patharg[1] == '?')
				? &sparams->patharg[0]
				: &sparams->patharg[1]
			: &sparams->patharg[0]
		: 0;

	/* quickly determine whether a large buffer is needed */
	for (asize=0,parg=sparams->argv; parg && *parg; asize++) {
		asize += __ntapi_tt_string_null_offset_multibyte(*parg++);

		if (asize > (__SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_DEF >> 4))
			parg = 0;
	}

	for (penv=sparams->envp; penv && *penv; asize++) {
		asize += __ntapi_tt_string_null_offset_multibyte(*penv++);

		if (asize > (__SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_DEF >> 4))
			penv = 0;
	}

	if (parg && penv) {
		asize += (parg - sparams->argv) * sizeof(char *);
		asize += (penv - sparams->envp) * sizeof(char *);
	}

	asize = (asize <= (__SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_DEF >> 4))
		? __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_DEF
		: __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_ALLOC_SIZE_MAX;

	/* rtblock, rdata */
	rtblock.addr		= 0;
	rtblock.size		= asize;
	rtblock.remote_addr	= 0;
	rtblock.remote_size	= 0;
	rtblock.flags		= 0;

	if ((status = __ntapi->zw_allocate_virtual_memory(
			NT_CURRENT_PROCESS_HANDLE,
			&rtblock.addr,0,
			&rtblock.size,
			NT_MEM_COMMIT,
			NT_PAGE_READWRITE)))
		return status;

	__ntapi->tt_aligned_block_memset(
		rtblock.addr,0,rtblock.size);

	__ntapi->tt_aligned_block_memcpy(
		(uintptr_t *)(rdata = (nt_runtime_data *)rtblock.addr),
		(const uintptr_t *)rtctx,
		sizeof(*rtctx));

	/* abi */
	if (!(__ntapi->tt_guid_compare(&rdata->abi,&(nt_guid)NT_PROCESS_GUID_UNSPEC)))
		__ntapi->tt_guid_copy(
			&rdata->abi,
			&(nt_guid)NT_PROCESS_GUID_RTDATA);

	/* imgbuf */
	imgbuf  = (wchar16_t *)rtblock.addr;
	imgbuf += rtblock.size / sizeof(*imgbuf);
	imgbuf -= __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_IMGBUF_SIZE / sizeof(*imgbuf);

	/* hfile */
	if (sparams->himage) {
		hfile    = sparams->himage;
		himgfile = 0;

	} else if ((status = __ntapi_tt_open_file_utf8(
			&hfile,hat,patharg,1,
			imgbuf,__SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_IMGBUF_SIZE))) {
		return __tt_spawn_return(
			&rtblock,0,0,0,
			status);

		himgfile = hfile;
	}

	/* imgname */
	if ((status = __ntapi->zw_query_object(
			hfile,
			NT_OBJECT_NAME_INFORMATION,
			imgbuf,__SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_IMGBUF_SIZE,
			&(uint32_t){0})))
		return __tt_spawn_return(
			&rtblock,himgfile,0,0,
			status);

	imgname = (nt_unicode_string *)imgbuf;

	/* imgbuf must remain intact until after creation of the child process */
	buflen  = rtblock.size;
	buflen -= sizeof(*rdata);
	buflen -= __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_IMGBUF_SIZE;

	/* argv, envp */
	if ((status = __ntapi->tt_array_copy_utf8(
			&rdata->argc,
			(const char **)sparams->argv,
			(const char **)sparams->envp,
			sparams->interp,
			sparams->optarg,
			sparams->script,
			rtblock.addr,
			rdata->buffer,
			buflen,&written)))
		return __tt_spawn_return(
			&rtblock,himgfile,0,0,
			status);

	rdata->argv   = (char **)&((nt_runtime_data *)0)->buffer;
	rdata->envp   = rdata->argv + rdata->argc + 1;

	rdata->wargv  = (wchar16_t **)rdata->buffer;
	rdata->wargv += written / sizeof(wchar16_t **) + 1;
	rdata->wenvp  = rdata->wargv + rdata->argc + 1;

	rargv = rdata->argv + ((uintptr_t)rtblock.addr / sizeof(char *));
	renvp = rdata->envp + ((uintptr_t)rtblock.addr / sizeof(char *));

	for (rdata->envc=0, parg=sparams->envp; *parg; parg++)
		rdata->envc++;

	pwarg = rdata->wenvp + rdata->envc + 1;
	wch   = (wchar16_t *)pwarg;

	wcap  = (wchar16_t *)rtblock.addr;
	wcap += rtblock.size / sizeof(wchar16_t);

	buflen  = (wcap - wch) * sizeof(wchar16_t);
	buflen -= __SPAWN_NATIVE_PROCESS_RUNTIME_BLOCK_IMGBUF_SIZE;

	if ((status = __ntapi->tt_array_convert_utf8_to_utf16(
			rargv,
			rdata->wargv,
			rdata,wch,
			buflen,&written)))
		return __tt_spawn_return(
			&rtblock,himgfile,0,0,
			status);

	wch    += written / sizeof(wchar16_t);
	buflen -= written;

	if ((status = __ntapi->tt_array_convert_utf8_to_utf16(
			renvp,
			rdata->wenvp,
			rdata,wch,
			buflen,&written)))
		return __tt_spawn_return(
			&rtblock,himgfile,0,0,
			status);

	rdata->wargv -= (uintptr_t)rtblock.addr / sizeof(wchar16_t *);
	rdata->wenvp -= (uintptr_t)rtblock.addr / sizeof(wchar16_t *);

	wenv    = wch;
	wch    += written / sizeof(wchar16_t);
	buflen -= written;

	/* w32 environment */
	if (rtctx->w32_envp) {
		addr  = (uintptr_t)wch;
		addr += sizeof(uintptr_t) - 1;
		addr /= sizeof(uintptr_t);
		addr *= sizeof(uintptr_t);

		penv   = rtctx->w32_envp;
		needed = sizeof(char *);

		for (envc=0; *penv; envc++) {
			needed += sizeof(char *);
			needed += __ntapi->tt_string_null_offset_multibyte(*penv);
			penv++;
		}

		needed += sizeof(uintptr_t) - 1;
		needed /= sizeof(uintptr_t);
		needed *= sizeof(uintptr_t);

		if (buflen < needed)
			return __tt_spawn_return(
				&rtblock,himgfile,0,0,
				NT_STATUS_BUFFER_TOO_SMALL);

		rdata->w32_envp = (char **)(addr - (uintptr_t)rtblock.addr);

		pref = (char **)addr;
		dst  = (char *)&pref[++envc];


		for (penv=rtctx->w32_envp; *penv; penv++) {
			*pref++ = dst - (uintptr_t)rtblock.addr;

			for (src=*penv; *src; src++,dst++)
				*dst = *src;

			*dst++ = 0;
		}

		wch    += needed / sizeof(wchar16_t);
		buflen -= needed;
	}

	/* session */
	if (sparams->hready) {
		if ((status = __ntapi->zw_duplicate_object(
				NT_CURRENT_PROCESS_HANDLE,
				sparams->hready,
				NT_CURRENT_PROCESS_HANDLE,
				&rdata->hready,
				0,0,
				NT_DUPLICATE_SAME_ACCESS|NT_DUPLICATE_SAME_ATTRIBUTES)))
			return __tt_spawn_return(
				&rtblock,himgfile,0,0,
				status);
	} else {
		if ((status = __ntapi->tt_create_inheritable_event(
				&rdata->hready,
				NT_NOTIFICATION_EVENT,
				NT_EVENT_NOT_SIGNALED)))
			return __tt_spawn_return(
				&rtblock,himgfile,0,0,
				status);
	}

	/* process flags */
	if (sparams->processflags & NT_PROCESS_CREATE_FLAGS_CREATE_THREAD_SUSPENDED)
		fsuspended = NT_CREATE_SUSPENDED;

	else if (sparams->threadflags & NT_CREATE_SUSPENDED)
		fsuspended = NT_CREATE_SUSPENDED;

	else if (sparams->spawnflags & NT_PROCESS_SPAWN_FLAG_DEBUG_SUSPENDED)
		fsuspended = NT_CREATE_SUSPENDED;

	else if (sparams->spawnflags & NT_PROCESS_SPAWN_FLAG_DEBUG_EXECUTION)
		fsuspended = NT_CREATE_SUSPENDED;

	/* cparams */
	__ntapi->tt_aligned_block_memset(
		&cparams,0,sizeof(cparams));

	__ntapi->tt_generic_memcpy(
		&crtblock,&rtblock,sizeof(rtblock));

	cparams.image_name		= imgname->buffer;
	cparams.creation_flags_process	= NT_PROCESS_CREATE_FLAGS_INHERIT_HANDLES;
	cparams.creation_flags_thread	= NT_PROCESS_CREATE_FLAGS_CREATE_THREAD_SUSPENDED;
	cparams.environment		= wenv;

	/* crtblock: alloc size: _needed_only_, round up to system granularity */
	crtblock.size    = (size_t)wch - (size_t)rdata;
	crtblock.size   += 0xFFFF;
	crtblock.size   |= 0xFFFF;
	crtblock.size   ^= 0xFFFF;
	cparams.rtblock  = &crtblock;

	/* hoppla */
	if ((status = __ntapi->tt_create_native_process(&cparams)))
		return __tt_spawn_return(
			&rtblock,himgfile,0,0,
			status);

	/* debug */
	if (sparams->spawnflags & NT_PROCESS_SPAWN_FLAG_DEBUG_MASK)
		if ((status = __ntapi->tt_debug_create_attach_object(
				&sparams->hdbgobj,
				cparams.hprocess,
				NT_DEBUG_KILL_ON_EXIT)))
			return __tt_spawn_return(
				&rtblock,
				himgfile,
				cparams.hprocess,
				cparams.hthread,
				status);

	/* additional context */
	if (rtctx->ctx_addr) {
		rdata->ctx_addr   = 0;
		rdata->ctx_commit = rtctx->ctx_size;

		rdata->ctx_commit += (__NT_INTERNAL_PAGE_SIZE - 1);
		rdata->ctx_commit |= (__NT_INTERNAL_PAGE_SIZE - 1);
		rdata->ctx_commit ^= (__NT_INTERNAL_PAGE_SIZE - 1);

		if ((status = __ntapi->zw_allocate_virtual_memory(
				cparams.hprocess,
				&rdata->ctx_addr,0,
				&rdata->ctx_commit,
				NT_MEM_COMMIT,
				NT_PAGE_READWRITE)))
			return __tt_spawn_return(
				&rtblock,
				himgfile,
				cparams.hprocess,
				cparams.hthread,
				status);

		if ((status = __ntapi->zw_write_virtual_memory(
				cparams.hprocess,
				rdata->ctx_addr,
				rtctx->ctx_addr,
				rtctx->ctx_size,
				&rdata->ctx_size)))
			return __tt_spawn_return(
				&rtblock,
				himgfile,
				cparams.hprocess,
				cparams.hthread,
				status);

		raddr  = crtblock.remote_addr;
		raddr += __offsetof(nt_runtime_data,ctx_addr);

		rsize  = __offsetof(nt_runtime_data,ctx_offset);
		rsize -= __offsetof(nt_runtime_data,ctx_addr);

		if ((status = __ntapi->zw_write_virtual_memory(
				cparams.hprocess,
				raddr,(char *)&rdata->ctx_addr,
				rsize,&rsize)))
			return __tt_spawn_return(
				&rtblock,
				himgfile,
				cparams.hprocess,
				cparams.hthread,
				status);
	}

	/* tty session (optional) */
	if (sparams->hsession) {
		if ((status = __ntapi->tty_client_process_register(
				sparams->hsession,
				cparams.pbi.unique_process_id,
				0,NT_TTY_INHERIT_HANDLES,0)))
			return __tt_spawn_return(
				&rtblock,
				himgfile,
				cparams.hprocess,
				cparams.hthread,
				status);

		session.pid    = rtctx->alt_cid_self.pid;
		session.pgid   = rtctx->alt_cid_self.pgid;
		session.sid    = rtctx->alt_cid_self.sid;
		session.syspid = (uint32_t)cparams.pbi.unique_process_id;

		if ((status = __ntapi->tty_client_session_set(
				sparams->hsession,
				&session)))
			return __tt_spawn_return(
				&rtblock,
				himgfile,
				cparams.hprocess,
				cparams.hthread,
				status);
	}

	/* output */
	sparams->hprocess = cparams.hprocess;
	sparams->hthread  = cparams.hthread;
	sparams->rdata    = crtblock.remote_addr;

	sparams->cid.process_id = cparams.pbi.unique_process_id;
	sparams->cid.thread_id  = cparams.cid.thread_id;

	__ntapi->tt_generic_memcpy(
		&sparams->pbi,
		&cparams.pbi,
		sizeof(nt_pbi));

	/* create suspended? */
	if (fsuspended)
		return __tt_spawn_return(
			&rtblock,himgfile,0,0,NT_STATUS_SUCCESS);

	/* tada */
	if ((status = __ntapi->zw_resume_thread(cparams.hthread,0)))
		return __tt_spawn_return(
			&rtblock,
			himgfile,
			cparams.hprocess,
			cparams.hthread,
			status);

	/* hready */
	hchild[1] = cparams.hprocess;
	hchild[0] = rdata->hready;

	__ntapi->zw_wait_for_multiple_objects(
		2,hchild,
		NT_WAIT_ANY,
		NT_SYNC_NON_ALERTABLE,
		sparams->timeout);

	if ((status = __ntapi->zw_query_event(
			rdata->hready,
			NT_EVENT_BASIC_INFORMATION,
			&sparams->eready,
			sizeof(sparams->eready),
			&(size_t){0})))
		return __tt_spawn_return(
			&rtblock,
			himgfile,
			cparams.hprocess,
			cparams.hthread,
			status);

	/* all done */
	return __tt_spawn_return(
		&rtblock,himgfile,0,0,NT_STATUS_SUCCESS);
}