Cleanup, react to signals (\^C), tidy up and sort code

1 files changed, 31 insertions, 7 deletions

diff --git a/src/main.go b/src/main.go
index aac4f02..4833896 100644
--- a/src/main.go
+++ b/src/main.go
@@ -3,6 +3,9 @@ package main
 import (
 	"flag"
 	"log"
+	"net/http"
+	"os"
+	"os/signal"
 )
 
 func main() {
@@ -28,12 +31,33 @@ func main() {
 	if startVPNError != nil {
 		log.Fatalln(startVPNError.Error())
 	}
-	defer stopVPN()
 
-	// Run HTTP server
+	// Set up HTTP server
 	log.Printf("Running HTTP server on port %d", *port)
-	runHTTPServerError := runHTTPServer()
-	if runHTTPServerError != nil {
-		log.Fatalln(runHTTPServerError)
-	}
-}
\ No newline at end of file
+	httpServer := setupHTTPServer()
+
+	// React to system signals
+	signalChannel := make(chan os.Signal, 1)
+	signal.Notify(signalChannel)
+	go cleanup(signalChannel, httpServer)
+
+	// Start our HTTP server
+	log.Fatalln(httpServer.ListenAndServe())
+}
+
+func cleanup(signalChannel chan os.Signal, server http.Server) {
+	// Block until we receive a signal
+	<- signalChannel
+	log.Println("Requested shutdown")
+
+	// Stop everything
+	log.Println("Stopping HTTP server")
+	server.Close()
+	log.Println("Stopping VPN container")
+	stopVPN()
+	log.Println("Deleting Docker networks")
+	deleteNetwork(vpnHostNetworkName)
+	deleteNetwork(containerNetworkName)
+	// TODO: stop challenge containers
+	os.Exit(0)
+}