Ensure half-written files are removed on SIGINT

1 files changed, 32 insertions, 0 deletions

diff --git a/util.c b/util.c
index 91cb060..6439390 100644
--- a/util.c
+++ b/util.c
@@ -7,6 +7,7 @@
 #ifndef WINVER
 #include <netdb.h>
 #include <netinet/in.h>
+#include <signal.h>
 #include <sys/socket.h>
 #else
 #include <winsock2.h>
@@ -199,6 +200,17 @@ int post(char *host, int ssl, char *path, char *sendcookie, char *savecookie, ch
 	return request(host, ssl, h, savecookie, body, istext);
 }
 
+// used by cleanupanddie() to clean up half-written files in the
+// case of an interrupt during a write.
+char *curfn;
+
+void cleanupanddie()
+{
+	unlink(curfn);
+	exit(0);
+}
+
+
 int gettofile(char *host, int ssl, char *url, char *sendcookie, char *savecookie, char *savepath, int istext)
 {
 	char *buf = 0;
@@ -209,6 +221,18 @@ int gettofile(char *host, int ssl, char *url, char *sendcookie, char *savecookie
 		fprintf(stderr, "Could not download %s\n", url);
 		return 1;
 	}
+
+#ifndef WINVER
+	struct sigaction a;
+	curfn = savepath;
+	a.sa_handler = cleanupanddie;
+	sigemptyset(&a.sa_mask);
+	if(sigaction(SIGINT, &a, NULL) != 0) {
+		fprintf(stderr, "Error setting up signal handler\n");
+		free(buf); return 1;
+	}
+#endif
+
 	if((f = fopen(savepath, "wb")) == NULL) {
 		fprintf(stderr, "Could not create file %s\n", savepath);
 		free(buf); return 1;
@@ -223,6 +247,14 @@ int gettofile(char *host, int ssl, char *url, char *sendcookie, char *savecookie
 	free(buf);
 	fclose(f);
 
+#ifndef WINVER
+	a.sa_handler = SIG_DFL;
+	if(sigaction(SIGINT, &a, NULL) != 0) {
+		fprintf(stderr, "Error resetting signal handler\n");
+		return 1;
+	}
+#endif
+
 	return 0;
 }