Switch console client and mobile libraries to use context

ConsoleClient/main.go

@@ -21,6 +21,7 @@ package main
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"flag"
 	"fmt"
@@ -271,26 +272,31 @@ func main() {
 		os.Exit(1)
 	}
 
-	controllerStopSignal := make(chan struct{}, 1)
-	shutdownBroadcast := make(chan struct{})
+	controllerCtx, stopController := context.WithCancel(context.Background())
+	defer stopController()
+
 	controllerWaitGroup := new(sync.WaitGroup)
 	controllerWaitGroup.Add(1)
 	go func() {
 		defer controllerWaitGroup.Done()
-		controller.Run(shutdownBroadcast)
-		controllerStopSignal <- *new(struct{})
-	}()
+		controller.Run(controllerCtx)
 
-	// Wait for an OS signal or a Run stop signal, then stop Psiphon and exit
+		// Signal the <-controllerCtx.Done() case below. If the <-systemStopSignal
+		// case already called stopController, this is a noop.
+		stopController()
+	}()
 
 	systemStopSignal := make(chan os.Signal, 1)
 	signal.Notify(systemStopSignal, os.Interrupt, os.Kill)
+
+	// Wait for an OS signal or a Run stop signal, then stop Psiphon and exit
+
 	select {
 	case <-systemStopSignal:
 		psiphon.NoticeInfo("shutdown by system")
-		close(shutdownBroadcast)
+		stopController()
 		controllerWaitGroup.Wait()
-	case <-controllerStopSignal:
+	case <-controllerCtx.Done():
 		psiphon.NoticeInfo("shutdown by controller")
 	}
 }

MobileLibrary/psi/psi.go

@@ -25,6 +25,7 @@ package psi
 // Start/Stop interface on top of a single Controller instance.
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"sync"
@@ -64,7 +65,8 @@ func NoticeUserLog(message string) {
 
 var controllerMutex sync.Mutex
 var controller *psiphon.Controller
-var shutdownBroadcast chan struct{}
+var controllerCtx context.Context
+var stopController context.CancelFunc
 var controllerWaitGroup *sync.WaitGroup
 
 func Start(
@@ -128,12 +130,13 @@ func Start(
 		return fmt.Errorf("error initializing controller: %s", err)
 	}
 
-	shutdownBroadcast = make(chan struct{})
+	controllerCtx, stopController = context.WithCancel(context.Background())
+
 	controllerWaitGroup = new(sync.WaitGroup)
 	controllerWaitGroup.Add(1)
 	go func() {
 		defer controllerWaitGroup.Done()
-		controller.Run(shutdownBroadcast)
+		controller.Run(controllerCtx)
 	}()
 
 	return nil
@@ -145,10 +148,11 @@ func Stop() {
 	defer controllerMutex.Unlock()
 
 	if controller != nil {
-		close(shutdownBroadcast)
+		stopController()
 		controllerWaitGroup.Wait()
 		controller = nil
-		shutdownBroadcast = nil
+		controllerCtx = nil
+		stopController = nil
 		controllerWaitGroup = nil
 	}
 }