Always flush sockets on shutdown()

The system shutdown() function doesn't drop buffered data, so neither
should we.

We had one fix in place, but that didn't cover all cases. Move this
handling to all socket like classes we have.

10 files changed, 87 insertions, 19 deletions

diff --git a/common/network/Socket.cxx b/common/network/Socket.cxx
index 0a35e267..49abbc84 100644
--- a/common/network/Socket.cxx
+++ b/common/network/Socket.cxx
@@ -43,8 +43,12 @@
 
 #include <network/Socket.h>
 
+#include <rfb/LogWriter.h>
+
 using namespace network;
 
+static rfb::LogWriter vlog("Socket");
+
 // -=- Socket initialisation
 static bool socketsInitialised = false;
 void network::initSockets() {
@@ -98,6 +102,17 @@ Socket::~Socket()
 // if shutdown() is overridden then the override MUST call on to here
 void Socket::shutdown()
 {
+  try {
+    if (outstream->hasBufferedData()) {
+      outstream->cork(false);
+      outstream->flush();
+      if (outstream->hasBufferedData())
+        vlog.error("Failed to flush remaining socket data on close");
+    }
+  } catch (std::exception& e) {
+    vlog.error("Failed to flush remaining socket data on close: %s", e.what());
+  }
+
   isShutdown_ = true;
   ::shutdown(getFd(), SHUT_RDWR);
 }
diff --git a/common/rfb/CSecurityRSAAES.cxx b/common/rfb/CSecurityRSAAES.cxx
index 56e1d702..0985d0f2 100644
--- a/common/rfb/CSecurityRSAAES.cxx
+++ b/common/rfb/CSecurityRSAAES.cxx
@@ -55,6 +55,8 @@ const int MaxKeyLength = 8192;
 
 using namespace rfb;
 
+static LogWriter vlog("CSecurityRSAAES");
+
 CSecurityRSAAES::CSecurityRSAAES(CConnection* cc_, uint32_t _secType,
                                  int _keySize, bool _isAllEncrypted)
   : CSecurity(cc_), state(ReadPublicKey),
@@ -74,6 +76,19 @@ CSecurityRSAAES::~CSecurityRSAAES()
 
 void CSecurityRSAAES::cleanup()
 {
+  if (raos) {
+    try {
+      if (raos->hasBufferedData()) {
+        raos->cork(false);
+        raos->flush();
+        if (raos->hasBufferedData())
+          vlog.error("Failed to flush remaining socket data on close");
+      }
+    } catch (std::exception& e) {
+      vlog.error("Failed to flush remaining socket data on close: %s", e.what());
+    }
+  }
+
   if (serverKeyN)
     delete[] serverKeyN;
   if (serverKeyE)
diff --git a/common/rfb/CSecurityRSAAES.h b/common/rfb/CSecurityRSAAES.h
index e5548c1d..af380bd3 100644
--- a/common/rfb/CSecurityRSAAES.h
+++ b/common/rfb/CSecurityRSAAES.h
@@ -34,6 +34,8 @@
 namespace rdr {
   class InStream;
   class OutStream;
+  class AESInStream;
+  class AESOutStream;
 }
 
 namespace rfb {
@@ -79,8 +81,8 @@ namespace rfb {
     uint8_t serverRandom[32];
     uint8_t clientRandom[32];
 
-    rdr::InStream* rais;
-    rdr::OutStream* raos;
+    rdr::AESInStream* rais;
+    rdr::AESOutStream* raos;
 
     rdr::InStream* rawis;
     rdr::OutStream* rawos;
diff --git a/common/rfb/CSecurityTLS.cxx b/common/rfb/CSecurityTLS.cxx
index bb7e3879..0c10a85d 100644
--- a/common/rfb/CSecurityTLS.cxx
+++ b/common/rfb/CSecurityTLS.cxx
@@ -85,6 +85,19 @@ CSecurityTLS::CSecurityTLS(CConnection* cc_, bool _anon)
 
 void CSecurityTLS::shutdown()
 {
+  if (tlsos) {
+    try {
+      if (tlsos->hasBufferedData()) {
+        tlsos->cork(false);
+        tlsos->flush();
+        if (tlsos->hasBufferedData())
+          vlog.error("Failed to flush remaining socket data on close");
+      }
+    } catch (std::exception& e) {
+      vlog.error("Failed to flush remaining socket data on close: %s", e.what());
+    }
+  }
+
   if (session) {
     int ret;
     // FIXME: We can't currently wait for the response, so we only send
diff --git a/common/rfb/CSecurityTLS.h b/common/rfb/CSecurityTLS.h
index 5091682b..2464cb6c 100644
--- a/common/rfb/CSecurityTLS.h
+++ b/common/rfb/CSecurityTLS.h
@@ -34,6 +34,8 @@
 namespace rdr {
   class InStream;
   class OutStream;
+  class TLSInStream;
+  class TLSOutStream;
 }
 
 namespace rfb {
@@ -61,8 +63,8 @@ namespace rfb {
     gnutls_certificate_credentials_t cert_cred;
     bool anon;
 
-    rdr::InStream* tlsis;
-    rdr::OutStream* tlsos;
+    rdr::TLSInStream* tlsis;
+    rdr::TLSOutStream* tlsos;
 
     rdr::InStream* rawis;
     rdr::OutStream* rawos;
diff --git a/common/rfb/SSecurityRSAAES.cxx b/common/rfb/SSecurityRSAAES.cxx
index 39b286de..fed213ad 100644
--- a/common/rfb/SSecurityRSAAES.cxx
+++ b/common/rfb/SSecurityRSAAES.cxx
@@ -74,6 +74,8 @@ BoolParameter SSecurityRSAAES::requireUsername
 ("RequireUsername", "Require username for the RSA-AES security types",
  false, ConfServer);
 
+static LogWriter vlog("CSecurityRSAAES");
+
 SSecurityRSAAES::SSecurityRSAAES(SConnection* sc_, uint32_t _secType,
                                  int _keySize, bool _isAllEncrypted)
   : SSecurity(sc_), state(SendPublicKey),
@@ -94,6 +96,19 @@ SSecurityRSAAES::~SSecurityRSAAES()
 
 void SSecurityRSAAES::cleanup()
 {
+  if (raos) {
+    try {
+      if (raos->hasBufferedData()) {
+        raos->cork(false);
+        raos->flush();
+        if (raos->hasBufferedData())
+          vlog.error("Failed to flush remaining socket data on close");
+      }
+    } catch (std::exception& e) {
+      vlog.error("Failed to flush remaining socket data on close: %s", e.what());
+    }
+  }
+
   if (serverKeyN)
     delete[] serverKeyN;
   if (serverKeyE)
diff --git a/common/rfb/SSecurityRSAAES.h b/common/rfb/SSecurityRSAAES.h
index a6fbc499..e3300cb7 100644
--- a/common/rfb/SSecurityRSAAES.h
+++ b/common/rfb/SSecurityRSAAES.h
@@ -32,6 +32,8 @@
 namespace rdr {
   class InStream;
   class OutStream;
+  class AESInStream;
+  class AESOutStream;
 }
 
 namespace rfb {
@@ -89,8 +91,8 @@ namespace rfb {
     char password[256];
     AccessRights accessRights;
 
-    rdr::InStream* rais;
-    rdr::OutStream* raos;
+    rdr::AESInStream* rais;
+    rdr::AESOutStream* raos;
 
     rdr::InStream* rawis;
     rdr::OutStream* rawos;
diff --git a/common/rfb/SSecurityTLS.cxx b/common/rfb/SSecurityTLS.cxx
index 4b036e27..b297242b 100644
--- a/common/rfb/SSecurityTLS.cxx
+++ b/common/rfb/SSecurityTLS.cxx
@@ -85,6 +85,19 @@ SSecurityTLS::SSecurityTLS(SConnection* sc_, bool _anon)
 
 void SSecurityTLS::shutdown()
 {
+  if (tlsos) {
+    try {
+      if (tlsos->hasBufferedData()) {
+        tlsos->cork(false);
+        tlsos->flush();
+        if (tlsos->hasBufferedData())
+          vlog.error("Failed to flush remaining socket data on close");
+      }
+    } catch (std::exception& e) {
+      vlog.error("Failed to flush remaining socket data on close: %s", e.what());
+    }
+  }
+
   if (session) {
     int ret;
     // FIXME: We can't currently wait for the response, so we only send
diff --git a/common/rfb/SSecurityTLS.h b/common/rfb/SSecurityTLS.h
index 8f6539d9..1dc33cfd 100644
--- a/common/rfb/SSecurityTLS.h
+++ b/common/rfb/SSecurityTLS.h
@@ -40,6 +40,8 @@
 namespace rdr {
   class InStream;
   class OutStream;
+  class TLSInStream;
+  class TLSOutStream;
 }
 
 namespace rfb {
@@ -69,8 +71,8 @@ namespace rfb {
 
     bool anon;
 
-    rdr::InStream* tlsis;
-    rdr::OutStream* tlsos;
+    rdr::TLSInStream* tlsis;
+    rdr::TLSOutStream* tlsos;
 
     rdr::InStream* rawis;
     rdr::OutStream* rawos;
diff --git a/common/rfb/VNCSConnectionST.cxx b/common/rfb/VNCSConnectionST.cxx
index cb36872a..a354f636 100644
--- a/common/rfb/VNCSConnectionST.cxx
+++ b/common/rfb/VNCSConnectionST.cxx
@@ -124,17 +124,6 @@ void VNCSConnectionST::close(const char* reason)
   else
     vlog.debug("Second close: %s (%s)", peerEndpoint.c_str(), reason);
 
-  try {
-    if (sock->outStream().hasBufferedData()) {
-      sock->outStream().cork(false);
-      sock->outStream().flush();
-      if (sock->outStream().hasBufferedData())
-        vlog.error("Failed to flush remaining socket data on close");
-    }
-  } catch (std::exception& e) {
-    vlog.error("Failed to flush remaining socket data on close: %s", e.what());
-  }
-
   // Just shutdown the socket and mark our state as closing.  Eventually the
   // calling code will call VNCServerST's removeSocket() method causing us to
   // be deleted.