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.

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);
 }

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)

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;

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

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;

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)

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;

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

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;

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.