From 99b8fc7e0fe2982985abf0856ede8b93e8ad32c3 Mon Sep 17 00:00:00 2001
From: Sri Ramanujam <sramanujam@datto.com>
Date: Fri, 1 Dec 2017 14:14:24 -0500
Subject: [PATCH] Properly cleanup shadow server when listener is stopped

---
 include/rdp/RDPListener.h |  9 +++++++--
 src/rdp/RDPListener.cpp   | 34 +++++++++++++++++++++++++++++-----
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/include/rdp/RDPListener.h b/include/rdp/RDPListener.h
index a22ae95..0384257 100644
--- a/include/rdp/RDPListener.h
+++ b/include/rdp/RDPListener.h
@@ -152,6 +152,11 @@ class RDPListener
      */
     std::string CredentialPath();
 
+    /**
+     * @brief Shutdown and unregister the listener.
+     */
+    void shutdown();
+
     /**
      * @brief The freerdp_listener struct this object manages.
      *
@@ -247,11 +252,11 @@ class RDPListener
     /**
      * @brief Mutex guarding stop.
      */
-    std::mutex stopMutex;
+    std::mutex listenerStopMutex;
     /**
      * @brief Flag to be set when the listener needs to clean up and exit.
      */
-    bool stop;
+    bool listener_running;
     /**
      * @brief DBus id.
      */
diff --git a/src/rdp/RDPListener.cpp b/src/rdp/RDPListener.cpp
index 2e76823..bc8fb30 100644
--- a/src/rdp/RDPListener.cpp
+++ b/src/rdp/RDPListener.cpp
@@ -48,11 +48,11 @@ RDPListener::RDPListener(std::string uuid, int vm_id, uint16_t port, RDPServerWo
                                                                      uuid(uuid),
                                                                      samfile(),
                                                                      vm_id(vm_id),
+                                                                     listener_running(false),
                                                                      targetFPS(30),
                                                                      credential_path()
 {
     WTSRegisterWtsApiFunctionTable(FreeRDP_InitWtsApi());
-    stop = false;
 
     shadow_subsystem_set_entry(RDPMux_ShadowSubsystemEntry);
     server = shadow_server_new();
@@ -68,9 +68,15 @@ RDPListener::RDPListener(std::string uuid, int vm_id, uint16_t port, RDPServerWo
 
 RDPListener::~RDPListener()
 {
-    shadow_server_stop(server);
-    dbus_conn->unregister_object(registered_id);
+    {
+        std::lock_guard<std::mutex> lock(listenerStopMutex);
+        if (listener_running) {
+            shadow_server_stop(server);
+        }
+        listener_running = false;
+    }
     shadow_server_free(server);
+    dbus_conn->unregister_object(registered_id);
     WSACleanup();
 }
 
@@ -121,6 +127,11 @@ void RDPListener::RunServer()
         goto cleanup;
     }
 
+    {
+        std::lock_guard<std::mutex> lock(listenerStopMutex);
+        listener_running = true;
+    }
+
     WaitForSingleObject(this->server->thread, INFINITE);
 
     if (!GetExitCodeThread(this->server->thread, &exitCode)) {
@@ -131,9 +142,22 @@ void RDPListener::RunServer()
 
     VLOG(1) << "LISTENER " << this << ": Main loop exited, exit code " << status;
 cleanup:
-    parent->UnregisterVM(this->uuid, this->port); // this will trigger destruction of the RDPListener object.
+    shutdown(); // this will trigger destruction of the RDPListener object.
 }
 
+void RDPListener::shutdown()
+{
+    {
+        std::lock_guard<std::mutex> lock(listenerStopMutex);
+        if (listener_running) {
+            shadow_server_stop(this->server);
+        }
+        listener_running = false;
+    }
+    parent->UnregisterVM(this->uuid, this->port);
+}
+
+
 void RDPListener::processOutgoingMessage(std::vector<uint16_t> vec)
 {
     QueueItem item = std::make_tuple(vec, this->uuid);
@@ -150,7 +174,7 @@ void RDPListener::processIncomingMessage(std::vector<uint32_t> rvec)
         processDisplaySwitch(rvec);
     } else if (rvec[0] == SHUTDOWN) {
         VLOG(2) << "LISTENER " << this << ": Shutdown event received!";
-        parent->UnregisterVM(this->uuid, this->port);
+        shutdown();
     } else {
         // what the hell have you sent me
         LOG(WARNING) << "Invalid message type sent.";