From cab19b21a45c0ce3f6fe5385087623eb741bde78 Mon Sep 17 00:00:00 2001
From: Guenter Obiltschnig <guenter.obiltschnig@appinf.com>
Date: Tue, 25 Jun 2013 13:56:28 +0200
Subject: [PATCH] improved WebSocket closing behavior

---
 Net/include/Poco/Net/WebSocket.h |  6 +++---
 Net/src/WebSocketImpl.cpp        | 29 ++++++++++++++++++++---------
 2 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/Net/include/Poco/Net/WebSocket.h b/Net/include/Poco/Net/WebSocket.h
index d738d45b1..0e19097a2 100644
--- a/Net/include/Poco/Net/WebSocket.h
+++ b/Net/include/Poco/Net/WebSocket.h
@@ -1,7 +1,7 @@
 //
 // WebSocket.h
 //
-// $Id: //poco/1.4/Net/include/Poco/Net/WebSocket.h#2 $
+// $Id: //poco/1.4/Net/include/Poco/Net/WebSocket.h#4 $
 //
 // Library: Net
 // Package: WebSocket
@@ -230,8 +230,8 @@ public:
 		/// terminated.
 		///
 		/// Returns the number of bytes received. 
-		/// A return value of 0 means a graceful shutdown 
-		/// of the connection from the peer.
+		/// A return value of 0 means that the peer has
+		/// shut down or closed the connection.
 		///
 		/// Throws a TimeoutException if a receive timeout has
 		/// been set and nothing is received within that interval.
diff --git a/Net/src/WebSocketImpl.cpp b/Net/src/WebSocketImpl.cpp
index e24128465..6864ec2fc 100644
--- a/Net/src/WebSocketImpl.cpp
+++ b/Net/src/WebSocketImpl.cpp
@@ -1,7 +1,7 @@
 //
 // WebSocketImpl.cpp
 //
-// $Id: //poco/1.4/Net/src/WebSocketImpl.cpp#8 $
+// $Id: //poco/1.4/Net/src/WebSocketImpl.cpp#9 $
 //
 // Library: Net
 // Package: WebSocket
@@ -120,6 +120,11 @@ int WebSocketImpl::receiveBytes(void* buffer, int length, int)
 {
 	char header[MAX_HEADER_LENGTH];
 	int n = receiveNBytes(header, 2);
+	if (n <= 0)
+	{
+		_frameFlags = 0;
+		return n;
+	}
 	poco_assert (n == 2);
 	Poco::UInt8 lengthByte = static_cast<Poco::UInt8>(header[1]);
 	int maskOffset = 0;
@@ -134,7 +139,8 @@ int WebSocketImpl::receiveBytes(void* buffer, int length, int)
 		n = receiveNBytes(header + 2, MAX_HEADER_LENGTH - 2);
 	}
 
-	poco_assert (n > 0);
+	if (n <= 0) throw WebSocketException("Incomplete frame received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
+
 	n += 2;
 	Poco::MemoryInputStream istr(header, n);
 	Poco::BinaryReader reader(istr, Poco::BinaryReader::NETWORK_BYTE_ORDER);
@@ -179,7 +185,9 @@ int WebSocketImpl::receiveBytes(void* buffer, int length, int)
 	}
 	if (received < payloadLength)
 	{
-		received += receiveNBytes(reinterpret_cast<char*>(buffer) + received, payloadLength - received);
+		n = receiveNBytes(reinterpret_cast<char*>(buffer) + received, payloadLength - received);
+		if (n <= 0) throw WebSocketException("Incomplete frame received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
+		received += n;
 	}
 	if (lengthByte & FRAME_FLAG_MASK)
 	{
@@ -196,13 +204,16 @@ int WebSocketImpl::receiveBytes(void* buffer, int length, int)
 int WebSocketImpl::receiveNBytes(void* buffer, int bytes)
 {
 	int received = _pStreamSocketImpl->receiveBytes(reinterpret_cast<char*>(buffer), bytes);
-	while (received < bytes)
+	if (received > 0)
 	{
-		int n = _pStreamSocketImpl->receiveBytes(reinterpret_cast<char*>(buffer) + received, bytes - received);
-		if (n > 0)
-			received += n;
-		else
-			throw WebSocketException("Incomplete frame received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
+		while (received < bytes)
+		{
+			int n = _pStreamSocketImpl->receiveBytes(reinterpret_cast<char*>(buffer) + received, bytes - received);
+			if (n > 0)
+				received += n;
+			else
+				throw WebSocketException("Incomplete frame received", WebSocket::WS_ERR_INCOMPLETE_FRAME);
+		}
 	}
 	return received;
 }