Bug fixes for short WiFi network transfers

Fixed bug that could lead to empty files on upload
Fixed another related bug causing small files to be truncated

2 files changed, 15 insertions, 7 deletions

diff --git a/src/Networking/ESP8266WiFi/WiFiSocket.cpp b/src/Networking/ESP8266WiFi/WiFiSocket.cpp
index 03eb751b..c6c46eee 100644
--- a/src/Networking/ESP8266WiFi/WiFiSocket.cpp
+++ b/src/Networking/ESP8266WiFi/WiFiSocket.cpp
@@ -15,7 +15,7 @@
 
 const unsigned int MaxBuffersPerSocket = 4;
 
-WiFiSocket::WiFiSocket(NetworkInterface *iface) noexcept : Socket(iface), receivedData(nullptr), state(SocketState::inactive), needsPolling(false)
+WiFiSocket::WiFiSocket(NetworkInterface *iface) noexcept : Socket(iface), receivedData(nullptr), hasMoreDataPending(false), state(SocketState::inactive), needsPolling(false)
 {
 }
 
@@ -35,6 +35,7 @@ void WiFiSocket::Close() noexcept
 		if (reply == ResponseEmpty)
 		{
 			state = (state == SocketState::connected) ? SocketState::closing : SocketState::inactive;
+			DiscardReceivedData();
 			return;
 		}
 	}
@@ -63,7 +64,7 @@ void WiFiSocket::Terminate() noexcept
 bool WiFiSocket::CanRead() const noexcept
 {
 	return (state == SocketState::connected)
-		|| (state == SocketState::clientDisconnecting && receivedData != nullptr && receivedData->TotalRemaining() != 0);
+		|| (state == SocketState::clientDisconnecting && (hasMoreDataPending || (receivedData != nullptr && receivedData->TotalRemaining() != 0)));
 }
 
 // Return true if we can send data to this socket
@@ -144,8 +145,11 @@ void WiFiSocket::Poll() noexcept
 
 		if (state == SocketState::clientDisconnecting)
 		{
-			// We already got here before, so close the connection once and for all
-			Close();
+			if (!CanRead())
+			{
+				// We already got here before, so close the connection once and for all
+				Close();
+			}
 			break;
 		}
 		else if (state != SocketState::inactive)
@@ -173,7 +177,6 @@ void WiFiSocket::Poll() noexcept
 			localPort = resp.Value().localPort;
 			remotePort = resp.Value().remotePort;
 			remoteIPAddress.SetV4LittleEndian(resp.Value().remoteIp);
-			DiscardReceivedData();
 			if (state != SocketState::waitingForResponder)
 			{
 				WiFiInterface *iface = static_cast<WiFiInterface *>(interface);
@@ -208,7 +211,7 @@ void WiFiSocket::Poll() noexcept
 		break;
 
 	case ConnState::aborted:
-		if (reprap.Debug( moduleNetwork))
+		if (reprap.Debug(moduleNetwork))
 		{
 			debugPrintf("Socket %u aborted\n", socketNum);
 		}
@@ -220,7 +223,7 @@ void WiFiSocket::Poll() noexcept
 		if (state == SocketState::connected || state == SocketState::waitingForResponder)
 		{
 			// Unexpected change of state
-			if (state != SocketState::clientDisconnecting && reprap.Debug(moduleNetwork))
+			if (reprap.Debug(moduleNetwork))
 			{
 				debugPrintf("Unexpected state change on socket %u\n", socketNum);
 			}
@@ -257,6 +260,7 @@ void WiFiSocket::ReceiveData(uint16_t bytesAvailable) noexcept
 			const int32_t ret = GetInterface()->SendCommand(NetworkCommand::connRead, socketNum, 0, 0, nullptr, 0, lastBuffer->UnwrittenData(), maxToRead);
 			if (ret > 0 && (size_t)ret <= maxToRead)
 			{
+				bytesAvailable -= ret;
 				lastBuffer->dataLength += (size_t)ret;
 				if (reprap.Debug(moduleNetwork))
 				{
@@ -273,6 +277,7 @@ void WiFiSocket::ReceiveData(uint16_t bytesAvailable) noexcept
 				const int32_t ret = GetInterface()->SendCommand(NetworkCommand::connRead, socketNum, 0, 0, nullptr, 0, buf->Data(), maxToRead);
 				if (ret > 0 && (size_t)ret <= maxToRead)
 				{
+					bytesAvailable -= ret;
 					buf->dataLength = (size_t)ret;
 					NetworkBuffer::AppendToList(&receivedData, buf);
 					if (reprap.Debug(moduleNetwork))
@@ -288,6 +293,7 @@ void WiFiSocket::ReceiveData(uint16_t bytesAvailable) noexcept
 //			else debugPrintf("no buffer\n");
 		}
 	}
+	hasMoreDataPending = (bytesAvailable != 0);
 }
 
 // Discard any received data for this transaction
@@ -297,6 +303,7 @@ void WiFiSocket::DiscardReceivedData() noexcept
 	{
 		receivedData = receivedData->Release();
 	}
+	hasMoreDataPending = false;
 }
 
 // Send the data, returning the length buffered
diff --git a/src/Networking/ESP8266WiFi/WiFiSocket.h b/src/Networking/ESP8266WiFi/WiFiSocket.h
index 3951f2a0..accbc0b1 100644
--- a/src/Networking/ESP8266WiFi/WiFiSocket.h
+++ b/src/Networking/ESP8266WiFi/WiFiSocket.h
@@ -55,6 +55,7 @@ private:
 	void DiscardReceivedData() noexcept;
 
 	NetworkBuffer *receivedData;						// List of buffers holding received data
+	bool hasMoreDataPending;							// If there is more data left to read when the buffered data has been processed
 	uint32_t whenConnected;
 	uint16_t txBufferSpace;								// How much free transmit buffer space the WiFi mofule reported
 	SocketNumber socketNum;								// The WiFi socket number we are using