From 10749de40f692faac3b890551c2477af323f8dae Mon Sep 17 00:00:00 2001
From: fwkz <f4wkes@gmail.com>
Date: Thu, 19 Nov 2020 11:09:06 +0100
Subject: [PATCH] Make `proxy.servers` a map in order to easily pick proxy to
 close

---
 main.go        |  4 ++--
 proxy/serve.go | 49 ++++++++++++++++++-------------------------------
 2 files changed, 20 insertions(+), 33 deletions(-)

diff --git a/main.go b/main.go
index 35729f2e9..73155fd81 100644
--- a/main.go
+++ b/main.go
@@ -397,8 +397,8 @@ func startServers(cfg *config.Config) {
 						ports = unique(ports)
 					}
 					for _, port := range difference(lastPorts, ports) {
-						log.Printf("[DEBUG] Dynamic TCP listener on port %s eligable for termination", port)
-						proxy.CloseDynamicProxy(port)
+						log.Printf("[DEBUG] Dynamic TCP listener on %s eligable for termination", port)
+						proxy.CloseProxy(port)
 					}
 					for _, port := range ports {
 						l := l
diff --git a/proxy/serve.go b/proxy/serve.go
index 41d9f5f5f..bb8209f86 100644
--- a/proxy/serve.go
+++ b/proxy/serve.go
@@ -28,23 +28,21 @@ type Server interface {
 var (
 	// mu guards servers which contains the list
 	// of running proxy servers.
-	mu                sync.Mutex
-	servers           []Server
-	dynamicProxyPorts = make(map[string]Server)
+	mu      sync.Mutex
+	servers = make(map[string]Server)
 )
 
-func CloseDynamicProxy(port string) error {
-	if srv, ok := dynamicProxyPorts[port]; ok {
+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 port %s has been terminated", port)
-		mu.Lock()
-		servers = removeServer(servers, srv)
-		delete(dynamicProxyPorts, port)
-		mu.Unlock()
+		log.Printf("[INFO] Dynamic TCP listener on %s has been terminated", address)
+		delete(servers, address)
 	}
+	mu.Unlock()
 	return nil
 }
 
@@ -53,15 +51,17 @@ func Close() {
 	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
@@ -146,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 {
@@ -174,15 +175,12 @@ func ListenAndServeTCP(l config.Listen, h tcp.Handler, cfg *tls.Config) error {
 		ReadTimeout:  l.ReadTimeout,
 		WriteTimeout: l.WriteTimeout,
 	}
-	if _, ok := h.(*tcp.DynamicProxy); ok {
-		dynamicProxyPorts[l.Addr] = srv
-	}
 	return serve(ln, srv)
 }
 
 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 {
@@ -197,14 +195,3 @@ func serve(ln net.Listener, srv Server) error {
 	}
 	return err
 }
-
-func removeServer(srvs []Server, s Server) []Server {
-	for i, srv := range srvs {
-		if s == srv {
-			log.Printf("[DEBUG] Removing server from the list")
-			srvs[i] = srvs[len(servers)-1]
-			return srvs[:len(srvs)-1]
-		}
-	}
-	return srvs
-}