Merge pull request #106 from getlantern/exitgracefully

Cleaned up exit logic in flashlight.go
getlantern · Mar 1, 2015 · 15d6acf · 15d6acf
2 parents d048dd8 + c15e176
commit 15d6acf
Show file tree

Hide file tree

Showing 7 changed files with 154 additions and 91 deletions.
diff --git a/src/github.com/getlantern/flashlight/config/config.go b/src/github.com/getlantern/flashlight/config/config.go
@@ -63,10 +63,14 @@ type CA struct {
 	Cert       string // PEM-encoded
 }
 
-// Start starts the configuration system.
-func Start(updateHandler func(updated *Config)) (*Config, error) {
+// Init initializes the configuration system.
+func Init() (*Config, error) {
+	configPath, err := InConfigDir("lantern.yaml")
+	if err != nil {
+		return nil, err
+	}
 	m = &yamlconf.Manager{
-		FilePath:         InConfigDir("lantern.yaml"),
+		FilePath:         configPath,
 		FilePollInterval: 1 * time.Second,
 		ConfigServerAddr: *configaddr,
 		EmptyConfig: func() yamlconf.Config {
@@ -104,29 +108,37 @@ func Start(updateHandler func(updated *Config)) (*Config, error) {
 	var cfg *Config
 	if err == nil {
 		cfg = initial.(*Config)
-		updateGlobals(cfg)
-		go func() {
-			// Read updates
-			for {
-				next := m.Next()
-				nextCfg := next.(*Config)
-				updateGlobals(nextCfg)
-				updateHandler(nextCfg)
-			}
-		}()
+		err = updateGlobals(cfg)
+		if err != nil {
+			return nil, err
+		}
 	}
 	return cfg, err
 }
 
-func updateGlobals(cfg *Config) {
+// Run runs the configuration system.
+func Run(updateHandler func(updated *Config)) error {
+	for {
+		next := m.Next()
+		nextCfg := next.(*Config)
+		err := updateGlobals(nextCfg)
+		if err != nil {
+			return err
+		}
+		updateHandler(nextCfg)
+	}
+}
+
+func updateGlobals(cfg *Config) error {
 	globals.InstanceId = cfg.InstanceId
 	loc := &geolookup.City{}
 	loc.Country.IsoCode = cfg.Country
 	globals.SetLocation(loc)
 	err := globals.SetTrustedCAs(cfg.TrustedCACerts())
 	if err != nil {
-		log.Fatalf("Unable to configure trusted CAs: %s", err)
+		return fmt.Errorf("Unable to configure trusted CAs: %s", err)
 	}
+	return nil
 }
 
 // Update updates the configuration using the given mutator function.
@@ -137,7 +149,7 @@ func Update(mutate func(cfg *Config) error) error {
 }
 
 // InConfigDir returns the path to the given filename inside of the configdir.
-func InConfigDir(filename string) string {
+func InConfigDir(filename string) (string, error) {
 	cdir := *configdir
 	if cdir == "" {
 		cdir = appdir.General("Lantern")
@@ -147,11 +159,11 @@ func InConfigDir(filename string) string {
 		if os.IsNotExist(err) {
 			// Create config dir
 			if err := os.MkdirAll(cdir, 0755); err != nil {
-				log.Fatalf("Unable to create configdir at %s: %s", cdir, err)
+				return "", fmt.Errorf("Unable to create configdir at %s: %s", cdir, err)
 			}
 		}
 	}
-	return filepath.Join(cdir, filename)
+	return filepath.Join(cdir, filename), nil
 }
 
 // TrustedCACerts returns a slice of PEM-encoded certs for the trusted CAs

diff --git a/src/github.com/getlantern/flashlight/flashlight.go b/src/github.com/getlantern/flashlight/flashlight.go
@@ -27,12 +27,6 @@ import (
 	"github.com/getlantern/flashlight/ui"
 )
 
-const (
-	// Exit Statuses
-	ConfigError    = 1
-	PortmapFailure = 50
-)
-
 var (
 	version   string
 	buildDate string
@@ -43,6 +37,7 @@ var (
 	parentPID = flag.Int("parentpid", 0, "the parent process's PID, used on Windows for killing flashlight when the parent disappears")
 
 	configUpdates = make(chan *config.Config)
+	exitCh        = make(chan error, 1)
 )
 
 func init() {
@@ -57,40 +52,72 @@ func init() {
 }
 
 func main() {
-	systray.Run(doMain)
+	systray.Run(_main)
 }
 
-func doMain() {
+func _main() {
+	err := doMain()
+	if err != nil {
+		log.Fatal(err)
+	}
+	log.Debug("Lantern stopped")
+	os.Exit(0)
+}
+
+func doMain() error {
+	err := logging.Init()
+	if err != nil {
+		return err
+	}
+
+	// Schedule cleanup actions
 	defer logging.Close()
+	defer pacOff()
+	defer systray.Quit()
+
 	i18nInit()
-	configureSystemTray()
+	err = configureSystemTray()
+	if err != nil {
+		return err
+	}
 	displayVersion()
 
 	flag.Parse()
 	configUpdates = make(chan *config.Config)
-	cfg, err := config.Start(func(updated *config.Config) {
-		configUpdates <- updated
-	})
+	cfg, err := config.Init()
 	if err != nil {
-		log.Fatalf("Unable to start configuration: %s", err)
+		return fmt.Errorf("Unable to initialize configuration: %v", err)
 	}
+	go func() {
+		err := config.Run(func(updated *config.Config) {
+			configUpdates <- updated
+		})
+		if err != nil {
+			exit(err)
+		}
+	}()
 	if *help || cfg.Addr == "" || (cfg.Role != "server" && cfg.Role != "client") {
 		flag.Usage()
-		os.Exit(ConfigError)
+		return fmt.Errorf("Wrong arguments")
 	}
 
 	finishProfiling := profiling.Start(cfg.CpuProfile, cfg.MemProfile)
 	defer finishProfiling()
 
 	// Configure stats initially
-	configureStats(cfg, true)
+	err = statreporter.Configure(cfg.Stats)
+	if err != nil {
+		return err
+	}
 
 	log.Debugf("Running proxy")
 	if cfg.IsDownstream() {
 		runClientProxy(cfg)
 	} else {
 		runServerProxy(cfg)
 	}
+
+	return waitForExit()
 }
 
 func i18nInit() {
@@ -99,32 +126,19 @@ func i18nInit() {
 	})
 	err := i18n.UseOSLocale()
 	if err != nil {
-		panic(err)
+		exit(err)
 	}
 }
 
 func displayVersion() {
 	log.Debugf("---- flashlight version %s (%s) ----", version, buildDate)
 }
 
-func configureStats(cfg *config.Config, failOnError bool) {
-	var err error
-
-	// Configuring statreporter
-	err = statreporter.Configure(cfg.Stats)
-	if err != nil {
-		log.Error(err)
-		if failOnError {
-			flag.Usage()
-			os.Exit(ConfigError)
-		}
-	}
-}
-
 // Runs the client-side proxy
 func runClientProxy(cfg *config.Config) {
-	if !setUpPacTool() {
-		exit()
+	err := setUpPacTool()
+	if err != nil {
+		exit(err)
 	}
 	client := &client.Client{
 		Addr:         cfg.Addr,
@@ -137,7 +151,8 @@ func runClientProxy(cfg *config.Config) {
 	if cfg.UIAddr != "" {
 		err := ui.Start(cfg.UIAddr)
 		if err != nil {
-			panic(fmt.Errorf("Unable to start UI: %v", err))
+			exit(fmt.Errorf("Unable to start UI: %v", err))
+			return
 		}
 		ui.Show()
 	}
@@ -160,7 +175,8 @@ func runClientProxy(cfg *config.Config) {
 			cfg := <-configUpdates
 
 			proxiedsites.Configure(cfg.ProxiedSites)
-			configureStats(cfg, false)
+			// Note - we deliberately ignore the error from statreporter.Configure here
+			statreporter.Configure(cfg.Stats)
 			hqfd = client.Configure(cfg.Client)
 			if hqfd != nil {
 				hqfdc := hqfd.DirectHttpClient()
@@ -171,24 +187,33 @@ func runClientProxy(cfg *config.Config) {
 		}
 	}()
 
-	err := client.ListenAndServe(pacOn)
-	if err != nil {
-		log.Fatalf("Unable to run client proxy: %s", err)
-	}
+	go func() {
+		exit(client.ListenAndServe(pacOn))
+	}()
+	log.Debug("Ran goroutine")
 }
 
 // Runs the server-side proxy
 func runServerProxy(cfg *config.Config) {
 	useAllCores()
 
+	pkFile, err := config.InConfigDir("proxypk.pem")
+	if err != nil {
+		log.Fatal(err)
+	}
+	certFile, err := config.InConfigDir("servercert.pem")
+	if err != nil {
+		log.Fatal(err)
+	}
+
 	srv := &server.Server{
 		Addr:         cfg.Addr,
 		Host:         cfg.Server.AdvertisedHost,
 		ReadTimeout:  0, // don't timeout
 		WriteTimeout: 0,
 		CertContext: &fronted.CertContext{
-			PKFile:         config.InConfigDir("proxypk.pem"),
-			ServerCertFile: config.InConfigDir("servercert.pem"),
+			PKFile:         pkFile,
+			ServerCertFile: certFile,
 		},
 		AllowedPorts: []int{80, 443, 8080, 8443, 5222},
 	}
@@ -199,12 +224,12 @@ func runServerProxy(cfg *config.Config) {
 	go func() {
 		for {
 			cfg := <-configUpdates
-			configureStats(cfg, false)
+			statreporter.Configure(cfg.Stats)
 			srv.Configure(cfg.Server)
 		}
 	}()
 
-	err := srv.ListenAndServe()
+	err = srv.ListenAndServe()
 	if err != nil {
 		log.Fatalf("Unable to run server proxy: %s", err)
 	}
@@ -216,10 +241,10 @@ func useAllCores() {
 	runtime.GOMAXPROCS(numcores)
 }
 
-func configureSystemTray() {
+func configureSystemTray() error {
 	icon, err := Asset("icons/16on.ico")
 	if err != nil {
-		log.Fatalf("Unable to load icon for system tray: %v", err)
+		return fmt.Errorf("Unable to load icon for system tray: %v", err)
 	}
 	systray.SetIcon(icon)
 	systray.SetTooltip("Lantern")
@@ -231,14 +256,22 @@ func configureSystemTray() {
 			case <-show.ClickedCh:
 				ui.Show()
 			case <-quit.ClickedCh:
-				pacOff()
-				exit()
+				exit(nil)
+				return
 			}
 		}
 	}()
+
+	return nil
 }
 
-func exit() {
-	systray.Quit()
-	os.Exit(0)
+// exit tells the application to exit, optionally supplying an error that caused
+// the exit.
+func exit(err error) {
+	exitCh <- err
+}
+
+// WaitForExit waits for a request to exit the application.
+func waitForExit() error {
+	return <-exitCh
 }
diff --git a/src/github.com/getlantern/flashlight/logging/logging.go b/src/github.com/getlantern/flashlight/logging/logging.go
@@ -43,14 +43,14 @@ var (
 	lastAddr string
 )
 
-func init() {
+func Init() error {
 	logdir := appdir.Logs("Lantern")
 	log.Debugf("Placing logs in %v", logdir)
 	if _, err := os.Stat(logdir); err != nil {
 		if os.IsNotExist(err) {
 			// Create log dir
 			if err := os.MkdirAll(logdir, 0755); err != nil {
-				log.Fatalf("Unable to create logdir at %s: %s", logdir, err)
+				return fmt.Errorf("Unable to create logdir at %s: %s", logdir, err)
 			}
 		}
 	}
@@ -65,6 +65,8 @@ func init() {
 	errorOut = timestamped(NonStopWriter(os.Stderr, logFile))
 	debugOut = timestamped(NonStopWriter(os.Stdout, logFile))
 	golog.SetOutputs(errorOut, debugOut)
+
+	return nil
 }
 
 func Configure(cfg *config.Config, version string, buildDate string) {
@@ -100,6 +102,7 @@ func Configure(cfg *config.Config, version string, buildDate string) {
 }
 
 func Close() error {
+	golog.ResetOutputs()
 	return logFile.Close()
 }