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

/*  toksvc: a framework-native token broker service      */

/*  Copyright (C) 2020  SysDeer Technologies, LLC        */

/*  Released under GPLv2 and GPLv3; see COPYING.TOKSVC.  */

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

#include <psxtypes/psxtypes.h>

#include <ntapi/ntapi.h>

#include <toksvc/toksvc.h>

#include "toksvc_daemon_impl.h"

#include "toksvc_driver_impl.h"

#include "toksvc_log_impl.h"

#define  TOKS_VTBL_ELEMENTS   TOKS_DAEMON_OPCODE_CAP - TOKS_DAEMON_OPCODE_BASE

#define  TOKS_OPCODE_IDX(IDX) TOKS_DAEMON_ ## IDX - TOKS_DAEMON_OPCODE_BASE

#define  TOKS_HANDLER(IDX,fn) [TOKS_OPCODE_IDX(IDX)] = fn

static toks_daemon_routine * toks_daemon_vtbl[TOKS_VTBL_ELEMENTS] = {

	TOKS_HANDLER(CONNECT,       toks_daemon_connect),

	TOKS_HANDLER(CIDQUERY,      toks_daemon_cidquery),

	TOKS_HANDLER(TTYSIGNAL,     toks_daemon_signal),

	TOKS_HANDLER(ACQUIRE,       toks_daemon_acquire),

	TOKS_HANDLER(RELEASE,       toks_daemon_release),

	TOKS_HANDLER(CANCEL,        toks_daemon_cancel),

	TOKS_HANDLER(ABORT,         toks_daemon_abort),

	TOKS_HANDLER(IOCTL,         toks_daemon_ioctl),

};

static toks_daemon_routine * toks_client_vtbl[TOKS_VTBL_ELEMENTS] = {

	TOKS_HANDLER(CONNECT,       toks_daemon_connect),

	TOKS_HANDLER(CIDQUERY,      toks_daemon_cidquery),

	TOKS_HANDLER(TTYSIGNAL,     toks_daemon_signal),

};

static const nt_guid    g_pidmap    = NT_PROCESS_GUID_PIDMAP;

static const wchar16_t  p_pidmap[6] = NT_PROCESS_OBJDIR_PREFIX_PIDMAP;

static const nt_guid    g_ntpipc    = NT_PROCESS_GUID_NTPIPC;

static const wchar16_t  p_ntpipc[6] = NT_PROCESS_OBJDIR_PREFIX_NTPIPC;

static void toks_daemon_ctrlpid_abort(struct toks_daemon_ctx * dctx)

{

	void *			hport;

	struct _nt_tty_sync_msg	msg;

	hport = dctx->hport_internal_client;

	ntapi->tt_aligned_block_memset(

		&msg,0,sizeof(msg));

	msg.header.msg_type	= NT_LPC_NEW_MESSAGE;

	msg.header.data_size	= sizeof(msg.data);

	msg.header.msg_size	= sizeof(msg);

	msg.data.ttyinfo.opcode	= TOKS_DAEMON_ABORT;

	ntapi->zw_request_wait_reply_port(

		hport,&msg,&msg;;

}

static int32_t toks_daemon_ctrlpid_wait(void * rapunzel)

{

	struct toks_daemon_ctx * dctx;

	dctx = (struct toks_daemon_ctx *)rapunzel;

	ntapi->zw_set_event(

		dctx->hswap,0);

	ntapi->zw_wait_for_single_object(

		dctx->hctrl,

		NT_SYNC_NON_ALERTABLE,

		0);

	toks_daemon_ctrlpid_abort(dctx);

	return ntapi->zw_terminate_thread(

		NT_CURRENT_THREAD_HANDLE,

		NT_STATUS_REMOTE_DISCONNECT);

}

static int32_t toks_daemon_ctrlpid_instance(struct toks_daemon_ctx * dctx)

{

	int32_t			status;

	nt_thread_params	params;

	nt_cid			cid;

	nt_oa			oa;

	if ((dctx->ctrlpid == 0) && (dctx->csyspid == 0))

		return NT_STATUS_SUCCESS;

	if (dctx->ctrlpid) {

		dctx->reqtokpid = dctx->ctrlpid;

		if ((status = toks_daemon_pidopen(dctx)))

			return status;

		dctx->csyspid = dctx->reqsyspid;

	}

	oa.len      = sizeof(oa);

	oa.root_dir = 0;

	oa.obj_name = 0;

	oa.obj_attr = 0;

	oa.sec_desc = 0;

	oa.sec_qos  = 0;

	cid.process_id = dctx->csyspid;

	cid.thread_id  = 0;

	if ((status = ntapi->zw_open_process(

			&dctx->hctrl,

			NT_PROCESS_SYNCHRONIZE,

			&oa,&cid)))

		return status;

	if ((status = ntapi->tt_create_private_event(

			&dctx->hswap,

			NT_NOTIFICATION_EVENT,

			NT_EVENT_NOT_SIGNALED)))

		return status;

	ntapi->tt_aligned_block_memset(

		&params,0,sizeof(params));

	params.hprocess		  = NT_CURRENT_PROCESS_HANDLE;

	params.start		  = toks_daemon_ctrlpid_wait;

	params.arg                = dctx;

	params.stack_size_commit  = 4 * 1024;

	params.stack_size_reserve = 4 * 1024;

	params.creation_flags	  = NT_CREATE_LOCAL_THREAD;

	if ((status = ntapi->tt_create_thread(&params))) {

		ntapi->zw_close(dctx->hctrl);

		ntapi->zw_close(dctx->hswap);

		return status;

	}

	status = ntapi->zw_wait_for_single_object(

		dctx->hswap,

		NT_SYNC_NON_ALERTABLE,

		0);

	ntapi->zw_close(dctx->hswap);

	ntapi->zw_close(params.hthread);

	if (status) {

		ntapi->zw_close(dctx->hctrl);

		return status;

	}

	return NT_STATUS_SUCCESS;

}

static int toks_server_init_pidmap_object_directory(nt_rtdata * rtdata)

{

	return rtdata->hpidmapdir

		? NT_STATUS_SUCCESS

		: ntapi->tt_open_ipc_object_directory(

			&rtdata->hpidmapdir,

			NT_SEC_READ_CONTROL

				| NT_DIRECTORY_QUERY

				| NT_DIRECTORY_TRAVERSE

				| NT_DIRECTORY_CREATE_OBJECT

				| NT_DIRECTORY_CREATE_SUBDIRECTORY,

			p_pidmap,&g_pidmap);

}

static int toks_server_init_ntpipc_object_directory(nt_rtdata * rtdata)

{

	return rtdata->hntpipcdir

		? NT_STATUS_SUCCESS

		: ntapi->tt_open_ipc_object_directory(

			&rtdata->hntpipcdir,

			NT_SEC_READ_CONTROL

				| NT_DIRECTORY_QUERY

				| NT_DIRECTORY_TRAVERSE

				| NT_DIRECTORY_CREATE_OBJECT

				| NT_DIRECTORY_CREATE_SUBDIRECTORY,

			p_ntpipc,&g_ntpipc);

}

static int toks_init_pidmap_target_symlink(struct toks_daemon_ctx * dctx)

{

	int     status;

	void *  hkeydir;

	struct toks_driver_ctx_impl * ictx = toks_get_driver_ictx(dctx->driver_ctx);

	if ((status = ntapi->tt_create_keyed_object_directory(

			&hkeydir,

			NT_SYMBOLIC_LINK_ALL_ACCESS,

			ictx->rtdata->hpidmapdir,

			pe_get_current_process_id())))

		return status;

	if ((status = ntapi->zw_set_information_object(

			hkeydir,

			NT_OBJECT_HANDLE_INFORMATION,

			&(nt_object_handle_information){0,0},

			sizeof(nt_object_handle_information))))

		return status;

	return ntapi->tt_create_keyed_object_directory_entry(

		&ictx->hntpipc,

		NT_SYMBOLIC_LINK_ALL_ACCESS,

		hkeydir,

		dctx->hport_daemon,0,

		pe_get_current_process_id());

}

static int toks_init_ntpipc_target_symlink(struct toks_daemon_ctx * dctx)

{

	struct toks_driver_ctx_impl * ictx = toks_get_driver_ictx(dctx->driver_ctx);

	return ntapi->tt_create_keyed_object_directory_entry(

		&ictx->hntpipc,

		NT_SYMBOLIC_LINK_ALL_ACCESS,

		ictx->rtdata->hntpipcdir,

		dctx->hport_daemon,0,

		pe_get_current_process_id());

}

int32_t __stdcall toks_daemon_loop(void * ctx)

{

	struct toks_daemon_ctx *	dctx;

	toks_daemon_routine **		svcvtbl;

	nt_rtdata *			rtdata;

	nt_tty_port_msg *		request;

	nt_tty_port_msg *		reply;

	intptr_t			port_id;

	int32_t				opcode;

	int32_t				status;

	/* runtime data */

	if (ntapi->tt_get_runtime_data(&rtdata,0))

		return NT_STATUS_INTERNAL_ERROR;

	/* daemon context and controlling (system) process */

	dctx = (struct toks_daemon_ctx *)ctx;

	dctx->ctrlpid = toks_get_driver_ctrlpid(dctx->driver_ctx);

	dctx->csyspid = toks_get_driver_csyspid(dctx->driver_ctx);

	dctx->ftokens = toks_get_driver_ntokens(dctx->driver_ctx);

	svcvtbl = (dctx->driver_ctx->cctx->drvflags & TOKS_DRIVER_MODE_SERVER)

		? toks_daemon_vtbl

		: toks_client_vtbl;

	if (svcvtbl == toks_daemon_vtbl)

		if ((status = toks_daemon_ctrlpid_instance(dctx)))

			ntapi->zw_terminate_process(

				NT_CURRENT_PROCESS_HANDLE,

				status);

	/* service synchronization */

	if (rtdata->hsync) {

		ntapi->zw_set_event(rtdata->hsync,0);

		ntapi->zw_close(rtdata->hsync);

	}

	/* service info */

	if (svcvtbl == toks_daemon_vtbl)

		toks_log_service_info(dctx);

	/* pidmap object directory */

	if ((status = toks_server_init_pidmap_object_directory(rtdata)))

		return status;

	/* pidmap daemon symlink */

	if ((status = toks_init_pidmap_target_symlink(dctx)))

		return status;

	/* ntpipc object directory */

	if ((status = toks_server_init_ntpipc_object_directory(rtdata)))

		return status;

	/* ntpipc daemon symlink */

	if ((status = toks_init_ntpipc_target_symlink(dctx)))

		return status;

	/* init */

	request = &dctx->request;

	reply   = &dctx->reply;

	ntapi->tt_aligned_block_memset(

		request,0,sizeof(*request));

	/* get first message */

	ntapi->zw_reply_wait_receive_port(

		dctx->hport_daemon,

		&port_id,

		0,(nt_port_message *)request);

	/* message loop */

	do {

		switch (request->header.msg_type) {

			case NT_LPC_REQUEST:

			case NT_LPC_DATAGRAM:

				opcode = request->ttyinfo.opcode;

				break;

			case NT_LPC_CONNECTION_REQUEST:

				opcode = TOKS_DAEMON_CONNECT;

				break;

			default:

				opcode = -1;

				break;

		}

		/* dispatch */

		dctx->opcode = opcode;

		dctx->noise++;

		request->syncinfo.ipcsvc.keys.padding = 0;

		if (svcvtbl == toks_daemon_vtbl)

			toks_log_lpc_request(

				dctx,request);

		ntapi->tt_aligned_block_memcpy(

			(uintptr_t *)reply,

			(uintptr_t *)request,

			sizeof(*reply));

		reply->header.msg_type = NT_LPC_REPLY;

		switch (opcode) {

			case TOKS_DAEMON_IOCTL:

				break;

			default:

				reply->ttyinfo.opdata = reply->header.msg_id;

		}

		if ((opcode >= TOKS_DAEMON_OPCODE_BASE) && (opcode < TOKS_DAEMON_OPCODE_CAP)) {

			opcode -= TOKS_DAEMON_OPCODE_BASE;

			reply->ttyinfo.status = svcvtbl[opcode]

				? svcvtbl[opcode](dctx)

				: NT_STATUS_NOT_IMPLEMENTED;

		} else {

			reply->ttyinfo.status   = NT_STATUS_LPC_INVALID_CONNECTION_USAGE;

		}

		if (svcvtbl == toks_daemon_vtbl)

			toks_log_lpc_reply(

				dctx,reply);

		ntapi->tt_aligned_block_memset(

			request,0,sizeof(*request));

		if (!reply->ttyinfo.opdata)

			ntapi->zw_reply_wait_receive_port(

				dctx->hport_daemon,

				&port_id,

				0,&request->header);

		else if (reply->header.client_id.process_id == rtdata->cid_self.process_id)

			ntapi->zw_reply_wait_receive_port(

				dctx->hport_daemon,

				&port_id,

				&reply->header,

				&request->header);

		else {

			ntapi->zw_reply_port(

				dctx->hport_daemon,

				&reply->header);

			ntapi->zw_reply_wait_receive_port(

				dctx->hport_daemon,

				&port_id,

				0,&request->header);

		}

	} while (request->header.msg_id);

	return NT_STATUS_INTERNAL_ERROR;

}