agent: add signal handlers

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,11 +58,61 @@ 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;
 
@@ -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",
@@ -220,7 +278,4 @@ int main(int count, char *strings[])
 		SSL_shutdown(ssl);
 		SSL_free(ssl);		/* release connection state */
 	} while (bytes);
-	close(server);			/* close socket */
-	SSL_CTX_free(ctx);		/* release context */
 }
-

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;
-	}
 	unlink(dump_path);
 	return 0;
 }