[SERVER] Close all client sockets on shutdown

3 files changed, 43 insertions, 13 deletions

diff --git a/src/server/net.c b/src/server/net.c
index 87af60f..2d6b9fc 100644
--- a/src/server/net.c
+++ b/src/server/net.c
@@ -19,6 +19,7 @@
  */
 
 #include <stdint.h>
+#include <stdlib.h>
 #include <stdio.h>
 #include <string.h>
 #include <unistd.h>
@@ -33,16 +34,16 @@
 #include "server.h"
 #include "hashtable.h"
 
-void *dnbd3_handle_query(void *client_socket)
+void *dnbd3_handle_query(void *dnbd3_client)
 {
+    dnbd3_client_t *client = (dnbd3_client_t *) (uintptr_t) dnbd3_client;
     int image_file = -1;
     off_t filesize = 0;
-    int sock = (int) (uintptr_t) client_socket;
     dnbd3_request_t request;
     dnbd3_reply_t reply;
     uint16_t cmd;
 
-    while (recv(sock, &request, sizeof(dnbd3_request_t), MSG_WAITALL) > 0)
+    while (recv(client->sock, &request, sizeof(dnbd3_request_t), MSG_WAITALL) > 0)
     {
         cmd = request.cmd;
         switch (cmd)
@@ -50,7 +51,7 @@ void *dnbd3_handle_query(void *client_socket)
         case CMD_PING:
             reply.cmd = request.cmd;
             memcpy(reply.handle, request.handle, sizeof(request.handle));
-            send(sock, (char *) &reply, sizeof(dnbd3_reply_t), 0);
+            send(client->sock, (char *) &reply, sizeof(dnbd3_reply_t), 0);
             break;
 
         case CMD_GET_SIZE:
@@ -70,7 +71,7 @@ void *dnbd3_handle_query(void *client_socket)
             }
             reply.cmd = request.cmd;
             reply.filesize = filesize;
-            send(sock, (char *) &reply, sizeof(dnbd3_reply_t), 0);
+            send(client->sock, (char *) &reply, sizeof(dnbd3_reply_t), 0);
             break;
 
         case CMD_GET_BLOCK:
@@ -79,9 +80,9 @@ void *dnbd3_handle_query(void *client_socket)
 
             reply.cmd = request.cmd;
             memcpy(reply.handle, request.handle, sizeof(request.handle));
-            send(sock, (char *) &reply, sizeof(dnbd3_reply_t), 0);
+            send(client->sock, (char *) &reply, sizeof(dnbd3_reply_t), 0);
 
-            if (sendfile(sock, image_file, (off_t *) &request.offset, request.size) < 0)
+            if (sendfile(client->sock, image_file, (off_t *) &request.offset, request.size) < 0)
                 printf("ERROR: sendfile returned -1\n");
 
             break;
@@ -92,8 +93,10 @@ void *dnbd3_handle_query(void *client_socket)
         }
 
     }
-    close(sock);
-    printf("INFO: Client exit.\n");
+    close(client->sock);
+    _dnbd3_clients = g_slist_remove(_dnbd3_clients, client);
+    free(client);
+    printf("INFO: Client %s exit\n", client->ip);
     pthread_exit((void *) 0);
 }
 
diff --git a/src/server/server.c b/src/server/server.c
index e15249b..ad4c20b 100644
--- a/src/server/server.c
+++ b/src/server/server.c
@@ -24,6 +24,7 @@
 #include <signal.h>
 #include <getopt.h>
 #include <pthread.h>
+#include <string.h>
 
 #include "../types.h"
 #include "../version.h"
@@ -38,6 +39,7 @@ int _sock;
 
 pthread_spinlock_t _spinlock;
 char *_config_file_name = DEFAULT_CONFIG_FILE;
+GSList *_dnbd3_clients = NULL;
 
 void dnbd3_print_help(char* argv_0)
 {
@@ -60,6 +62,17 @@ void dnbd3_print_version()
 
 void dnbd3_cleanup()
 {
+    printf("INFO: Cleanup...\n");
+    GSList *iterator = NULL;
+    for (iterator = _dnbd3_clients; iterator; iterator = iterator->next)
+    {
+        dnbd3_client_t *client =  iterator->data;
+        shutdown(client->sock, SHUT_RDWR);
+        pthread_join(*client->thread, NULL);
+    }
+
+    g_slist_free(_dnbd3_clients);
+
     close(_sock);
     dnbd3_delete_pid_file();
     exit(EXIT_SUCCESS);
@@ -148,15 +161,20 @@ int main(int argc, char* argv[])
             printf("ERROR: Accept failure\n");
             continue;
         }
-        printf("INFO: Client: %s connected\n", inet_ntoa(client.sin_addr));
+        printf("INFO: Client %s connected\n", inet_ntoa(client.sin_addr));
 
         setsockopt(fd, SOL_SOCKET, SO_RCVTIMEO, (char *) &timeout, sizeof(timeout));
         setsockopt(fd, SOL_SOCKET, SO_SNDTIMEO, (char *) &timeout, sizeof(timeout));
 
-        // FIXME: catch SIGKILL/SIGTERM and close all socket before exit
         pthread_t thread;
-        pthread_create(&(thread), NULL, dnbd3_handle_query, (void *) (uintptr_t) fd);
-        pthread_detach(thread);
+        dnbd3_client_t *dnbd3_client = (dnbd3_client_t *) malloc(sizeof(dnbd3_client_t));
+        strcpy(dnbd3_client->ip, inet_ntoa(client.sin_addr));
+        dnbd3_client->sock = fd;
+        dnbd3_client->thread = &thread;
+
+        _dnbd3_clients = g_slist_append (_dnbd3_clients, dnbd3_client);
+
+        pthread_create(&(thread), NULL, dnbd3_handle_query, (void *) (uintptr_t) dnbd3_client);
     }
 
     dnbd3_cleanup();
diff --git a/src/server/server.h b/src/server/server.h
index 3918faf..9372c55 100644
--- a/src/server/server.h
+++ b/src/server/server.h
@@ -22,12 +22,21 @@
 #define SERVER_H_
 
 #include <stdint.h>
+#include <glib-2.0/glib.h>
 
 #include "config.h"
 #include "../types.h"
 
+typedef struct
+{
+    int sock;
+    char ip[16];
+    pthread_t *thread;
+} dnbd3_client_t;
+
 extern pthread_spinlock_t _spinlock;
 extern char *_config_file_name;
+extern GSList *_dnbd3_clients;
 
 void dnbd3_cleanup();