Commit effc6675 authored by Cecylia Bocovich's avatar Cecylia Bocovich
Browse files

Wait until all goroutines finish before shutdown

parent b9cc54b3
Pipeline #1848 canceled with stage
...@@ -12,6 +12,7 @@ import ( ...@@ -12,6 +12,7 @@ import (
"os/signal" "os/signal"
"path/filepath" "path/filepath"
"strings" "strings"
"sync"
"syscall" "syscall"
"time" "time"
...@@ -27,7 +28,7 @@ const ( ...@@ -27,7 +28,7 @@ const (
) )
// Accept local SOCKS connections and pass them to the handler. // 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() defer ln.Close()
for { for {
conn, err := ln.AcceptSocks() conn, err := ln.AcceptSocks()
...@@ -40,6 +41,8 @@ func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan struc ...@@ -40,6 +41,8 @@ func socksAcceptLoop(ln *pt.SocksListener, tongue sf.Tongue, shutdown chan struc
} }
log.Printf("SOCKS accepted: %v", conn.Req) log.Printf("SOCKS accepted: %v", conn.Req)
go func() { go func() {
wg.Add(1)
defer wg.Done()
defer conn.Close() defer conn.Close()
err := conn.Grant(&net.TCPAddr{IP: net.IPv4zero, Port: 0}) err := conn.Grant(&net.TCPAddr{IP: net.IPv4zero, Port: 0})
...@@ -173,6 +176,7 @@ func main() { ...@@ -173,6 +176,7 @@ func main() {
} }
listeners := make([]net.Listener, 0) listeners := make([]net.Listener, 0)
shutdown := make(chan struct{}) shutdown := make(chan struct{})
var wg sync.WaitGroup
for _, methodName := range ptInfo.MethodNames { for _, methodName := range ptInfo.MethodNames {
switch methodName { switch methodName {
case "snowflake": case "snowflake":
...@@ -183,7 +187,7 @@ func main() { ...@@ -183,7 +187,7 @@ func main() {
break break
} }
log.Printf("Started SOCKS listener at %v.", ln.Addr()) 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()) pt.Cmethod(methodName, ln.Version(), ln.Addr())
listeners = append(listeners, ln) listeners = append(listeners, ln)
default: default:
...@@ -216,6 +220,7 @@ func main() { ...@@ -216,6 +220,7 @@ func main() {
ln.Close() ln.Close()
} }
close(shutdown) close(shutdown)
wg.Wait()
log.Println("snowflake is done.") log.Println("snowflake is done.")
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment