server: proxy: disconnect all active sessions on shutdown
authorKobi Mizrachi <kmizrachi18@gmail.com>
Sun, 19 Jan 2020 09:51:34 +0000 (11:51 +0200)
committerakallabeth <akallabeth@users.noreply.github.com>
Tue, 21 Jan 2020 11:08:39 +0000 (12:08 +0100)
server/proxy/freerdp_proxy.c
server/proxy/pf_config.c
server/proxy/pf_config.h
server/proxy/pf_context.c
server/proxy/pf_server.c
server/proxy/pf_server.h

index fad58e7..db5a4a7 100644 (file)
 
 #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_internal(&config);
+       pf_server_config_free(server->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";
-       int status = 0;
+       proxyConfig* config = NULL;
+       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);
+
+       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;
 
-       pf_server_config_print(&config);
-       status = pf_server_start(&config);
+       status = 0;
 fail:
+       pf_server_free(server);
        pf_modules_free();
-       pf_server_config_free_internal(&config);
+       pf_server_config_free(config);
        return status;
 }
index 8aa77de..227b65b 100644 (file)
@@ -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);
 }
index 8270936..f33e88a 100644 (file)
@@ -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 */
index ac002b3..c143129 100644 (file)
@@ -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)
index 696eac9..0aadb0e 100644 (file)
  * 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);
+       return 0;
 }
 
-int pf_server_start(proxyConfig* config)
+BOOL pf_server_start(proxyServer* server)
 {
        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;
-       }
+               goto error;
 
-       if (listener->Open(listener, config->Host, config->Port))
-       {
-               pf_server_mainloop(listener);
-       }
+       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;
 
-       freerdp_listener_free(listener);
+       return TRUE;
+
+error:
        WSACleanup();
-       return 0;
+       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);
 }
index 1da72dd..6d920e0 100644 (file)
 #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 */