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

5 files changed, 162 insertions, 75 deletions

diff --git a/src/access.go b/src/access.go
index 4ae96ed..6f0d848 100644
--- a/src/access.go
+++ b/src/access.go
@@ -16,13 +16,10 @@ import (
 )
 
 const(
-	vpnHostNetworkName = "vpnhostnet"
 	vpnContainerName = "circus-vpn"
 )
 
 var vpnContainerID string
-var vpnNetworkID string
-var vpnHostNetworkID string
 var remoteAddress* string
 var remotePort* int
 
@@ -32,24 +29,35 @@ func registerAccessFlags() {
 }
 
 func startVPN() (err error) {
+	// First: try to stop an existing container with that name, to avoid collision
+	// This container may exist if we didn't properly clean up on the last run, e.g. if the companion crashed.
+	stopVPN()
+	// stopVPN() as first command in startVPN()... This looks strange, but is perfectly fine ;)
+
 	// Set up our context and Docker CLI connection
 	setupContext()
 	setupDockerCLI()
-	// Set up network
-	err = setupNetwork()
 
-	if(err != nil) {
-		return err
+	// Set up VPN host network
+	if vpnHostNetworkID == "" {
+		id, err := setupNetwork(vpnHostNetworkName, false)
+		if (err != nil) {
+			return err
+		}
+		vpnHostNetworkID = id
 	}
 
-	err = setupVPNHostNetwork()
-
-	if err != nil {
-		return err
+	// Set up container network
+	if containerNetworkID == "" {
+		id, err := setupNetwork(containerNetworkName, true)
+		if (err != nil) {
+			return err
+		}
+		containerNetworkID = id
 	}
 
 	// Get subnet of challenge container network, to hand it over to our VPN container for routes
-	inspectResp, err := dockerCli.NetworkInspect(dockerCtx, vpnNetworkID, types.NetworkInspectOptions{})
+	inspectResp, err := dockerCli.NetworkInspect(dockerCtx, containerNetworkID, types.NetworkInspectOptions{})
 	if err != nil {
 		return err
 	}
@@ -95,7 +103,7 @@ func startVPN() (err error) {
 	}
 
 	// Attach container network to VPN container
-	err = dockerCli.NetworkConnect(dockerCtx, vpnNetworkID, resp.ID, &network.EndpointSettings{})
+	err = dockerCli.NetworkConnect(dockerCtx, containerNetworkID, resp.ID, &network.EndpointSettings{})
 	if err != nil {
 		return err
 	}
@@ -110,7 +118,7 @@ func startVPN() (err error) {
 	// However, getCertificate() requires that port 9999 of the VPN container hosts the configuration files for our client.
 	// That means we need to attach our own container - thanks to --privileged mode - into the VPN container network.
 	// We get the ID of our container from the "hostname" environment variable. That's a bit dirty, but works for the moment. TODO: solve this better.
-	err = dockerCli.NetworkConnect(dockerCtx, vpnNetworkID, os.Getenv("HOSTNAME"), &network.EndpointSettings{})
+	err = dockerCli.NetworkConnect(dockerCtx, vpnHostNetworkID, os.Getenv("HOSTNAME"), &network.EndpointSettings{})
 	if err != nil {
 		return err
 	}
@@ -130,44 +138,6 @@ func stopVPN() {
 	vpnContainerID = ""
 }
 
-func setupNetwork() (error) {
-	setupContext()
-	setupDockerCLI()
-
-	if vpnNetworkID == "" {
-		response, err := dockerCli.NetworkCreate(dockerCtx, VPNNetworkName, types.NetworkCreate{
-			Internal: true,
-		})
-
-		if err != nil {
-			return err
-		}
-
-		vpnNetworkID = response.ID
-	}
-
-	return nil
-}
-
-func setupVPNHostNetwork() (error) {
-	setupContext()
-	setupDockerCLI()
-
-	if vpnHostNetworkID == "" {
-		response, err := dockerCli.NetworkCreate(dockerCtx, vpnHostNetworkName, types.NetworkCreate{
-			Internal: false,
-		})
-
-		if err != nil {
-			return err
-		}
-
-		vpnHostNetworkID = response.ID
-	}
-
-	return nil
-}
-
 func getCertificate() (string, error) {
 	if vpnContainerID == "" {
 		return "", errors.New("VPN container not up")
@@ -182,16 +152,22 @@ func getCertificate() (string, error) {
 	// get certificate
 	var certResponse *http.Response
 
+	// retry for 10 seconds to dial to the VPN container
 	for i := 0; i < 10; i++ {
-		certResponse, err = http.Get(fmt.Sprintf("http://%s:9999/", inspectJSON.NetworkSettings.Networks[VPNNetworkName].IPAddress))
-
-		if err == nil {
-			break
+		// Check if the VPN container is already part of our challenge container network
+		if inspectJSON.NetworkSettings.Networks[vpnHostNetworkName] != nil {
+			// it is - get the IP address and dial to it
+			certResponse, err = http.Get(fmt.Sprintf("http://%s:9999/", inspectJSON.NetworkSettings.Networks[vpnHostNetworkName].IPAddress))
+
+			if err == nil {
+				break
+			}
 		}
+
 		time.Sleep(time.Second)
 	}
 
-	if err != nil {
+	if err != nil || certResponse == nil {
 		return "", err
 	}
 
diff --git a/src/container.go b/src/container.go
index c9a918f..29046e2 100644
--- a/src/container.go
+++ b/src/container.go
@@ -9,7 +9,7 @@ import (
 )
 
 const (
-	VPNNetworkName = "circus-vpnnet"
+	containerNetworkName = "circus-vpnnet"
 )
 
 type ChallengeContainer struct {
@@ -23,11 +23,14 @@ func (cc ChallengeContainer) startContainer() (address string, containerID strin
 	// Set up our context and Docker CLI connection
 	setupContext()
 	setupDockerCLI()
-	// Set up network
-	err = setupNetwork()
 
-	if err != nil {
-		return "", "", err
+	// Set up container network
+	if containerNetworkID == "" {
+		id, err := setupNetwork(containerNetworkName, true)
+		if (err != nil) {
+			return "", "", err
+		}
+		containerNetworkID = id
 	}
 
 	// Create container
@@ -37,8 +40,8 @@ func (cc ChallengeContainer) startContainer() (address string, containerID strin
 		Tty: false,
 	}, nil, &network.NetworkingConfig{
 		EndpointsConfig: map[string]*network.EndpointSettings{
-			VPNNetworkName: {
-				NetworkID: vpnNetworkID,
+			containerNetworkName: {
+				NetworkID: containerNetworkID,
 			},
 		},
 	}, "")
@@ -60,7 +63,7 @@ func (cc ChallengeContainer) startContainer() (address string, containerID strin
 	}
 
 	// Return IP, Container ID and error
-	return inspectJSON.NetworkSettings.Networks[VPNNetworkName].IPAddress, resp.ID,nil
+	return inspectJSON.NetworkSettings.Networks[containerNetworkName].IPAddress, resp.ID,nil
 }
 
 // Stops the container with a timeout of one second
diff --git a/src/http.go b/src/http.go
index 91d1c87..d1e3507 100644
--- a/src/http.go
+++ b/src/http.go
@@ -21,7 +21,7 @@ func registerHTTPFlags() {
 	port = flag.Int("port", 8080, "The port for HTTP")
 }
 
-func runHTTPServer() (error) {
+func setupHTTPServer() (http.Server) {
 	r := mux.NewRouter()
 
 	r.HandleFunc("/", indexHandler)
@@ -37,8 +37,10 @@ func runHTTPServer() (error) {
 	r.HandleFunc("/api/stopContainer", stopContainerHandler).Methods("POST")
 	r.HandleFunc("/api/getAccess", getAccessHandler).Methods("GET")
 
-	address := fmt.Sprintf("0.0.0.0:%d", *port)
-	return http.ListenAndServe(address, r)
+	return http.Server{
+		Addr: fmt.Sprintf("0.0.0.0:%d", *port),
+		Handler: r,
+	}
 }
 
 // Host the index file
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)
+}
diff --git a/src/network.go b/src/network.go
new file mode 100644
index 0000000..93abaca
--- /dev/null
+++ b/src/network.go
@@ -0,0 +1,82 @@
+package main
+
+import (
+	"github.com/docker/docker/api/types"
+	"time"
+)
+
+
+const(
+	vpnHostNetworkName = "vpnhostnet"
+)
+
+var (
+	containerNetworkID string
+	vpnHostNetworkID string
+)
+
+func setupNetwork(networkName string, internalNetwork bool) (string, error) {
+	setupContext()
+	setupDockerCLI()
+
+	// Cleanup old network(s) with that name
+	err := deleteNetwork(networkName)
+	if err != nil {
+		return "", err
+	}
+
+	// create our new network
+	response, err := dockerCli.NetworkCreate(dockerCtx, networkName, types.NetworkCreate{
+		Internal: internalNetwork,
+	})
+
+	if err != nil {
+		return "", err
+	}
+
+	return response.ID, nil
+}
+
+func deleteNetwork(networkName string) (error) {
+	// try to cleanup old instances of that network
+	// BUT we need to get the ID used on the old network, and we only have the name of the network
+	// so we get a list of all networks and skip through that array...:
+	networks, err := dockerCli.NetworkList(dockerCtx, types.NetworkListOptions{})
+	// TODO: we sure as hell can filter that with NetworkListOptions!
+	if err != nil {
+		return err
+	}
+
+	// Iterate over all networks
+	for _, network := range networks {
+		if network.Name == networkName {
+			// Found a network with that name
+			// We need to stop all containers attached to it before we can remove the network
+
+			// Get all containers attached to the network
+			resource, err := dockerCli.NetworkInspect(dockerCtx, network.ID, types.NetworkInspectOptions{})
+			if err != nil {
+				return err
+			}
+
+			// Loop over containers attached to that network and stop them
+			timeout := time.Second
+			for _, container := range resource.Containers {
+				// Stop container
+				dockerCli.ContainerStop(dockerCtx, container.EndpointID, &timeout)
+				time.Sleep(timeout)
+				// We ignore the err return variable here.
+				// It may happen that this container is already stopped, but is still part of the network
+				// TODO: Inspect the container to check that - we don't want to ignore errors!
+			}
+
+			// Now we can remove the network
+			err = dockerCli.NetworkRemove(dockerCtx, network.ID)
+			if err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
\ No newline at end of file