daemon service routine: set daemon status using atomic updates.

2 files changed, 101 insertions, 34 deletions

diff --git a/src/daemon/ntapi_dsr_init.c b/src/daemon/ntapi_dsr_init.c
index 58888ed..f028efd 100644
--- a/src/daemon/ntapi_dsr_init.c
+++ b/src/daemon/ntapi_dsr_init.c
@@ -4,6 +4,7 @@
 /*  Released under GPLv2 and GPLv3; see COPYING.NTAPI.  */
 /********************************************************/
 
+#include <ntapi/nt_atomic.h>
 #include <ntapi/nt_status.h>
 #include <ntapi/nt_thread.h>
 #include <ntapi/nt_port.h>
@@ -20,6 +21,19 @@ int32_t __stdcall __ntapi_dsr_init(nt_daemon_params * params)
 	nt_thread_params	tparams;
 	nt_large_integer	timeout;
 
+	/* report status */
+	at_store_32(
+		&params->exit_code_daemon_loop,
+		NT_STATUS_PENDING);
+
+	at_store_32(
+		&params->exit_code_daemon_start,
+		NT_STATUS_PENDING);
+
+	at_store_32(
+		&params->exit_code_internal_client,
+		NT_STATUS_PENDING);
+
 	/* port_keys */
 	if (params->flags & NT_DSR_INIT_GENERATE_KEYS)
 		if ((status = __ntapi->tt_port_generate_keys(params->port_keys)))
@@ -141,6 +155,10 @@ int32_t __stdcall __ntapi_dsr_start_impl(nt_daemon_params * params)
 
 int32_t __stdcall __ntapi_dsr_start(nt_daemon_params * params)
 {
+	at_store_32(
+		&params->exit_code_daemon_loop,
+		NT_STATUS_MORE_PROCESSING_REQUIRED);
+
 	return __ntapi->zw_terminate_thread(
 		NT_CURRENT_THREAD_HANDLE,
 		__ntapi_dsr_start_impl(params));
@@ -156,21 +174,32 @@ static int32_t __stdcall __ntapi_dsr_once(nt_daemon_params * params)
 		return NT_STATUS_SUCCESS;
 
 	if ((status = params->daemon_once_routine(params->daemon_loop_context)))
-		params->exit_code_daemon_start = status;
+		at_store_32(
+			&params->exit_code_daemon_start,
+			status);
 
 	return status;
 }
 
 /* __ntapi_dsr_create_port executes in the daemon's dedicated thread */
+static int32_t __stdcall __ntapi_dsr_create_port_exit(
+	nt_daemon_params *	params,
+	int32_t			status)
+{
+	at_store_32(
+		&params->exit_code_daemon_start,
+		status);
+
+	return status;
+}
+
 int32_t __stdcall __ntapi_dsr_create_port(nt_daemon_params * params)
 {
-	int32_t *			pstatus;
+	int32_t				status;
 	nt_object_attributes		oa;
 	nt_security_quality_of_service	sqos;
 	nt_unicode_string		server_name;
 
-	pstatus = &params->exit_code_daemon_start;
-
 	/* init server_name */
 	server_name.strlen = (uint16_t)__ntapi->tt_string_null_offset_short((const int16_t *)params->port_name);
 	server_name.maxlen = 0;
@@ -191,20 +220,24 @@ int32_t __stdcall __ntapi_dsr_create_port(nt_daemon_params * params)
 	oa.sec_qos	= &sqos;
 
 	/* create the port */
-	if ((*pstatus = __ntapi->zw_create_port(
+	if ((status = __ntapi->zw_create_port(
 			&params->hport_daemon,
 			&oa,0,(uint32_t)params->port_msg_size,
 			0)))
-		return *pstatus;
+		return __ntapi_dsr_create_port_exit(
+			params,status);
 
 	/* return port info */
 	if (params->pport_daemon)
-		*(params->pport_daemon) = params->hport_daemon;
+		at_store(
+			(intptr_t *)params->pport_daemon,
+			(intptr_t)params->hport_daemon);
 
 	/* signal the daemon-is-ready event */
-	*pstatus = __ntapi->zw_set_event(
+	status = __ntapi->zw_set_event(
 		params->hevent_daemon_ready,
 		(int32_t *)0);
 
-	return *pstatus;
+	return __ntapi_dsr_create_port_exit(
+		params,status);
 }
diff --git a/src/daemon/ntapi_dsr_internal_connection.c b/src/daemon/ntapi_dsr_internal_connection.c
index df80d46..7246eff 100644
--- a/src/daemon/ntapi_dsr_internal_connection.c
+++ b/src/daemon/ntapi_dsr_internal_connection.c
@@ -4,6 +4,7 @@
 /*  Released under GPLv2 and GPLv3; see COPYING.NTAPI.  */
 /********************************************************/
 
+#include <ntapi/nt_atomic.h>
 #include <ntapi/nt_status.h>
 #include <ntapi/nt_thread.h>
 #include <ntapi/nt_port.h>
@@ -11,60 +12,97 @@
 #include <ntapi/ntapi.h>
 #include "ntapi_impl.h"
 
+static int32_t __stdcall __ntapi_dsr_connect_internal_client_exit(
+	nt_daemon_params *	params,
+	int32_t			status)
+{
+	at_store_32(
+		&params->exit_code_daemon_start,
+		status);
+
+	if (status == NT_STATUS_SUCCESS)
+		at_store_32(
+			&params->exit_code_daemon_loop,
+			NT_STATUS_SERVICE_NOTIFICATION);
+
+	return status;
+}
+
+static int32_t __stdcall __ntapi_dsr_internal_client_connect_exit(
+	nt_daemon_params *	params,
+	int32_t			status)
+{
+	at_store_32(
+		&params->exit_code_internal_client,
+		status);
+
+	return __ntapi->zw_terminate_thread(
+		NT_CURRENT_THREAD_HANDLE,
+		status);
+}
+
 /* __ntapi_dsr_connect_internal_client executes in the daemon's dedicated thread */
 int32_t __stdcall __ntapi_dsr_connect_internal_client(nt_daemon_params * params)
 {
-	int32_t *		pstatus;
+	int32_t			status;
 	intptr_t		port_id;
 	nt_port_message		port_msg;
 	nt_large_integer	timeout;
 	void *			_hport_client;
 
-	pstatus = &params->exit_code_daemon_start;
+	/* report state */
+	at_store_32(
+		&params->exit_code_daemon_start,
+		NT_STATUS_MORE_PROCESSING_REQUIRED);
 
 	/* timeout-enabled first connection */
 	timeout.quad = NT_DSR_INIT_MAX_WAIT;
 
-	if ((*pstatus = __ntapi->zw_reply_wait_receive_port_ex(
+	if ((status = __ntapi->zw_reply_wait_receive_port_ex(
 			params->hport_daemon,
 			&port_id,
 			(nt_port_message *)0,
 			(nt_port_message *)&port_msg,
 			&timeout)))
-		return *pstatus;
+		return __ntapi_dsr_connect_internal_client_exit(
+			params,status);
 
 	/* the internal client must be first */
 	if (port_msg.client_id.process_id != pe_get_current_process_id())
 		return NT_STATUS_PORT_CONNECTION_REFUSED;
 
 	/* accept connection request */
-	if ((*pstatus = __ntapi->zw_accept_connect_port(
+	if ((status = __ntapi->zw_accept_connect_port(
 			&_hport_client,
 			port_msg.client_id.process_id,
 			(nt_port_message *)&port_msg,
 			NT_LPC_ACCEPT_CONNECTION,
 			(nt_port_section_write *)0,
 			(nt_port_section_read *)0)))
-		return *pstatus;
+		return __ntapi_dsr_connect_internal_client_exit(
+			params,status);
 
 	/* finalize connection */
-	*pstatus = __ntapi->zw_complete_connect_port(
+	status = __ntapi->zw_complete_connect_port(
 		_hport_client);
 
-	return *pstatus;
+	return __ntapi_dsr_connect_internal_client_exit(
+		params,status);
 }
 
 
 /* __ntapi_dsr_internal_client_connect executes in its own temporary thread */
 int32_t __stdcall __ntapi_dsr_internal_client_connect(nt_daemon_params * params)
 {
-	int32_t *			pstatus;
-
+	int32_t				status;
 	nt_unicode_string		server_name;
 	nt_object_attributes		oa;
 	nt_security_quality_of_service	sqos;
 
-	pstatus = &params->exit_code_internal_client;
+	/* report status */
+	at_store_32(
+		&params->exit_code_internal_client,
+		NT_STATUS_MORE_PROCESSING_REQUIRED);
 
 	/* init server_name */
 	server_name.strlen = (uint16_t)__ntapi->tt_string_null_offset_short((const int16_t *)params->port_name);
@@ -86,28 +124,24 @@ int32_t __stdcall __ntapi_dsr_internal_client_connect(nt_daemon_params * params)
 	oa.sec_qos	= &sqos;
 
 	/* establish internal connection */
-	*pstatus = __ntapi->zw_connect_port(
-		&params->hport_internal_client,
-		&server_name,
-		&sqos,
-		0,0,0,0,0);
-
-	if (*pstatus != NT_STATUS_SUCCESS)
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			*pstatus);
+	if ((status = __ntapi->zw_connect_port(
+			&params->hport_internal_client,
+			&server_name,
+			&sqos,
+			0,0,0,0,0)))
+		return __ntapi_dsr_internal_client_connect_exit(
+			params,status);
 
 	/* return port info */
 	if (params->pport_internal_client)
 		*(params->pport_internal_client) = params->hport_internal_client;
 
 	/* signal the 'internal-client-is-ready' event */
-	*pstatus = __ntapi->zw_set_event(
+	status = __ntapi->zw_set_event(
 		params->hevent_internal_client_ready,
 		0);
 
 	/* exit the task-specific thread */
-	return __ntapi->zw_terminate_thread(
-		NT_CURRENT_THREAD_HANDLE,
-		*pstatus);
+	return __ntapi_dsr_internal_client_connect_exit(
+		params,status);
 }