Make sure TLS is allowed to terminate gracefully

Some systems (like TLS) need to send some final data before closing
a connection. Make sure this is properly handled by cleaning up the
security object before closing the underlying network socket.

5 files changed, 39 insertions, 17 deletions

diff --git a/common/rfb/CConnection.cxx b/common/rfb/CConnection.cxx
index bdde3253..f541a969 100644
--- a/common/rfb/CConnection.cxx
+++ b/common/rfb/CConnection.cxx
@@ -60,14 +60,7 @@ CConnection::CConnection()
 
 CConnection::~CConnection()
 {
-  setFramebuffer(NULL);
-  if (csecurity)
-    delete csecurity;
-  delete reader_;
-  reader_ = 0;
-  delete writer_;
-  writer_ = 0;
-  strFree(serverClipboard);
+  close();
 }
 
 void CConnection::setStreams(rdr::InStream* is_, rdr::OutStream* os_)
@@ -336,6 +329,21 @@ void CConnection::securityCompleted()
   writer_->writeClientInit(shared);
 }
 
+void CConnection::close()
+{
+  state_ = RFBSTATE_CLOSING;
+
+  setFramebuffer(NULL);
+  delete csecurity;
+  csecurity = NULL;
+  delete reader_;
+  reader_ = NULL;
+  delete writer_;
+  writer_ = NULL;
+  strFree(serverClipboard);
+  serverClipboard = NULL;
+}
+
 void CConnection::setDesktopSize(int w, int h)
 {
   decoder.flush();
diff --git a/common/rfb/CConnection.h b/common/rfb/CConnection.h
index f01d5d36..3857be4d 100644
--- a/common/rfb/CConnection.h
+++ b/common/rfb/CConnection.h
@@ -86,6 +86,11 @@ namespace rfb {
     // NB: In either case, you must have called initialiseProtocol() first.
     void processMsg();
 
+    // close() gracefully shuts down the connection to the server and
+    // should be called before terminating the underlying network
+    // connection
+    void close();
+
 
     // Methods overridden from CMsgHandler
 
@@ -213,6 +218,7 @@ namespace rfb {
       RFBSTATE_SECURITY_RESULT,
       RFBSTATE_INITIALISATION,
       RFBSTATE_NORMAL,
+      RFBSTATE_CLOSING,
       RFBSTATE_INVALID
     };
 
diff --git a/common/rfb/SConnection.cxx b/common/rfb/SConnection.cxx
index 9e6d0322..b5a69d4c 100644
--- a/common/rfb/SConnection.cxx
+++ b/common/rfb/SConnection.cxx
@@ -66,13 +66,7 @@ SConnection::SConnection()
 
 SConnection::~SConnection()
 {
-  if (ssecurity)
-    delete ssecurity;
-  delete reader_;
-  reader_ = 0;
-  delete writer_;
-  writer_ = 0;
-  strFree(clientClipboard);
+  cleanup();
 }
 
 void SConnection::setStreams(rdr::InStream* is_, rdr::OutStream* os_)
@@ -464,6 +458,7 @@ void SConnection::clientInit(bool shared)
 void SConnection::close(const char* reason)
 {
   state_ = RFBSTATE_CLOSING;
+  cleanup();
 }
 
 void SConnection::setPixelFormat(const PixelFormat& pf)
@@ -552,6 +547,18 @@ void SConnection::sendClipboardData(const char* data)
   }
 }
 
+void SConnection::cleanup()
+{
+  delete ssecurity;
+  ssecurity = NULL;
+  delete reader_;
+  reader_ = NULL;
+  delete writer_;
+  writer_ = NULL;
+  strFree(clientClipboard);
+  clientClipboard = NULL;
+}
+
 void SConnection::writeFakeColourMap(void)
 {
   int i;
diff --git a/common/rfb/SConnection.h b/common/rfb/SConnection.h
index a7c4e0a6..e7bbf2c3 100644
--- a/common/rfb/SConnection.h
+++ b/common/rfb/SConnection.h
@@ -230,6 +230,7 @@ namespace rfb {
     void setWriter(SMsgWriter *w) { writer_ = w; }
 
   private:
+    void cleanup();
     void writeFakeColourMap(void);
 
     bool readyForSetColourMapEntries;
diff --git a/common/rfb/VNCSConnectionST.cxx b/common/rfb/VNCSConnectionST.cxx
index 6ac9edbe..00f640b3 100644
--- a/common/rfb/VNCSConnectionST.cxx
+++ b/common/rfb/VNCSConnectionST.cxx
@@ -108,6 +108,8 @@ bool VNCSConnectionST::accessCheck(AccessRights ar) const
 
 void VNCSConnectionST::close(const char* reason)
 {
+  SConnection::close(reason);
+
   // Log the reason for the close
   if (!closeReason.buf)
     closeReason.buf = strDup(reason);
@@ -129,8 +131,6 @@ void VNCSConnectionST::close(const char* reason)
   // calling code will call VNCServerST's removeSocket() method causing us to
   // be deleted.
   sock->shutdown();
-
-  SConnection::close(reason);
 }