[SERVER] Fix premature connection close for RPC replies

3 files changed, 40 insertions, 7 deletions

diff --git a/src/server/rpc.c b/src/server/rpc.c
index af3507c..d02ff92 100644
--- a/src/server/rpc.c
+++ b/src/server/rpc.c
@@ -4,6 +4,7 @@
 #include "uplink.h"
 #include "log.h"
 #include "locks.h"
+#include "sockhelper.h"
 #include "helper.h"
 #include "image.h"
 
@@ -16,13 +17,12 @@ static void clientsToJson(json_t *jsonClients);
 
 void rpc_sendStatsJson(int sock)
 {
-	uint64_t receivedBytes = uplink_getTotalBytesReceived();
-	uint64_t sentBytes = net_getTotalBytesSent();
-	int uptime = dnbd3_serverUptime();
-
+	const uint64_t receivedBytes = uplink_getTotalBytesReceived();
+	const uint64_t sentBytes = net_getTotalBytesSent();
+	const int uptime = dnbd3_serverUptime();
 	json_t *jsonClients = json_array();
-	clientsToJson( jsonClients );
 
+	clientsToJson( jsonClients );
 
 	json_t *statisticsJson = json_pack( "{sIsI}", "receivedBytes", (json_int_t) receivedBytes, "sentBytes", (json_int_t) sentBytes );
 	json_object_set_new( statisticsJson, "clients", jsonClients );
@@ -31,11 +31,14 @@ void rpc_sendStatsJson(int sock)
 	char *jsonString = json_dumps( statisticsJson, 0 );
 
 	char buffer[500];
-	snprintf(buffer, sizeof(buffer), "HTTP/1.1 200 OK\r\nConnection: Close\r\nContent-Length: %d\r\nContent-Type: application/json\r\n\r\n",
+	snprintf(buffer, sizeof buffer , "HTTP/1.1 200 OK\r\nConnection: Close\r\nContent-Length: %d\r\nContent-Type: application/json\r\n\r\n",
 			(int) strlen( jsonString ) );
 	write( sock, buffer, strlen( buffer ) );
-	write( sock, jsonString, strlen( jsonString ) );
+	sock_sendAll( sock, jsonString, strlen( jsonString ), 10 );
 	json_decref( statisticsJson );
+	// Wait for flush
+	shutdown( sock, SHUT_WR );
+	while ( read( sock, buffer, sizeof buffer ) > 0 );
 	free( jsonString );
 }
 
diff --git a/src/server/sockhelper.c b/src/server/sockhelper.c
index e45efba..fb09ec2 100644
--- a/src/server/sockhelper.c
+++ b/src/server/sockhelper.c
@@ -189,3 +189,26 @@ bool sock_append(poll_list_t *list, const int sock, bool wantRead, bool wantWrit
 	list->count++;
 	return true;
 }
+
+ssize_t sock_sendAll(int sock, void *buffer, size_t len, int maxtries)
+{
+	size_t done = 0;
+	ssize_t ret = 0;
+	while ( done < len ) {
+		if ( maxtries >= 0 && --maxtries == -1 ) break;
+		ret = write( sock, (char*)buffer + done, len - done );
+		if ( ret < 0 ) {
+			if ( errno == EINTR ) continue;
+			if ( errno == EAGAIN || errno == EWOULDBLOCK ) {
+				usleep( 1000 );
+				continue;
+			}
+			break;
+		}
+		if ( ret == 0 ) break;
+		done += ret;
+	}
+	if ( done == 0 ) return ret;
+	return done;
+}
+
diff --git a/src/server/sockhelper.h b/src/server/sockhelper.h
index b6b0b11..3f4d485 100644
--- a/src/server/sockhelper.h
+++ b/src/server/sockhelper.h
@@ -76,4 +76,11 @@ void sock_set_block(int sock);
  */
 bool sock_append(poll_list_t *list, const int sock, bool wantRead, bool wantWrite);
 
+/**
+ * Send the whole buffer, calling write() multiple times if neccessary.
+ * Give up after calling write() maxtries times.
+ * Set maxtries < 0 to try infinitely.
+ */
+ssize_t sock_sendAll(int sock, void *buffer, size_t len, int maxtries);
+
 #endif /* SOCKHELPER_H_ */