server: proxy: disconnect all active sessions on shutdown

server/proxy/freerdp_proxy.c

@@ -31,18 +31,21 @@
 
 #define TAG PROXY_TAG("server")
 
-static proxyConfig config = { 0 };
+static proxyServer* server = NULL;
 
 static void cleanup_handler(int signum)
 {
 	printf("\n");
 	WLog_INFO(TAG, "[%s]: caught signal %d, starting cleanup...", __FUNCTION__, signum);
 
+	WLog_INFO(TAG, "stopping all connections.");
+	pf_server_stop(server);
+
 	WLog_INFO(TAG, "freeing loaded modules and plugins.");
 	pf_modules_free();
 
-	WLog_INFO(TAG, "freeing config.");
+	pf_server_config_free(server->config);
-	pf_server_config_free_internal(&config);
+	pf_server_free(server);
 
 	WLog_INFO(TAG, "exiting.");
 	exit(signum);
@@ -50,11 +53,12 @@ static void cleanup_handler(int signum)
 
 int main(int argc, char* argv[])
 {
-	const char* cfg = "config.ini";
+	proxyConfig* config = NULL;
-	int status = 0;
+	const char* config_path = "config.ini";
+	int status = -1;
 
 	if (argc > 1)
-		cfg = argv[1];
+		config_path = argv[1];
 
 	/* Register cleanup handler for graceful termination */
 	signal(SIGINT, cleanup_handler);
@@ -72,13 +76,26 @@ int main(int argc, char* argv[])
 
 	pf_modules_list_loaded_plugins();
 
-	if (!pf_server_config_load(cfg, &config))
+	config = pf_server_config_load(config_path);
+	if (!config)
 		goto fail;
 
-	pf_server_config_print(&config);
+	pf_server_config_print(config);
-	status = pf_server_start(&config);
+
+	server = pf_server_new(config);
+	if (!server)
+		goto fail;
+
+	if (!pf_server_start(server))
+		goto fail;
+
+	if (WaitForSingleObject(server->thread, INFINITE) != WAIT_OBJECT_0)
+		goto fail;
+
+	status = 0;
 fail:
+	pf_server_free(server);
 	pf_modules_free();
-	pf_server_config_free_internal(&config);
+	pf_server_config_free(config);
 	return status;
 }

server/proxy/pf_config.c

@@ -242,9 +242,9 @@ static BOOL pf_config_load_captures(wIniFile* ini, proxyConfig* config)
 	return TRUE;
 }
 
-BOOL pf_server_config_load(const char* path, proxyConfig* config)
+proxyConfig* pf_server_config_load(const char* path)
 {
-	BOOL ok = FALSE;
+	proxyConfig* config = NULL;
 	wIniFile* ini = IniFile_New();
 
 	if (!ini)
@@ -259,6 +259,8 @@ BOOL pf_server_config_load(const char* path, proxyConfig* config)
 		goto out;
 	}
 
+	config = calloc(1, sizeof(proxyConfig));
+
 	if (!pf_config_load_server(ini, config))
 		goto out;
 
@@ -283,10 +285,13 @@ BOOL pf_server_config_load(const char* path, proxyConfig* config)
 	if (!pf_config_load_captures(ini, config))
 		goto out;
 
-	ok = TRUE;
+	IniFile_Free(ini);
+	return config;
+
 out:
 	IniFile_Free(ini);
-	return ok;
+	pf_server_config_free(config);
+	return NULL;
 }
 
 void pf_server_config_print(proxyConfig* config)
@@ -336,7 +341,7 @@ void pf_server_config_print(proxyConfig* config)
 	CONFIG_PRINT_STR(config, CapturesDirectory);
 }
 
-void pf_server_config_free_internal(proxyConfig* config)
+void pf_server_config_free(proxyConfig* config)
 {
 	if (config == NULL)
 		return;
@@ -344,4 +349,5 @@ void pf_server_config_free_internal(proxyConfig* config)
 	free(config->CapturesDirectory);
 	free(config->TargetHost);
 	free(config->Host);
+	free(config);
 }

server/proxy/pf_config.h

@@ -77,8 +77,8 @@ FREERDP_API BOOL pf_config_get_uint32(wIniFile* ini, const char* section, const
 FREERDP_API BOOL pf_config_get_bool(wIniFile* ini, const char* section, const char* key);
 FREERDP_API const char* pf_config_get_str(wIniFile* ini, const char* section, const char* key);
 
-BOOL pf_server_config_load(const char* path, proxyConfig* config);
+proxyConfig* pf_server_config_load(const char* path);
 void pf_server_config_print(proxyConfig* config);
-void pf_server_config_free_internal(proxyConfig* config);
+void pf_server_config_free(proxyConfig* config);
 
 #endif /* FREERDP_SERVER_PROXY_PFCONFIG_H */

server/proxy/pf_context.c

@@ -56,11 +56,13 @@ error:
 /* Proxy context free callback */
 static void client_to_proxy_context_free(freerdp_peer* client, pServerContext* context)
 {
-	WINPR_UNUSED(client);
+	proxyServer* server;
 
-	if (!context)
+	if (!client || !context)
 		return;
 
+	server = (proxyServer*)client->ContextExtra;
+
 	WTSCloseServer((HANDLE)context->vcm);
 
 	if (context->dynvcReady)

server/proxy/pf_server.c

@@ -19,27 +19,17 @@
  * limitations under the License.
  */
 
-#include <errno.h>
-#include <signal.h>
-
-#include <freerdp/freerdp.h>
-#include <freerdp/listener.h>
-
 #include <winpr/crt.h>
 #include <winpr/ssl.h>
 #include <winpr/synch.h>
 #include <winpr/string.h>
-#include <winpr/path.h>
 #include <winpr/winsock.h>
 #include <winpr/thread.h>
 
+#include <freerdp/freerdp.h>
 #include <freerdp/channels/wtsvc.h>
 #include <freerdp/channels/channels.h>
 
-#include <freerdp/constants.h>
-#include <freerdp/server/rdpsnd.h>
-#include <freerdp/server/rdpgfx.h>
-
 #include "pf_server.h"
 #include "pf_log.h"
 #include "pf_config.h"
@@ -208,9 +198,10 @@ static BOOL pf_server_adjust_monitor_layout(freerdp_peer* peer)
 static BOOL pf_server_initialize_peer_connection(freerdp_peer* peer)
 {
 	pServerContext* ps = (pServerContext*)peer->context;
+	rdpSettings* settings = peer->settings;
 	proxyData* pdata;
 	proxyConfig* config;
-	rdpSettings* settings = peer->settings;
+	proxyServer* server;
 
 	if (!ps)
 		return FALSE;
@@ -220,7 +211,9 @@ static BOOL pf_server_initialize_peer_connection(freerdp_peer* peer)
 		return FALSE;
 
 	proxy_data_set_server_context(pdata, ps);
-	config = pdata->config = peer->ContextExtra;
+	server = (proxyServer*)peer->ContextExtra;
+
+	config = pdata->config = server->config;
 
 	/* currently not supporting GDI orders */
 	ZeroMemory(settings->OrderSupport, 32);
@@ -262,6 +255,11 @@ static BOOL pf_server_initialize_peer_connection(freerdp_peer* peer)
 	peer->Activate = pf_server_activate;
 	peer->AdjustMonitorsLayout = pf_server_adjust_monitor_layout;
 	peer->settings->MultifragMaxRequestSize = 0xFFFFFF; /* FIXME */
+
+	if (ArrayList_Add(server->clients, pdata) < 0)
+		return FALSE;
+
+	CountdownEvent_AddCount(server->waitGroup, 1);
 	return TRUE;
 }
 
@@ -281,6 +279,7 @@ static DWORD WINAPI pf_server_handle_peer(LPVOID arg)
 	rdpContext* pc;
 	proxyData* pdata;
 	freerdp_peer* client = (freerdp_peer*)arg;
+	proxyServer* server = (proxyServer*)client->ContextExtra;
 
 	if (!pf_context_init_server_context(client))
 		goto out_free_peer;
@@ -317,7 +316,7 @@ static DWORD WINAPI pf_server_handle_peer(LPVOID arg)
 
 		if (status == WAIT_FAILED)
 		{
-			WLog_ERR(TAG, "WaitForMultipleObjects failed (errno: %d)", errno);
+			WLog_ERR(TAG, "WaitForMultipleObjects failed (status: %d)", status);
 			break;
 		}
 
@@ -376,6 +375,7 @@ fail:
 	LOG_INFO(TAG, ps, "freeing server's channels");
 	pf_server_channels_free(ps);
 	LOG_INFO(TAG, ps, "freeing proxy data");
+	ArrayList_Remove(server->clients, pdata);
 	proxy_data_free(pdata);
 	freerdp_client_context_free(pc);
 	client->Close(client);
@@ -383,6 +383,8 @@ fail:
 out_free_peer:
 	freerdp_peer_context_free(client);
 	freerdp_peer_free(client);
+	CountdownEvent_Signal(server->waitGroup, 1);
+	ExitThread(0);
 	return 0;
 }
 
@@ -398,11 +400,13 @@ static BOOL pf_server_peer_accepted(freerdp_listener* listener, freerdp_peer* cl
 	return TRUE;
 }
 
-static void pf_server_mainloop(freerdp_listener* listener)
+static DWORD WINAPI pf_server_mainloop(LPVOID arg)
 {
 	HANDLE eventHandles[32];
 	DWORD eventCount;
 	DWORD status;
+	proxyServer* server = (proxyServer*)arg;
+	freerdp_listener* listener = server->listener;
 
 	while (1)
 	{
@@ -414,8 +418,12 @@ static void pf_server_mainloop(freerdp_listener* listener)
 			break;
 		}
 
+		eventHandles[eventCount++] = server->stopEvent;
 		status = WaitForMultipleObjects(eventCount, eventHandles, FALSE, INFINITE);
 
+		if (WaitForSingleObject(server->stopEvent, 0) == WAIT_OBJECT_0)
+			break;
+
 		if (WAIT_FAILED == status)
 		{
 			WLog_ERR(TAG, "select failed");
@@ -430,33 +438,113 @@ static void pf_server_mainloop(freerdp_listener* listener)
 	}
 
 	listener->Close(listener);
-}
+	ExitThread(0);
-
-int pf_server_start(proxyConfig* config)
-{
-	WSADATA wsaData;
-	freerdp_listener* listener = freerdp_listener_new();
-
-	if (!listener)
-		return -1;
-
-	WTSRegisterWtsApiFunctionTable(FreeRDP_InitWtsApi());
-	winpr_InitializeSSL(WINPR_SSL_INIT_DEFAULT);
-	listener->info = config;
-	listener->PeerAccepted = pf_server_peer_accepted;
-
-	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
-	{
-		freerdp_listener_free(listener);
-		return -1;
-	}
-
-	if (listener->Open(listener, config->Host, config->Port))
-	{
-		pf_server_mainloop(listener);
-	}
-
-	freerdp_listener_free(listener);
-	WSACleanup();
 	return 0;
 }
+
+BOOL pf_server_start(proxyServer* server)
+{
+	WSADATA wsaData;
+	WTSRegisterWtsApiFunctionTable(FreeRDP_InitWtsApi());
+	winpr_InitializeSSL(WINPR_SSL_INIT_DEFAULT);
+
+	if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
+		goto error;
+
+	if (!server->listener->Open(server->listener, server->config->Host, server->config->Port))
+		goto error;
+
+	server->thread = CreateThread(NULL, 0, pf_server_mainloop, (void*)server, 0, NULL);
+	if (!server->thread)
+		goto error;
+
+	return TRUE;
+
+error:
+	WSACleanup();
+	return FALSE;
+}
+
+static void pf_server_clients_list_client_free(void* obj)
+{
+	proxyData* pdata = (proxyData*)obj;
+	proxy_data_abort_connect(pdata);
+}
+
+proxyServer* pf_server_new(proxyConfig* config)
+{
+	proxyServer* server;
+
+	if (!config)
+		return NULL;
+
+	server = calloc(1, sizeof(proxyServer));
+	if (!server)
+		return NULL;
+
+	server->config = config;
+
+	server->stopEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
+	if (!server->stopEvent)
+		goto out;
+
+	server->clients = ArrayList_New(TRUE);
+	if (!server->clients)
+		goto out;
+
+	server->clients->object.fnObjectFree = pf_server_clients_list_client_free;
+
+	server->waitGroup = CountdownEvent_New(0);
+	if (!server->waitGroup)
+		goto out;
+
+	server->listener = freerdp_listener_new();
+	if (!server->listener)
+		goto out;
+
+	server->listener->info = server;
+	server->listener->PeerAccepted = pf_server_peer_accepted;
+	return server;
+
+out:
+	pf_server_free(server);
+	return NULL;
+}
+
+void pf_server_stop(proxyServer* server)
+{
+	HANDLE waitHandle = INVALID_HANDLE_VALUE;
+
+	if (!server)
+		return;
+
+	/* clear clients list, also disconnects every client */
+	ArrayList_Clear(server->clients);
+
+	/* block until all clients are disconnected */
+	waitHandle = CountdownEvent_WaitHandle(server->waitGroup);
+	if (WaitForSingleObject(waitHandle, INFINITE) != WAIT_OBJECT_0)
+		WLog_ERR(TAG, "[%s]: WaitForSingleObject failed!", __FUNCTION__);
+
+	/* signal main thread to stop and wait for the thread to exit */
+	SetEvent(server->stopEvent);
+	WaitForSingleObject(server->thread, INFINITE);
+}
+
+void pf_server_free(proxyServer* server)
+{
+	if (!server)
+		return;
+
+	freerdp_listener_free(server->listener);
+	ArrayList_Free(server->clients);
+	CountdownEvent_Free(server->waitGroup);
+
+	if (server->stopEvent)
+		CloseHandle(server->stopEvent);
+
+	if (server->thread)
+		CloseHandle(server->thread);
+
+	free(server);
+}

server/proxy/pf_server.h

@@ -22,8 +22,26 @@
 #ifndef FREERDP_SERVER_PROXY_SERVER_H
 #define FREERDP_SERVER_PROXY_SERVER_H
 
+#include <winpr/collections.h>
+#include <freerdp/listener.h>
+
 #include "pf_config.h"
 
-int pf_server_start(proxyConfig* config);
+typedef struct proxy_server
+{
+	proxyConfig* config;
+
+	freerdp_listener* listener;
+	wArrayList* clients;        /* maintain a list of active sessions, for stats */
+	wCountdownEvent* waitGroup; /* wait group used for gracefull shutdown */
+	HANDLE thread;              /* main server thread - freerdp listener thread */
+	HANDLE stopEvent;           /* an event used to signal the main thread to stop */
+} proxyServer;
+
+proxyServer* pf_server_new(proxyConfig* config);
+void pf_server_free(proxyServer* server);
+
+BOOL pf_server_start(proxyServer* server);
+void pf_server_stop(proxyServer* server);
 
 #endif /* FREERDP_SERVER_PROXY_SERVER_H */