Clean up/refactor the shutdown/termination handling code.

This combines the old signal processing code with the parent monitor,
into a new termination monitor structure, which also now handles keeping
track of outstanding sessions.

3 files changed, 84 insertions, 83 deletions

diff --git a/obfs4proxy/obfs4proxy.go b/obfs4proxy/obfs4proxy.go
index b27d75d..9b452ac 100644
--- a/obfs4proxy/obfs4proxy.go
+++ b/obfs4proxy/obfs4proxy.go
@@ -38,7 +38,6 @@ import (
 	"net"
 	"net/url"
 	"os"
-	"os/signal"
 	"path"
 	"sync"
 	"syscall"
@@ -60,7 +59,7 @@ const (
 var enableLogging bool
 var unsafeLogging bool
 var stateDir string
-var handlerChan chan int
+var termMon *termMonitor
 
 // DialFn is a function pointer to a function that matches the net.Dialer.Dial
 // interface.
@@ -176,10 +175,8 @@ func clientAcceptLoop(f base.ClientFactory, ln *pt.SocksListener, proxyURI *url.
 
 func clientHandler(f base.ClientFactory, conn *pt.SocksConn, proxyURI *url.URL) {
 	defer conn.Close()
-	handlerChan <- 1
-	defer func() {
-		handlerChan <- -1
-	}()
+	termMon.onHandlerStart()
+	defer termMon.onHandlerFinish()
 
 	name := f.Transport().Name()
 	addrStr := elideAddr(conn.Req.Target)
@@ -298,10 +295,8 @@ func serverAcceptLoop(f base.ServerFactory, ln net.Listener, info *pt.ServerInfo
 
 func serverHandler(f base.ServerFactory, conn net.Conn, info *pt.ServerInfo) {
 	defer conn.Close()
-	handlerChan <- 1
-	defer func() {
-		handlerChan <- -1
-	}()
+	termMon.onHandlerStart()
+	defer termMon.onHandlerFinish()
 
 	name := f.Transport().Name()
 	addrStr := elideAddr(conn.RemoteAddr().String())
@@ -386,8 +381,8 @@ func getVersion() string {
 }
 
 func main() {
-	// Initialize parent process monitoring as early as possible.
-	pmonErr := initParentMonitor()
+	// Initialize the termination state monitor as soon as possible.
+	termMon = newTermMonitor()
 
 	// Handle the command line arguments.
 	_, execName := path.Split(os.Args[0])
@@ -405,10 +400,8 @@ func main() {
 		log.Fatalf("[ERROR]: failed to set log level: %s", err)
 	}
 
-	// Determine if this is a client or server, initialize logging, and finish
-	// the pt configuration.
+	// Determine if this is a client or server, initialize the common state.
 	var ptListeners []net.Listener
-	handlerChan = make(chan int)
 	launched := false
 	isClient, err := ptIsClient()
 	if err != nil {
@@ -419,12 +412,10 @@ func main() {
 	}
 	if err = ptInitializeLogging(enableLogging); err != nil {
 		log.Fatalf("[ERROR]: %s - failed to initialize logging", execName)
-	} else {
-		noticef("%s - launched", getVersion())
-		if pmonErr != nil {
-			warnf("%s - failed to initialize parent monitor: %s", execName, pmonErr)
-		}
 	}
+	noticef("%s - launched", getVersion())
+
+	// Do the managed pluggable transport protocol configuration.
 	if isClient {
 		infof("%s - initializing client transport listeners", execName)
 		launched, ptListeners = clientSetup()
@@ -444,39 +435,18 @@ func main() {
 	}()
 
 	// At this point, the pt config protocol is finished, and incoming
-	// connections will be processed.  Per the pt spec, on sane platforms
-	// termination is signaled via SIGINT (or SIGTERM), so wait on tor to
-	// request a shutdown of some sort.
-
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-
-	// Wait for the first SIGINT (close listeners).
-	var sig os.Signal
-	numHandlers := 0
-	for sig == nil {
-		select {
-		case n := <-handlerChan:
-			numHandlers += n
-		case sig = <-sigChan:
-			if sig == syscall.SIGTERM {
-				// SIGTERM causes immediate termination.
-				return
-			}
-		}
+	// connections will be processed.  Wait till the parent dies
+	// (immediate exit), a SIGTERM is received (immediate exit),
+	// or a SIGINT is received.
+	if sig := termMon.wait(false); sig == syscall.SIGTERM {
+		return
 	}
+
+	// Ok, it was the first SIGINT, close all listeners, and wait till,
+	// the parent dies, all the current connections are closed, or either
+	// a SIGINT/SIGTERM is received, and exit.
 	for _, ln := range ptListeners {
 		ln.Close()
 	}
-
-	// Wait for the 2nd SIGINT (or a SIGTERM), or for all current sessions to
-	// finish.
-	sig = nil
-	for sig == nil && numHandlers != 0 {
-		select {
-		case n := <-handlerChan:
-			numHandlers += n
-		case sig = <-sigChan:
-		}
-	}
+	termMon.wait(true)
 }
diff --git a/obfs4proxy/parentMonitor.go b/obfs4proxy/termmon.go
index e2f078d..eac7e20 100644
--- a/obfs4proxy/parentMonitor.go
+++ b/obfs4proxy/termmon.go
@@ -28,36 +28,47 @@
 package main
 
 import (
-	"fmt"
 	"os"
+	"os/signal"
 	"runtime"
 	"syscall"
 	"time"
 )
 
-var parentMonitorOSInit func() error
+var termMonitorOSInit func(*termMonitor) error
 
-func initParentMonitor() error {
-	// Until #15435 is implemented, there is no reliable way to see if
-	// the parent has died that is portable/platform independent/reliable.
-	//
-	// Do the next best thing and use various kludges and hacks:
-	//  * Linux - Platform specific code that should always work.
-	//  * Other U*IX - Somewhat generic code, that works unless the parent
-	//    dies before the monitor is initialized.
-	//  * Windows - Log an error, can't be bothered to figure out how
-	//    to handle this there.
-	if parentMonitorOSInit != nil {
-		return parentMonitorOSInit()
-	} else if runtime.GOOS != "windows" {
-		ppid := os.Getppid()
-		go parentMonitorPpidChange(ppid)
-		return nil
+type termMonitor struct {
+	sigChan     chan os.Signal
+	handlerChan chan int
+	numHandlers int
+}
+
+func (m *termMonitor) onHandlerStart() {
+	m.handlerChan <- 1
+}
+
+func (m *termMonitor) onHandlerFinish() {
+	m.handlerChan <- -1
+}
+
+func (m *termMonitor) wait(termOnNoHandlers bool) os.Signal {
+	// Block until a signal has been received, or (optionally) the
+	// number of pending handlers has hit 0.  In the case of the
+	// latter, treat it as if a SIGTERM has been received.
+	for {
+		select {
+		case n := <-m.handlerChan:
+			m.numHandlers += n
+		case sig := <-m.sigChan:
+			return sig
+		}
+		if termOnNoHandlers && m.numHandlers == 0 {
+			return syscall.SIGTERM
+		}
 	}
-	return fmt.Errorf("unsupported on: %s", runtime.GOOS)
 }
 
-func parentMonitorPpidChange(ppid int) {
+func (m *termMonitor) termOnPPIDChange(ppid int) {
 	// Under most if not all U*IX systems, the parent PID will change
 	// to that of init once the parent dies.  There are several notable
 	// exceptions (Slowlaris/Android), but the parent PID changes
@@ -72,17 +83,37 @@ func parentMonitorPpidChange(ppid int) {
 		time.Sleep(ppidPollInterval)
 	}
 
-	// If possible SIGTERM ourself so that the normal shutdown code
-	// gets invoked.  If any of that fails, exit anyway, we are a
-	// defunt process.
+	// Treat the parent PID changing as the same as having received
+	// a SIGTERM.
 	noticef("Parent pid changed: %d (was %d)", os.Getppid(), ppid)
-	if p, err := os.FindProcess(os.Getpid()); err == nil {
-		if err := p.Signal(syscall.SIGTERM); err == nil {
-			return
+	m.sigChan <- syscall.SIGTERM
+}
+
+func newTermMonitor() *termMonitor {
+	ppid := os.Getppid()
+	m := new(termMonitor)
+	m.sigChan = make(chan os.Signal)
+	m.handlerChan = make(chan int)
+	signal.Notify(m.sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	// Until #15435 is implemented, there is no reliable way to see if
+	// the parent has died that is portable/platform independent/reliable.
+	//
+	// Do the next best thing and use various kludges and hacks:
+	//  * Linux - Platform specific code that should always work.
+	//  * Other U*IX - Somewhat generic code, that works unless the parent
+	//    dies before the monitor is initialized.
+	//  * Windows - Don't specifically monitor for parent termination.
+	if termMonitorOSInit != nil {
+		// Errors here are non-fatal, since it might still be possible
+		// to fall back to a generic implementation.
+		if err := termMonitorOSInit(m); err == nil {
+			return m
 		}
-		warnf("Failed to SIGTERM ourself: %v", err)
-	} else {
-		warnf("Failed to find our own process: %v", err)
 	}
-	os.Exit(-1)
+	if runtime.GOOS != "windows" {
+		go m.termOnPPIDChange(ppid)
+	}
+
+	return m
 }
diff --git a/obfs4proxy/parentMonitor_linux.go b/obfs4proxy/termmon_linux.go
index 65fd307..9711cfc 100644
--- a/obfs4proxy/parentMonitor_linux.go
+++ b/obfs4proxy/termmon_linux.go
@@ -32,7 +32,7 @@ import (
 	"syscall"
 )
 
-func parentMonitorInitLinux() error {
+func termMonitorInitLinux(m *termMonitor) error {
 	// Use prctl() to have the kernel deliver a SIGTERM if the parent
 	// process dies.  This beats anything else that can be done before
 	// #15435 is implemented.
@@ -45,5 +45,5 @@ func parentMonitorInitLinux() error {
 }
 
 func init() {
-	parentMonitorOSInit = parentMonitorInitLinux
+	termMonitorOSInit = termMonitorInitLinux
 }