[SERVER] Fix use-after-free, improve cleanup

6 files changed, 90 insertions, 36 deletions

diff --git a/src/server/globals.h b/src/server/globals.h
index b77c547..1dea849 100644
--- a/src/server/globals.h
+++ b/src/server/globals.h
@@ -113,6 +113,7 @@ struct _dnbd3_client
 	dnbd3_image_t *image;
 	pthread_spinlock_t lock;
 	pthread_mutex_t sendMutex;
+	int running;
 };
 
 // #######################################################
@@ -143,7 +144,6 @@ extern int _serverPenalty;
  */
 extern int _clientPenalty;
 
-
 extern int _shutdown;
 
 void globals_loadConfig();
diff --git a/src/server/image.c b/src/server/image.c
index af752a9..1393d5b 100644
--- a/src/server/image.c
+++ b/src/server/image.c
@@ -33,7 +33,7 @@ static int64_t image_pad(const char *path, const int64_t currentSize);
 /**
  * Returns TRUE if the given image is complete
  */
-int image_is_complete(dnbd3_image_t *image)
+int image_isComplete(dnbd3_image_t *image)
 {
 	assert( image != NULL );
 	if ( image->working && image->cache_map == NULL ) {
@@ -68,7 +68,7 @@ int image_is_complete(dnbd3_image_t *image)
  * Update cache-map of given image for the given byte range
  * Locks on: images[].lock
  */
-void image_update_cachemap(dnbd3_image_t *image, uint64_t start, uint64_t end, const int set)
+void image_updateCachemap(dnbd3_image_t *image, uint64_t start, uint64_t end, const int set)
 {
 	assert( image != NULL );
 	// This should always be block borders due to how the protocol works, but better be safe
@@ -115,13 +115,43 @@ void image_update_cachemap(dnbd3_image_t *image, uint64_t start, uint64_t end, c
 }
 
 /**
+ * Mark image as complete by freeing the cache_map and deleting the map file on disk
+ * DOES NOT LOCK ON THE IMAGE
+ */
+void image_markComplete(dnbd3_image_t *image)
+{
+	assert( image != NULL );
+	if ( image->cache_map == NULL ) return;
+	free( image->cache_map );
+	image->cache_map = NULL;
+	char mapfile[strlen( image->path ) + 4 + 1];
+	sprintf( mapfile, "%s.map", image->path );
+	remove( mapfile );
+}
+
+/**
  * Saves the cache map of the given image.
  * Return TRUE on success.
- * DOES NOT lock
+ * Locks on: image.lock
  */
-int image_save_cache_map(dnbd3_image_t *image)
+int image_saveCacheMap(dnbd3_image_t *image)
 {
 	if ( image == NULL || image->cache_map == NULL ) return TRUE;
+	spin_lock( &image->lock );
+	// Lock and get a copy of the cache map, as it could be freed by another thread that is just about to
+	// figure out that this image's cache copy is complete
+	if ( image->cache_map == NULL || image->filesize < DNBD3_BLOCK_SIZE ) {
+		spin_unlock( &image->lock );
+		return TRUE;
+	}
+	const size_t size = IMGSIZE_TO_MAPBYTES(image->filesize);
+	uint8_t *map = malloc( size );
+	memcpy( map, image->cache_map, size );
+	// Unlock. Use path and cacheFd without locking. path should never change after initialization of the image,
+	// cacheFd is written to and we don't hold a spinlock during I/O
+	// By increasing the user count we make sure the image is not freed in the meantime
+	image->users++;
+	spin_unlock( &image->lock );
 	assert( image->path != NULL );
 	char mapfile[strlen( image->path ) + 4 + 1];
 	int fd;
@@ -129,15 +159,23 @@ int image_save_cache_map(dnbd3_image_t *image)
 	strcat( mapfile, ".map" );
 
 	fd = open( mapfile, O_WRONLY | O_CREAT, 0640 );
-	if ( fd < 0 ) return FALSE;
+	if ( fd < 0 ) {
+		spin_lock( &image->lock );
+		image->users--;
+		spin_unlock( &image->lock );
+		return FALSE;
+	}
 
-	write( fd, image->cache_map, ((image->filesize + (1 << 15) - 1) >> 15) * sizeof(char) );
+	write( fd, map, ((image->filesize + (1 << 15) - 1) >> 15) * sizeof(char) );
 	if ( image->cacheFd != -1 ) {
 		fsync( image->cacheFd );
 	}
 	fsync( fd );
 	close( fd );
 
+	spin_lock( &image->lock );
+	image->users--;
+	spin_unlock( &image->lock );
 	return TRUE;
 }
 
@@ -266,14 +304,14 @@ void image_killUplinks()
 
 /**
  * Free image. DOES NOT check if it's in use.
- * Indirectly locks on uplink.queueLock
+ * Indirectly locks on image.lock, uplink.queueLock
  * DO NOT lock on the image when calling.
  */
 dnbd3_image_t* image_free(dnbd3_image_t *image)
 {
 	assert( image != NULL );
 	//
-	image_save_cache_map( image );
+	image_saveCacheMap( image );
 	uplink_shutdown( image );
 	free( image->cache_map );
 	free( image->crc32 );
@@ -291,7 +329,7 @@ dnbd3_image_t* image_free(dnbd3_image_t *image)
  * Load all images in given path recursively.
  * Pass NULL to use path from config.
  */
-int image_load_all(char *path)
+int image_loadAll(char *path)
 {
 	if ( path == NULL ) {
 		return image_load_all_internal( _basePath, _basePath );
@@ -523,10 +561,8 @@ static int image_try_load(char *base, char *path)
 	image->working = (image->cache_map == NULL );
 	spin_init( &image->lock, PTHREAD_PROCESS_PRIVATE );
 	// Get rid of cache map if image is complete
-	if ( image->cache_map != NULL && image_is_complete( image ) ) {
-		remove( mapFile );
-		free( image->cache_map );
-		image->cache_map = NULL;
+	if ( image->cache_map != NULL && image_isComplete( image ) ) {
+		image_markComplete( image );
 		image->working = TRUE;
 	}
 	if ( image->cache_map != NULL ) {
@@ -643,7 +679,7 @@ int image_create(char *image, int revision, uint64_t size)
  * This function wants a plain file name instead of a dnbd3_image_t,
  * as it can be used directly from the command line.
  */
-int image_generate_crc_file(char *image)
+int image_generateCrcFile(char *image)
 {
 	int fdImage = open( image, O_RDONLY );
 	if ( fdImage < 0 ) {
diff --git a/src/server/image.h b/src/server/image.h
index 66d2563..e1eddd3 100644
--- a/src/server/image.h
+++ b/src/server/image.h
@@ -8,11 +8,13 @@ extern dnbd3_image_t *_images[SERVER_MAX_IMAGES];
 extern int _num_images;
 extern pthread_spinlock_t _images_lock;
 
-int image_is_complete(dnbd3_image_t *image);
+int image_isComplete(dnbd3_image_t *image);
 
-void image_update_cachemap(dnbd3_image_t *image, uint64_t start, uint64_t end, const int set);
+void image_updateCachemap(dnbd3_image_t *image, uint64_t start, uint64_t end, const int set);
 
-int image_save_cache_map(dnbd3_image_t *image);
+void image_markComplete(dnbd3_image_t *image);
+
+int image_saveCacheMap(dnbd3_image_t *image);
 
 dnbd3_image_t* image_get(char *name, uint16_t revision);
 
@@ -22,11 +24,11 @@ void image_killUplinks();
 
 dnbd3_image_t* image_free(dnbd3_image_t *image);
 
-int image_load_all(char *path);
+int image_loadAll(char *path);
 
 int image_create(char *image, int revision, uint64_t size);
 
-int image_generate_crc_file(char *image);
+int image_generateCrcFile(char *image);
 
 // one byte in the map covers 8 4kib blocks, so 32kib per byte
 // "+ (1 << 15) - 1" is required to account for the last bit of
diff --git a/src/server/net.c b/src/server/net.c
index f0f9bf8..8309672 100644
--- a/src/server/net.c
+++ b/src/server/net.c
@@ -351,8 +351,8 @@ void *net_client_handler(void *dnbd3_client)
 	}
 	exit_client_cleanup: ;
 	if ( image_file != -1 ) close( image_file );
+	client->running = FALSE;
 	dnbd3_remove_client( client );
 	client = dnbd3_free_client( client );
-	pthread_detach( client->thread );
 	return NULL ;
 }
diff --git a/src/server/server.c b/src/server/server.c
index 48a4f1e..7571cbc 100644
--- a/src/server/server.c
+++ b/src/server/server.c
@@ -101,7 +101,7 @@ void dnbd3_print_version()
  */
 void dnbd3_cleanup()
 {
-	int i;
+	int i, count;
 
 	_shutdown = TRUE;
 	debug_locks_stop_watchdog();
@@ -124,13 +124,30 @@ void dnbd3_cleanup()
 		dnbd3_client_t * const client = _clients[i];
 		spin_lock( &client->lock );
 		if ( client->sock >= 0 ) shutdown( client->sock, SHUT_RDWR );
-		if ( client->thread != 0 ) pthread_join( client->thread, NULL );
-		_clients[i] = NULL;
+		client->sock = -1;
 		spin_unlock( &client->lock );
-		free( client );
 	}
-	_num_clients = 0;
 	spin_unlock( &_clients_lock );
+	count = -1;
+	while ( count != 0 ) {
+		count = 0;
+		spin_lock( &_clients_lock );
+		for (i = 0; i < _num_clients; ++i) {
+			if ( _clients[i] == NULL ) continue;
+			if ( _clients[i]->running ) {
+				count++;
+			} else {
+				_clients[i] = NULL;
+				dnbd3_free_client( _clients[i] );
+			}
+		}
+		spin_unlock( &_clients_lock );
+		if ( count != 0 ) {
+			printf( "%d clients still active...\n", count );
+			sleep( 1 );
+		}
+	}
+	_num_clients = 0;
 
 	// Clean up images
 	spin_lock( &_images_lock );
@@ -200,7 +217,7 @@ int main(int argc, char *argv[])
 			dnbd3_print_version();
 			break;
 		case 'crc4':
-			return image_generate_crc_file( optarg ) ? 0 : EXIT_FAILURE;
+			return image_generateCrcFile( optarg ) ? 0 : EXIT_FAILURE;
 		case 'asrt':
 			printf( "Testing a failing assertion:\n" );
 			assert( 4 == 5 );
@@ -262,7 +279,7 @@ int main(int argc, char *argv[])
 
 	printf( "Loading images....\n" );
 	// Load all images in base path
-	if ( !image_load_all( NULL ) ) {
+	if ( !image_loadAll( NULL ) ) {
 		printf( "[ERROR] Could not load images.\n" );
 		return EXIT_FAILURE;
 	}
@@ -320,6 +337,7 @@ int main(int argc, char *argv[])
 			dnbd3_client = dnbd3_free_client( dnbd3_client );
 			continue;
 		}
+		pthread_detach( dnbd3_client->thread );
 	}
 
 	dnbd3_cleanup();
@@ -332,7 +350,7 @@ int main(int argc, char *argv[])
 dnbd3_client_t* dnbd3_init_client(struct sockaddr_storage *client, int fd)
 {
 	dnbd3_client_t *dnbd3_client = calloc( 1, sizeof(dnbd3_client_t) );
-	if ( dnbd3_client == NULL ) {
+	if ( dnbd3_client == NULL ) { // This will never happen thanks to memory overcommit
 		memlogf( "[ERROR] Could not alloc dnbd3_client_t for new client." );
 		return NULL ;
 	}
@@ -352,6 +370,7 @@ dnbd3_client_t* dnbd3_init_client(struct sockaddr_storage *client, int fd)
 		free( dnbd3_client );
 		return NULL ;
 	}
+	dnbd3_client->running = TRUE;
 	dnbd3_client->sock = fd;
 	spin_init( &dnbd3_client->lock, PTHREAD_PROCESS_PRIVATE );
 	pthread_mutex_init( &dnbd3_client->sendMutex, NULL );
@@ -441,5 +460,5 @@ static void dnbd3_handle_sigterm(int signum)
 void dnbd3_handle_sigusr1(int signum)
 {
 	memlogf( "INFO: SIGUSR1 (%s) received, re-scanning image directory", strsignal( signum ) );
-	image_load_all( NULL );
+	image_loadAll( NULL );
 }
diff --git a/src/server/uplink.c b/src/server/uplink.c
index fea48a0..4dbe75a 100644
--- a/src/server/uplink.c
+++ b/src/server/uplink.c
@@ -261,13 +261,10 @@ static void* uplink_mainloop(void *data)
 				nextAltCheck = now + SERVER_RTT_DELAY_MAX;
 			} else if ( now >= nextAltCheck ) {
 				// It seems it's time for a check
-				if ( image_is_complete( link->image ) ) {
+				if ( image_isComplete( link->image ) ) {
 					// Quit work if image is complete
 					if ( spin_trylock( &link->image->lock ) == 0 ) {
-						if ( link->image->cache_map != NULL ) {
-							free( link->image->cache_map );
-							link->image->cache_map = NULL;
-						}
+						image_markComplete(link->image);
 						link->image->uplink = NULL;
 						link->shutdown = TRUE;
 						free( link->recvBuffer );
@@ -382,7 +379,7 @@ static void uplink_handle_receive(dnbd3_connection_t *link)
 		memlogf( "[ERROR] lseek() failed when writing to cache for %s", link->image->path );
 	} else {
 		ret = (int)write( link->image->cacheFd, link->recvBuffer, reply.size );
-		if ( ret > 0 ) image_update_cachemap( link->image, start, start + ret, TRUE );
+		if ( ret > 0 ) image_updateCachemap( link->image, start, start + ret, TRUE );
 	}
 	// 2) Figure out which clients are interested in it
 	struct iovec iov[2];