diff --git a/main.go b/main.go index e37fb9b30..73155fd81 100644 --- a/main.go +++ b/main.go @@ -378,6 +378,7 @@ func startServers(cfg *config.Config) { case "tcp-dynamic": go func() { var buffer strings.Builder + lastPorts := []string{} for { time.Sleep(l.Refresh) table := route.GetTable() @@ -395,6 +396,10 @@ func startServers(cfg *config.Config) { } ports = unique(ports) } + for _, port := range difference(lastPorts, ports) { + log.Printf("[DEBUG] Dynamic TCP listener on %s eligable for termination", port) + proxy.CloseProxy(port) + } for _, port := range ports { l := l port := port @@ -419,6 +424,7 @@ func startServers(cfg *config.Config) { } }() } + lastPorts = ports } }() case "https+tcp+sni": @@ -655,6 +661,21 @@ func unique(strSlice []string) []string { return list } +// difference returns elements in `a` that aren't in `b` +func difference(a, b []string) []string { + mb := make(map[string]struct{}, len(b)) + for _, x := range b { + mb[x] = struct{}{} + } + var diff []string + for _, x := range a { + if _, found := mb[x]; !found { + diff = append(diff, x) + } + } + return diff +} + func tableSchemes(r route.Routes) []string { schemes := []string{} for _, rt := range r { diff --git a/proxy/serve.go b/proxy/serve.go index 54a98a73b..bb8209f86 100644 --- a/proxy/serve.go +++ b/proxy/serve.go @@ -4,6 +4,7 @@ import ( "context" "crypto/tls" "errors" + "log" "net" "net/http" "sync" @@ -28,23 +29,39 @@ var ( // mu guards servers which contains the list // of running proxy servers. mu sync.Mutex - servers []Server + servers = make(map[string]Server) ) +func CloseProxy(address string) error { + mu.Lock() + if srv, ok := servers[address]; ok { + err := srv.Close() + if err != nil { + return err + } + log.Printf("[INFO] Dynamic TCP listener on %s has been terminated", address) + delete(servers, address) + } + mu.Unlock() + return nil +} + func Close() { mu.Lock() for _, srv := range servers { srv.Close() } - servers = []Server{} + servers = make(map[string]Server) mu.Unlock() } func Shutdown(timeout time.Duration) { mu.Lock() - srvs := make([]Server, len(servers)) - copy(srvs, servers) - servers = []Server{} + srvs := make(map[string]Server, len(servers)) + for k, v := range servers { + srvs[k] = v + } + servers = make(map[string]Server) mu.Unlock() var wg sync.WaitGroup @@ -129,8 +146,9 @@ func ListenAndServeHTTPSTCPSNI(l config.Listen, h http.Handler, p tcp.Handler, c }) // tcpproxy creates its own listener from the configuration above so we can - // safely pass nil here. - return serve(nil, tps) + // safely pass nil here, nonetheless we are passing `httpsListener` to + // extract it's address and save server in the `servers` map. + return serve(httpsListener, tps) } func ListenAndServeGRPC(l config.Listen, opts []grpc.ServerOption, cfg *tls.Config) error { @@ -162,7 +180,7 @@ func ListenAndServeTCP(l config.Listen, h tcp.Handler, cfg *tls.Config) error { func serve(ln net.Listener, srv Server) error { mu.Lock() - servers = append(servers, srv) + servers[ln.Addr().String()] = srv mu.Unlock() err := srv.Serve(ln) if err != nil {