1 files changed, 55 insertions, 16 deletions

diff --git a/src/server/rpc.c b/src/server/rpc.c
index 5dbcafe..5daf20c 100644
--- a/src/server/rpc.c
+++ b/src/server/rpc.c
@@ -9,6 +9,7 @@
 #include "fileutil.h"
 #include "picohttpparser/picohttpparser.h"
 #include "urldecode.h"
+#include "reference.h"
 
 #include <jansson.h>
 #include <sys/types.h>
@@ -43,7 +44,9 @@ _Static_assert( sizeof("test") == 5 && sizeof("test2") == 6, "Stringsize messup
 DEFSTR(STR_CONNECTION, "connection")
 DEFSTR(STR_CLOSE, "close")
 DEFSTR(STR_QUERY, "/query")
+DEFSTR(STR_CACHEMAP, "/cachemap")
 DEFSTR(STR_Q, "q")
+DEFSTR(STR_ID, "id")
 
 static inline bool equals(struct string *s1,struct string *s2)
 {
@@ -75,13 +78,13 @@ static json_int_t randomRunId;
 static pthread_mutex_t aclLock;
 #define MAX_CLIENTS 50
 #define CUTOFF_START 40
-static pthread_mutex_t statusLock;
 static struct {
-	int count;
-	bool overloaded;
+	atomic_int count;
+	atomic_bool overloaded;
 } status;
 
 static bool handleStatus(int sock, int permissions, struct field *fields, size_t fields_num, int keepAlive);
+static bool handleCacheMap(int sock, int permissions, struct field *fields, size_t fields_num, int keepAlive);
 static bool sendReply(int sock, const char *status, const char *ctype, const char *payload, ssize_t plen, int keepAlive);
 static void parsePath(struct string *path, struct string *file, struct field *getv, size_t *getc);
 static bool hasHeaderValue(struct phr_header *headers, size_t numHeaders, struct string *name, struct string *value);
@@ -91,8 +94,7 @@ static void loadAcl();
 
 void rpc_init()
 {
-	mutex_init( &aclLock );
-	mutex_init( &statusLock );
+	mutex_init( &aclLock, LOCK_RPC_ACL );
 	randomRunId = (((json_int_t)getpid()) << 16) | (json_int_t)time(NULL);
 	// </guard>
 	if ( sizeof(randomRunId) > 4 ) {
@@ -123,10 +125,8 @@ void rpc_sendStatsJson(int sock, dnbd3_host_t* host, const void* data, const int
 		return;
 	}
 	do {
-		mutex_lock( &statusLock );
 		const int curCount = ++status.count;
 		UPDATE_LOADSTATE( curCount );
-		mutex_unlock( &statusLock );
 		if ( curCount > MAX_CLIENTS ) {
 			sendReply( sock, "503 Service Temporarily Unavailable", "text/plain", "Too many HTTP clients", -1, HTTP_CLOSE );
 			goto func_return;
@@ -141,13 +141,13 @@ void rpc_sendStatsJson(int sock, dnbd3_host_t* host, const void* data, const int
 	bool hasName = false;
 	bool ok;
 	int keepAlive = HTTP_KEEPALIVE;
-	do {
+	while ( !_shutdown ) {
 		// Read request from client
 		struct phr_header headers[100];
 		size_t numHeaders, prevLen = 0, consumed;
 		struct string method, path;
 		int minorVersion;
-		do {
+		while ( !_shutdown ) {
 			// Parse before calling recv, there might be a complete pipelined request in the buffer already
 			// If the request is incomplete, we allow exactly one additional recv() to complete it.
 			// This should suffice for real world scenarios as I don't know of any
@@ -192,15 +192,15 @@ void rpc_sendStatsJson(int sock, dnbd3_host_t* host, const void* data, const int
 				sendReply( sock, "400 Bad Request", "text/plain", "Server cannot understand what you're trying to say", -1, HTTP_CLOSE );
 				goto func_return;
 			}
-		} while ( true );
+		} // Loop while request header incomplete
+		if ( _shutdown )
+			break;
 		if ( keepAlive == HTTP_KEEPALIVE ) {
 			// Only keep the connection alive (and indicate so) if the client seems to support this
 			if ( minorVersion == 0 || hasHeaderValue( headers, numHeaders, &STR_CONNECTION, &STR_CLOSE ) ) {
 				keepAlive = HTTP_CLOSE;
 			} else { // And if there aren't too many active HTTP sessions
-				mutex_lock( &statusLock );
 				if ( status.overloaded ) keepAlive = HTTP_CLOSE;
-				mutex_unlock( &statusLock );
 			}
 		}
 		if ( method.s != NULL && path.s != NULL ) {
@@ -216,10 +216,13 @@ void rpc_sendStatsJson(int sock, dnbd3_host_t* host, const void* data, const int
 			// Don't care if GET or POST
 			if ( equals( &file, &STR_QUERY ) ) {
 				ok = handleStatus( sock, permissions, getv, getc, keepAlive );
+			} else if ( equals( &file, &STR_CACHEMAP ) ) {
+				ok = handleCacheMap( sock, permissions, getv, getc, keepAlive );
 			} else {
 				ok = sendReply( sock, "404 Not found", "text/plain", "Nothing", -1, keepAlive );
 			}
-			if ( !ok ) break;
+			if ( !ok )
+				break;
 		}
 		// hoff might be beyond end if the client sent another request (burst)
 		const ssize_t extra = hoff - consumed;
@@ -231,13 +234,11 @@ void rpc_sendStatsJson(int sock, dnbd3_host_t* host, const void* data, const int
 			hasName = true;
 			setThreadName( "HTTP" );
 		}
-	} while (true);
+	} // Loop while more requests
 func_return:;
 	do {
-		mutex_lock( &statusLock );
 		const int curCount = --status.count;
 		UPDATE_LOADSTATE( curCount );
-		mutex_unlock( &statusLock );
 	} while (0);
 }
 
@@ -347,6 +348,44 @@ static bool handleStatus(int sock, int permissions, struct field *fields, size_t
 	return ok;
 }
 
+static bool handleCacheMap(int sock, int permissions, struct field *fields, size_t fields_num, int keepAlive)
+{
+	if ( !(permissions & ACL_IMAGE_LIST) ) {
+		return sendReply( sock, "403 Forbidden", "text/plain", "No permission to access image list", -1, keepAlive );
+	}
+	int imgId = -1;
+	static const char one = (char)0xff;
+	for (size_t i = 0; i < fields_num; ++i) {
+		if ( equals( &fields[i].name, &STR_ID ) ) {
+			char *broken;
+			imgId = (int)strtol( fields[i].value.s, &broken, 10 );
+			if ( broken != fields[i].value.s )
+				break;
+			imgId = -1;
+		}
+	}
+	if ( imgId == -1 )
+		return sendReply( sock, "400 Bad Request", "text/plain", "Missing parameter 'id'", -1, keepAlive );
+	dnbd3_image_t *image = image_byId( imgId );
+	if ( image == NULL )
+		return sendReply( sock, "404 Not found", "text/plain", "Image not found", -1, keepAlive );
+	dnbd3_cache_map_t *cache = ref_get_cachemap( image );
+	image_release( image );
+	int len;
+	const char *map;
+	if ( cache == NULL ) {
+		map = &one;
+		len = 1;
+	} else {
+		_Static_assert( sizeof(const char) == sizeof(_Atomic uint8_t), "Atomic assumption exploded" );
+		map = (const char*)cache->map;
+		len = IMGSIZE_TO_MAPBYTES( image->virtualFilesize );
+	}
+	bool ok = sendReply( sock, "200 OK", "application/octet-stream", map, len, keepAlive );
+	ref_put( &cache->reference );
+	return ok;
+}
+
 static bool sendReply(int sock, const char *status, const char *ctype, const char *payload, ssize_t plen, int keepAlive)
 {
 	if ( plen == -1 ) plen = strlen( payload );