Implement server-side downgrade from udp

common/network/Udp.cxx

@@ -125,7 +125,7 @@ static uint8_t udpsend(WuClient *client, const uint8_t *data, unsigned len, uint
 	return 0;
 }
 
-UdpStream::UdpStream(): OutStream(), client(NULL), total_len(0), id(0) {
+UdpStream::UdpStream(): OutStream(), client(NULL), total_len(0), id(0), failed(false) {
 	ptr = data;
 	end = data + UDPSTREAM_BUFSIZE;
 
@@ -137,8 +137,10 @@ void UdpStream::flush() {
 	total_len += len;
 
 	if (client) {
-		if (udpsend(client, data, len, &id))
+		if (udpsend(client, data, len, &id)) {
 			vlog.error("Error sending udp, client gone?");
+			failed = true;
+		}
 	} else {
 		vlog.error("Tried to send udp without a client");
 	}
@@ -151,6 +153,14 @@ void UdpStream::overrun(size_t needed) {
 	abort();
 }
 
+bool UdpStream::isFailed() const {
+	return failed;
+}
+
+void UdpStream::clearFailed() {
+	failed = false;
+}
+
 void wuGotHttp(const char msg[], const uint32_t msglen, char resp[]) {
 	WuGotHttp(host, msg, msglen, resp);
 }

common/network/Udp.h

@@ -39,11 +39,15 @@ namespace network {
 			void setClient(WuClient *cli) {
 				client = cli;
 			}
+
+			bool isFailed() const;
+			void clearFailed();
 		private:
 			uint8_t data[UDPSTREAM_BUFSIZE];
 			WuClient *client;
 			size_t total_len;
 			uint32_t id;
+			bool failed;
 	};
 }
 

common/rfb/EncodeManager.cxx

@@ -1716,3 +1716,7 @@ unsigned EncodeManager::scaledQuality(const Rect& rect) const {
 
   return dynamic;
 }
+
+void EncodeManager::resetZlib() {
+  ((TightEncoder *) encoders[encoderTight])->resetZlib();
+}

common/rfb/EncodeManager.h

@@ -79,6 +79,8 @@ namespace rfb {
         return scalingTime;
     };
 
+    void resetZlib();
+
     struct codecstats_t {
       uint32_t ms;
       uint32_t area;

common/rfb/SMsgHandler.h

@@ -96,7 +96,7 @@ namespace rfb {
     virtual void supportsQEMUKeyEvent();
 
     virtual void udpUpgrade(const char *resp) = 0;
-    virtual void udpDowngrade() = 0;
+    virtual void udpDowngrade(const bool) = 0;
 
     ConnParams cp;
   };

common/rfb/SMsgReader.cxx

@@ -346,7 +346,7 @@ void SMsgReader::readUpgradeToUdp()
   }
 
   if (!len) {
-    handler->udpDowngrade();
+    handler->udpDowngrade(false);
     return;
   }
 

common/rfb/TightEncoder.cxx

@@ -57,7 +57,7 @@ static const TightConf conf[10] = {
 };
 
 TightEncoder::TightEncoder(SConnection* conn) :
-  Encoder(conn, encodingTight, EncoderPlain, 256)
+  Encoder(conn, encodingTight, EncoderPlain, 256), zlibNeedsReset(false)
 {
   setCompressLevel(-1);
 }
@@ -166,7 +166,7 @@ void TightEncoder::writeFullColourRect(const PixelBuffer* pb, const Palette& pal
   int stride, h;
 
   os = conn->getOutStream(conn->cp.supportsUdp);
-  if (conn->cp.supportsUdp)
+  if (conn->cp.supportsUdp || zlibNeedsReset)
     os->writeU8((streamId << 4) | (1 << streamId));
   else
     os->writeU8(streamId << 4);
@@ -247,7 +247,7 @@ rdr::OutStream* TightEncoder::getZlibOutStream(int streamId, int level, size_t l
 
   zlibStreams[streamId].setUnderlying(&memStream);
   zlibStreams[streamId].setCompressionLevel(level);
-  if (conn->cp.supportsUdp)
+  if (conn->cp.supportsUdp || zlibNeedsReset)
     zlibStreams[streamId].resetDeflate();
 
   return &zlibStreams[streamId];
@@ -272,6 +272,11 @@ void TightEncoder::flushZlibOutStream(rdr::OutStream* os_)
   memStream.clear();
 }
 
+void TightEncoder::resetZlib()
+{
+  zlibNeedsReset = true;
+}
+
 //
 // Including BPP-dependent implementation of the encoder.
 //

common/rfb/TightEncoder.h

@@ -39,6 +39,7 @@ namespace rfb {
     virtual void writeSolidRect(int width, int height,
                                 const PixelFormat& pf,
                                 const rdr::U8* colour);
+    void resetZlib();
 
   protected:
     void writeMonoRect(const PixelBuffer* pb, const Palette& palette);
@@ -76,6 +77,7 @@ namespace rfb {
     rdr::MemOutStream memStream;
 
     int idxZlibLevel, monoZlibLevel, rawZlibLevel;
+    bool zlibNeedsReset;
   };
 
 }

common/rfb/TightEncoderBPP.cxx

@@ -40,7 +40,7 @@ void TightEncoder::writeMonoRect(int width, int height,
 
   os = conn->getOutStream(conn->cp.supportsUdp);
 
-  if (conn->cp.supportsUdp)
+  if (conn->cp.supportsUdp || zlibNeedsReset)
     os->writeU8(((streamId | tightExplicitFilter) << 4) | (1 << streamId));
   else
     os->writeU8((streamId | tightExplicitFilter) << 4);
@@ -130,7 +130,7 @@ void TightEncoder::writeIndexedRect(int width, int height,
 
   os = conn->getOutStream(conn->cp.supportsUdp);
 
-  if (conn->cp.supportsUdp)
+  if (conn->cp.supportsUdp || zlibNeedsReset)
     os->writeU8(((streamId | tightExplicitFilter) << 4) | (1 << streamId));
   else
     os->writeU8((streamId | tightExplicitFilter) << 4);

common/rfb/VNCSConnectionST.cxx

@@ -1769,10 +1769,12 @@ void VNCSConnectionST::udpUpgrade(const char *resp)
   writer()->writeUdpUpgrade(resp);
 }
 
-void VNCSConnectionST::udpDowngrade()
+void VNCSConnectionST::udpDowngrade(const bool byServer)
 {
   cp.supportsUdp = false;
   cp.useCopyRect = true;
+  encodeManager.resetZlib();
 
-  vlog.info("Client %s downgrading from udp", sock->getPeerAddress());
+  vlog.info("Client %s downgrading from udp by %s", sock->getPeerAddress(),
+            byServer ? "the server" : "its own request");
 }

common/rfb/VNCSConnectionST.h

@@ -196,6 +196,8 @@ namespace rfb {
       return encodeManager.getScalingTime();
     }
 
+    virtual void udpDowngrade(const bool byServer);
+
     bool upgradingToUdp;
 
   private:
@@ -220,7 +222,6 @@ namespace rfb {
     virtual void handleClipboardAnnounce(bool available);
     virtual void handleClipboardAnnounceBinary(const unsigned num, const char mimes[][32]);
     virtual void udpUpgrade(const char *resp);
-    virtual void udpDowngrade();
     virtual void supportsLocalCursor();
     virtual void supportsFence();
     virtual void supportsContinuousUpdates();

common/rfb/VNCServerST.cxx

@@ -1064,6 +1064,11 @@ void VNCServerST::writeUpdate()
     (*ci)->add_changed(ui.changed);
     (*ci)->writeFramebufferUpdateOrClose();
 
+    if (((network::UdpStream *)(*ci)->getOutStream(true))->isFailed()) {
+      ((network::UdpStream *)(*ci)->getOutStream(true))->clearFailed();
+      (*ci)->udpDowngrade(true);
+    }
+
     if (apimessager) {
       (*ci)->sendStats(false);
       const EncodeManager::codecstats_t subjpeg = (*ci)->getJpegStats();