2 files changed, 53 insertions, 64 deletions

diff --git a/src/daemon/ntapi_dsr_init.c b/src/daemon/ntapi_dsr_init.c
index be72502..58888ed 100644
--- a/src/daemon/ntapi_dsr_init.c
+++ b/src/daemon/ntapi_dsr_init.c
@@ -11,7 +11,7 @@
 #include <ntapi/ntapi.h>
 #include "ntapi_impl.h"
 
-static void __stdcall __ntapi_dsr_once(nt_daemon_params * params);
+static int32_t __stdcall __ntapi_dsr_once(nt_daemon_params * params);
 
 int32_t __stdcall __ntapi_dsr_init(nt_daemon_params * params)
 {
@@ -120,29 +120,45 @@ int32_t __stdcall __ntapi_dsr_init(nt_daemon_params * params)
 
 
 /* __ntapi_dsr_start executes in the daemon's dedicated thread */
-int32_t __stdcall __ntapi_dsr_start(nt_daemon_params * params)
+int32_t __stdcall __ntapi_dsr_start_impl(nt_daemon_params * params)
 {
-	__ntapi_dsr_once(params);
-	__ntapi_dsr_create_port(params);
-	__ntapi_dsr_connect_internal_client(params);
-	params->daemon_loop_routine(params->daemon_loop_context);
+	int32_t status;
+
+	if ((status = __ntapi_dsr_once(params)))
+		return status;
+
+	if ((status = __ntapi_dsr_create_port(params)))
+		return status;
+
+	if ((status = __ntapi_dsr_connect_internal_client(params)))
+		return status;
+
+	if ((status = params->daemon_loop_routine(params->daemon_loop_context)))
+		return status;
+
+	return NT_STATUS_SUCCESS;
+}
 
-	/* (no return) */
-	return NT_STATUS_INTERNAL_ERROR;
+int32_t __stdcall __ntapi_dsr_start(nt_daemon_params * params)
+{
+	return __ntapi->zw_terminate_thread(
+		NT_CURRENT_THREAD_HANDLE,
+		__ntapi_dsr_start_impl(params));
 }
 
+
 /* __ntapi_dsr_once executes in the daemon's dedicated thread */
-static void __stdcall __ntapi_dsr_once(nt_daemon_params * params)
+static int32_t __stdcall __ntapi_dsr_once(nt_daemon_params * params)
 {
 	int32_t status;
 
 	if (!params->daemon_once_routine)
-		return;
+		return NT_STATUS_SUCCESS;
 
-	if ((status = params->daemon_once_routine(params->daemon_loop_context))) {
+	if ((status = params->daemon_once_routine(params->daemon_loop_context)))
 		params->exit_code_daemon_start = status;
-		__ntapi->zw_terminate_thread(NT_CURRENT_THREAD_HANDLE,status);
-	}
+
+	return status;
 }
 
 /* __ntapi_dsr_create_port executes in the daemon's dedicated thread */
@@ -175,15 +191,11 @@ int32_t __stdcall __ntapi_dsr_create_port(nt_daemon_params * params)
 	oa.sec_qos	= &sqos;
 
 	/* create the port */
-	*pstatus = __ntapi->zw_create_port(
-		&params->hport_daemon,
-		&oa,0,(uint32_t)params->port_msg_size,
-		0);
-
-	if (*pstatus != NT_STATUS_SUCCESS)
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			*pstatus);
+	if ((*pstatus = __ntapi->zw_create_port(
+			&params->hport_daemon,
+			&oa,0,(uint32_t)params->port_msg_size,
+			0)))
+		return *pstatus;
 
 	/* return port info */
 	if (params->pport_daemon)
@@ -194,10 +206,5 @@ int32_t __stdcall __ntapi_dsr_create_port(nt_daemon_params * params)
 		params->hevent_daemon_ready,
 		(int32_t *)0);
 
-	if (*pstatus != NT_STATUS_SUCCESS)
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			*pstatus);
-
 	return *pstatus;
 }
diff --git a/src/daemon/ntapi_dsr_internal_connection.c b/src/daemon/ntapi_dsr_internal_connection.c
index 1ff0846..df80d46 100644
--- a/src/daemon/ntapi_dsr_internal_connection.c
+++ b/src/daemon/ntapi_dsr_internal_connection.c
@@ -15,7 +15,6 @@
 int32_t __stdcall __ntapi_dsr_connect_internal_client(nt_daemon_params * params)
 {
 	int32_t *		pstatus;
-
 	intptr_t		port_id;
 	nt_port_message		port_msg;
 	nt_large_integer	timeout;
@@ -26,45 +25,31 @@ int32_t __stdcall __ntapi_dsr_connect_internal_client(nt_daemon_params * params)
 	/* timeout-enabled first connection */
 	timeout.quad = NT_DSR_INIT_MAX_WAIT;
 
-	*pstatus = __ntapi->zw_reply_wait_receive_port_ex(
-		params->hport_daemon,
-		&port_id,
-		(nt_port_message *)0,
-		(nt_port_message *)&port_msg,
-		&timeout);
-
-	if (*pstatus != NT_STATUS_SUCCESS)
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			*pstatus);
+	if ((*pstatus = __ntapi->zw_reply_wait_receive_port_ex(
+			params->hport_daemon,
+			&port_id,
+			(nt_port_message *)0,
+			(nt_port_message *)&port_msg,
+			&timeout)))
+		return *pstatus;
 
 	/* the internal client must be first */
 	if (port_msg.client_id.process_id != pe_get_current_process_id())
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			NT_STATUS_PORT_CONNECTION_REFUSED);
+		return NT_STATUS_PORT_CONNECTION_REFUSED;
 
 	/* accept connection request */
-	*pstatus = __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);
-
-	if (*pstatus != NT_STATUS_SUCCESS)
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			*pstatus);
+	if ((*pstatus = __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;
 
 	/* finalize connection */
-	*pstatus = __ntapi->zw_complete_connect_port(_hport_client);
-
-	if (*pstatus != NT_STATUS_SUCCESS)
-		__ntapi->zw_terminate_thread(
-			NT_CURRENT_THREAD_HANDLE,
-			*pstatus);
+	*pstatus = __ntapi->zw_complete_connect_port(
+		_hport_client);
 
 	return *pstatus;
 }
@@ -122,10 +107,7 @@ int32_t __stdcall __ntapi_dsr_internal_client_connect(nt_daemon_params * params)
 		0);
 
 	/* exit the task-specific thread */
-	__ntapi->zw_terminate_thread(
+	return __ntapi->zw_terminate_thread(
 		NT_CURRENT_THREAD_HANDLE,
 		*pstatus);
-
-	/* (no return) */
-	return NT_STATUS_INTERNAL_ERROR;
 }