From fe354d9f7471234c73bb7eab5f127fecb099e6ef Mon Sep 17 00:00:00 2001
From: Bogomil Vasilev <smirky@smirky.net>
Date: Sun, 13 Jan 2019 12:26:15 +0200
Subject: [PATCH] agent: add signal handlers

---
 agent/agent.c | 75 ++++++++++++++++++++++++++++++++++++++++++++-------
 agent/job.c   | 11 +++++---
 2 files changed, 72 insertions(+), 14 deletions(-)

diff --git a/agent/agent.c b/agent/agent.c
index 5a43485..f1788a7 100644
--- a/agent/agent.c
+++ b/agent/agent.c
@@ -23,14 +23,20 @@
 #include <unistd.h>
 #include <pthread.h>
 #include <string.h>
+#include <signal.h>
 #include "job.h"
 #include "../src/protocol.h"
 #include "agent_ssl.h"
 
 #define FAIL -1
 
+static SSL_CTX *ctx;
+static SSL *ssl;
+static int server;
 static struct job_args *args;
 static pthread_t *job_thread;
+static void agent_shutdown(void);
+static void signal_handler(int sig);
 
 static short get_job_slot();
 
@@ -52,14 +58,64 @@ static short find_job(unsigned id)
 	return -1;
 }
 
+static void signal_handler(int sig)
+{
+	switch (sig) {
+	case SIGHUP:
+		printf("Received SIGHUP signal. Ignoring...\n");
+		break;
+	case SIGINT:
+	case SIGTERM:
+		printf("Received SIGTERM signal!\n");
+		agent_shutdown();
+		printf("Agent has been stopped properly.\n");
+		_exit(EXIT_SUCCESS);
+		break;
+	default:
+		printf("Unhandled signal %s", strsignal(sig));
+		break;
+	}
+}
+
+static void agent_shutdown(void)
+{
+	printf("Shutting down agent...\n");
+	SSL_shutdown(ssl);
+	SSL_free(ssl);
+	SSL_CTX_free(ctx);		/* release context */
+	close(server);			/* close socket */
+}
+
+static void set_env(void)
+{
+	struct sigaction new_sigaction;
+	sigset_t new_sigset;
+
+	/* Set signal mask - signals we want to block */
+	sigemptyset(&new_sigset);
+	sigaddset(&new_sigset, SIGCHLD);  /* ignore child */
+	sigaddset(&new_sigset, SIGTSTP);  /* ignore Tty stop signals */
+	sigaddset(&new_sigset, SIGTTOU);  /* ignore Tty background writes */
+	sigaddset(&new_sigset, SIGTTIN);  /* ignore Tty background reads */
+	sigprocmask(SIG_BLOCK, &new_sigset, NULL); /* Block above signals */
+
+	/* Set up a signal handler */
+	new_sigaction.sa_handler = signal_handler;
+	sigemptyset(&new_sigaction.sa_mask);
+	new_sigaction.sa_flags = 0;
+
+	/* Signals to handle */
+	sigaction(SIGHUP, &new_sigaction, NULL);  /* catch hangup signal */
+	sigaction(SIGTERM, &new_sigaction, NULL); /* catch term signal */
+	sigaction(SIGINT, &new_sigaction, NULL);  /* catch interrupt signal */
+	signal(SIGPIPE, SIG_IGN); /* prevent crashing from bad writes */
+}
+
 int main(int count, char *strings[])
-{   
-	SSL_CTX *ctx;
-	int server;
-	SSL *ssl;
+{
 	int bytes, i;
 	char *hostname, *portnum;
-    
+
 	if (count != 6) {
 		printf("usage: %s <hostname> <portnum> <key> <cert> <CA>\n", strings[0]);
 		_exit(EXIT_SUCCESS);
@@ -76,7 +132,7 @@ int main(int count, char *strings[])
 	load_certs(ctx, strings[3], strings[4], strings[5]);
 	ssl = SSL_new(ctx);		/* create new SSL connection state */
 	SSL_set_fd(ssl, server);	/* attach the socket descriptor */
-    
+
 	if (SSL_connect(ssl) == FAIL) {	/* perform the connection */
 		ERR_print_errors_fp(stderr);
 		close(server);
@@ -85,6 +141,8 @@ int main(int count, char *strings[])
 	}
 	printf("Connected with %s encryption\n", SSL_get_cipher(ssl));
 	show_certs(ssl);
+	set_env();
+	atexit(agent_shutdown);
 	if (!(args = calloc(1, sizeof(*args) * MAX_AGENT_JOBS))) {
 		fprintf( stderr,
 			 "Failed to calloc() %d bytes for job_args! Exiting...\n",
@@ -216,11 +274,8 @@ int main(int count, char *strings[])
 					memcpy(&args[index].buf, &buf, sizeof(struct msg_t));
 			}
 		}
-		
+
 		SSL_shutdown(ssl);
 		SSL_free(ssl);		/* release connection state */
 	} while (bytes);
-	close(server);			/* close socket */
-	SSL_CTX_free(ctx);		/* release context */
 }
-
diff --git a/agent/job.c b/agent/job.c
index 12d2e0e..7a45e63 100644
--- a/agent/job.c
+++ b/agent/job.c
@@ -104,17 +104,20 @@ void* get_os(void *args)
 	system(cmd); /* TODO: set proper path */
 	if ((fd = open(dump_path, O_RDONLY)) == -1) {
 		fprintf(stderr, "Cannot open file: %s\n", dump_path);
-		job->buf.meta.len = sprintf((char*)job->buf.chunk.data, "Unknown");
-		SSL_write(job->ssl, &job->buf, sizeof(struct msg_t));
+		job->buf.meta.len = sprintf((char*)job->buf.chunk.data,
+				"I/O error");
+		//SSL_write(job->ssl, &job->buf, sizeof(struct msg_t));
 	} else {
 		job->buf.meta.len = read(fd, job->buf.chunk.data, sizeof(job->buf.chunk.data));
 		if ((signed)job->buf.meta.len == -1) {
 			perror("Failed to read from file: ");
 		}
 		close(fd);
-		SSL_write(job->ssl, &job->buf, sizeof(struct msg_t));
-		job->slot = FREE;
+		//SSL_write(job->ssl, &job->buf, sizeof(struct msg_t));
+		//job->slot = FREE;
 	}
+	SSL_write(job->ssl, &job->buf, sizeof(struct msg_t));
+	job->slot = FREE;
 	unlink(dump_path);
 	return 0;
 }