Add graceful shutdown to psiphonClient; move log file initialization out of the "library" package as log.SetOutput has global effect

psiphon/controller.go

@@ -27,9 +27,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"log"
 	"net"
-	"os"
 	"sync"
 	"time"
 )
@@ -85,6 +83,9 @@ func NewController(config *Config) (controller *Controller) {
 // - a local SOCKS proxy that port forwards through the pool of tunnels
 // - a local HTTP proxy that port forwards through the pool of tunnels
 func (controller *Controller) Run(shutdownBroadcast <-chan struct{}) {
+
+	Notice(NOTICE_VERSION, VERSION)
+
 	socksProxy, err := NewSocksProxy(controller.config, controller)
 	if err != nil {
 		Notice(NOTICE_ALERT, "error initializing local SOCKS proxy: %s", err)
@@ -609,23 +610,3 @@ func (controller *Controller) establishTunnelWorker() {
 	}
 	Notice(NOTICE_INFO, "stopped establish worker")
 }
-
-// RunForever executes the main loop of the Psiphon client. It launches
-// the controller with a shutdown that is never signaled.
-func RunForever(config *Config) {
-
-	if config.LogFilename != "" {
-		logFile, err := os.OpenFile(config.LogFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
-		if err != nil {
-			Fatal("error opening log file: %s", err)
-		}
-		defer logFile.Close()
-		log.SetOutput(logFile)
-	}
-
-	Notice(NOTICE_VERSION, VERSION)
-
-	controller := NewController(config)
-	shutdownBroadcast := make(chan struct{})
-	controller.Run(shutdownBroadcast)
-}

psiphonClient.go

@@ -23,9 +23,13 @@ import (
 	"flag"
 	psiphon "github.com/Psiphon-Labs/psiphon-tunnel-core/psiphon"
 	"log"
+	"os"
+	"os/signal"
+	"sync"
 )
 
 func main() {
+
 	var configFilename string
 	flag.StringVar(&configFilename, "config", "", "configuration file")
 	flag.Parse()
@@ -36,5 +40,30 @@ func main() {
 	if err != nil {
 		log.Fatalf("error loading configuration file: %s", err)
 	}
-	psiphon.RunForever(config)
+
+	if config.LogFilename != "" {
+		logFile, err := os.OpenFile(config.LogFilename, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
+		if err != nil {
+			log.Fatalf("error opening log file: %s", err)
+		}
+		defer logFile.Close()
+		log.SetOutput(logFile)
+	}
+
+	controller := psiphon.NewController(config)
+	shutdownBroadcast := make(chan struct{})
+	controllerWaitGroup := new(sync.WaitGroup)
+	controllerWaitGroup.Add(1)
+	go func() {
+		defer controllerWaitGroup.Done()
+		controller.Run(shutdownBroadcast)
+	}()
+
+	systemStopSignal := make(chan os.Signal, 1)
+	signal.Notify(systemStopSignal, os.Interrupt, os.Kill)
+	<-systemStopSignal
+
+	psiphon.Notice(psiphon.NOTICE_INFO, "shutdown by system")
+	close(shutdownBroadcast)
+	controllerWaitGroup.Wait()
 }