commit effc6675448a3a6e62d2784557d67e2c46e376d5 Author: Cecylia Bocovich cohosh@torproject.org Date: Fri Dec 4 10:50:00 2020 -0500
Wait until all goroutines finish before shutdown --- client/snowflake.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-)
diff --git a/client/snowflake.go b/client/snowflake.go index a1a679e..55addc1 100644 --- a/client/snowflake.go +++ b/client/snowflake.go @@ -12,6 +12,7 @@ import ( "os/signal" "path/filepath" "strings" + "sync" "syscall" "time"
@@ -27,7 +28,7 @@ const ( )
// Accept local SOCKS connections and pass them to the handler. -func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan struct{}) { +func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan struct{}, wg sync.WaitGroup) { defer ln.Close() for { conn, err := ln.AcceptSocks() @@ -40,6 +41,8 @@ func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan struc } log.Printf("SOCKS accepted: %v", conn.Req) go func() { + wg.Add(1) + defer wg.Done() defer conn.Close()
err := conn.Grant(&net.TCPAddr{IP: net.IPv4zero, Port: 0}) @@ -173,6 +176,7 @@ func main() { } listeners := make([]net.Listener, 0) shutdown := make(chan struct{}) + var wg sync.WaitGroup for _, methodName := range ptInfo.MethodNames { switch methodName { case "snowflake": @@ -183,7 +187,7 @@ func main() { break } log.Printf("Started SOCKS listener at %v.", ln.Addr()) - go socksAcceptLoop(ln, dialer, shutdown) + go socksAcceptLoop(ln, dialer, shutdown, wg) pt.Cmethod(methodName, ln.Version(), ln.Addr()) listeners = append(listeners, ln) default: @@ -216,6 +220,7 @@ func main() { ln.Close() } close(shutdown) + wg.Wait() log.Println("snowflake is done.") }