1 files changed, 37 insertions, 10 deletions

diff --git a/multipart-replace-http-server.c b/multipart-replace-http-server.c
index 29b0ff5..726714f 100644
--- a/multipart-replace-http-server.c
+++ b/multipart-replace-http-server.c
@@ -14,6 +14,7 @@
 #include <unistd.h>
 
 #include "util.h"
+#include "wg.h"
 #include "multipart-replace.h"
 
 struct httpfile {
@@ -31,6 +32,20 @@ struct reader_thread_args {
 	const char *boundary;
 };
 
+
+bool running = true;
+int sock;
+struct wg wg;
+
+void stop(int sig) {
+	log("Caught %d", sig);
+	log("STOPPING=1");
+	running = false;
+	close(sock);
+}
+
+#define threaderror(stat, errnum, ...) do { error(0, errnum, __VA_ARGS__); stop(0); } while(0)
+
 void *reader_thread(void *args_anon) {
 	struct reader_thread_args *args = args_anon;
 
@@ -42,7 +57,8 @@ void *reader_thread(void *args_anon) {
 
 	multipart_replace_reader(args->stream, args->fd, args->boundary);
 	sleep(5);
-	error(EXIT_FAILURE, 0, "multipart_replace stream ended: %s", args->filename);
+	threaderror(EXIT_FAILURE, 0, "multipart_replace stream ended: %s", args->filename);
+	wg_sub(&wg, 1);
 	return NULL;
 }
 
@@ -50,9 +66,10 @@ void start_multipart_replace_reader(struct multipart_replace_stream *s, int fd,
 	struct reader_thread_args *args = xrealloc(NULL, sizeof(struct reader_thread_args));
 	args->stream = s;
 	args->fd = fd;
-	args->filename = filename;
+	args->filename = strdup(filename);
 	args->boundary = boundary;
 	pthread_t thread;
+	wg_add(&wg, 1);
 	pthread_create(&thread, NULL, reader_thread, args);
 }
 
@@ -60,8 +77,10 @@ void file_add(const char *filename) {
 	struct multipart_replace_stream *stream = xrealloc(NULL, sizeof(struct multipart_replace_stream));
 	init_multipart_replace_stream(stream);
 	int fd = open(filename, O_RDONLY);
-	if (fd < 0)
-		error(EXIT_FAILURE, errno, "opening file: %s", filename);
+	if (fd < 0) {
+		threaderror(EXIT_FAILURE, errno, "opening file: %s", filename);
+		return;
+	}
 	filev = xrealloc(filev, (++filec)*sizeof(*filev));
 	char *shortname = strrchr(filename, '/');
 	if (shortname == NULL) {
@@ -161,6 +180,7 @@ void *connection_thread(void *arg_anon) {
 	log("Connection %d opened", fd);
 	connection_handler(fd);
 	log("Connection %d closed", fd);
+	wg_sub(&wg, 1);
 	return NULL;
 }
 
@@ -286,7 +306,6 @@ void usage() {
 }
 
 void cleanup() {
-	log("STOPPING=1");
 	fflush(stderr);
 	sleep(5); /* work around systemd bug dropping log messages */
 }
@@ -316,25 +335,33 @@ int main(int argc, char *argv[]) {
 	error_print_progname = progname;
 	atexit(cleanup);
 
-	int sock = sockstream_listen(argv[1], argv[2]);
+	sock = sockstream_listen(argv[1], argv[2]);
 	if (sock < 0) {
 		if (sock == -EAI_SYSTEM)
-			error(1, errno, "Opening socket %s %s", argv[1], argv[2]);
+			error(EXIT_FAILURE, errno, "Opening socket %s %s", argv[1], argv[2]);
 		else
-			error(1, 0, "Opening socket %s %s: %s", argv[1], argv[2], gai_strerror(-sock));
+			error(EXIT_FAILURE, 0, "Opening socket %s %s: %s", argv[1], argv[2], gai_strerror(-sock));
 	}
 
+	wg_init(&wg);
 	for (int i = 3; i < argc; i++)
 		file_add(argv[i]);
 
 	signal(SIGPIPE, SIG_IGN);
-	while (1) {
+	signal(SIGTERM, stop);
+	signal(SIGQUIT, stop);
+	signal(SIGINT, stop);
+	while (running) {
 		int conn = accept(sock, NULL, NULL);
 		if (conn < 0)
 			continue;
 		pthread_t thread;
+		wg_add(&wg, 1);
 		pthread_create(&thread, NULL, connection_thread, (void*)(intptr_t)conn);
 	}
-
+	wg_wait(&wg);
+	for (size_t i = 0; i < filec; i++)
+		destroy_multipart_replace_stream(filev[i].stream);
+	free(filev);
 	return 0;
 }