[tor-commits] [flashproxy/master] Do the proper thing on first and second SIGINT.

[dcf at torproject.org]

commit c3c5c05d581291df9567d8d85c861155193be113
Author: David Fifield <david at bamsoftware.com>
Date:   Sun Nov 11 19:28:09 2012 -0800

    Do the proper thing on first and second SIGINT.
---
 websocket-transport/websocket-client.go |   39 ++++++++++++++++++++++++++++++-
 1 files changed, 38 insertions(+), 1 deletions(-)

diff --git a/websocket-transport/websocket-client.go b/websocket-transport/websocket-client.go
index ca551dd..6cce020 100644
--- a/websocket-transport/websocket-client.go
+++ b/websocket-transport/websocket-client.go
@@ -14,6 +14,10 @@ import (
 const socksTimeout = 2
 const bufSiz = 1500
 
+// When a connection handler starts, +1 is written to this channel; when it
+// ends, -1 is written.
+var handlerChan = make(chan int)
+
 func logDebug(format string, v ...interface{}) {
 	fmt.Fprintf(os.Stderr, format+"\n", v...)
 }
@@ -94,6 +98,11 @@ func proxy(local *net.TCPConn, ws *websocket.Conn) {
 func handleConnection(conn *net.TCPConn) error {
 	defer conn.Close()
 
+	handlerChan <- 1
+	defer func() {
+		handlerChan <- -1
+	}()
+
 	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
 	dest, err := readSocks4aConnect(conn)
 	if err != nil {
@@ -167,16 +176,44 @@ func main() {
 
 	ptClientSetup([]string{ptMethodName})
 
+	listeners := make([]*net.TCPListener, 0)
 	for _, socksAddrStr := range socksAddrStrs {
 		ln, err := startListener(socksAddrStr)
 		if err != nil {
 			ptCmethodError(ptMethodName, err.Error())
 		}
 		ptCmethod(ptMethodName, "socks4", ln.Addr())
+		listeners = append(listeners, ln)
 	}
 	ptCmethodsDone()
 
+	var numHandlers int = 0
+
 	signalChan := make(chan os.Signal, 1)
 	signal.Notify(signalChan, os.Interrupt)
-	<-signalChan
+	var sigint bool = false
+	for !sigint {
+		select {
+		case n := <-handlerChan:
+			numHandlers += n
+		case <-signalChan:
+			logDebug("SIGINT")
+			sigint = true
+		}
+	}
+
+	for _, ln := range listeners {
+		ln.Close()
+	}
+
+	sigint = false
+	for numHandlers != 0 && !sigint {
+		select {
+		case n := <-handlerChan:
+			numHandlers += n
+		case <-signalChan:
+			logDebug("SIGINT")
+			sigint = true
+		}
+	}
 }