From fb756e03de5ff1411664d9dbcd65bc093a060127 Mon Sep 17 00:00:00 2001
From: haidy <haidy-cao@qq.com>
Date: Thu, 7 Mar 2019 21:45:38 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8D=E5=86=85=E5=AD=98=E6=B3=84?=
 =?UTF-8?q?=E6=BC=8F=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 client/control.go      |  6 +++++-
 client/service.go      |  3 ++-
 server/ports/ports.go  | 29 +++++++++++++++++++++++------
 server/service.go      | 14 ++++++++++++--
 utils/net/websocket.go |  8 ++++++++
 5 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/client/control.go b/client/control.go
index d304e367..ba687f94 100644
--- a/client/control.go
+++ b/client/control.go
@@ -145,7 +145,11 @@ func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
 
 func (ctl *Control) Close() error {
 	if ctl.conn != nil {
-		ctl.conn.Close()
+		_ = ctl.conn.Close()
+		ctl.conn = nil
+		ctl.msgHandlerShutdown = nil
+		ctl.readerShutdown = nil
+		ctl.session = nil
 		log.Info("conn closed")
 	}
 	if ctl.pm != nil {
diff --git a/client/service.go b/client/service.go
index 7221eb86..19bb90cd 100644
--- a/client/service.go
+++ b/client/service.go
@@ -239,9 +239,10 @@ func (svr *Service) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs
 func (svr *Service) Close() {
 	atomic.StoreUint32(&svr.exit, 1)
 	if svr.ctl != nil {
-		svr.ctl.Close()
+		_ = svr.ctl.Close()
 	}
 	svr.closedCh <- true
+	svr.ctl = nil
 }
 
 func (svr *Service) IsClosed() bool {
diff --git a/server/ports/ports.go b/server/ports/ports.go
index 3544e873..324b81da 100644
--- a/server/ports/ports.go
+++ b/server/ports/ports.go
@@ -38,6 +38,7 @@ type PortManager struct {
 	bindAddr string
 	netType  string
 	mu       sync.Mutex
+	stop     chan bool
 }
 
 func NewPortManager(netType string, bindAddr string, allowPorts map[int]struct{}) *PortManager {
@@ -47,6 +48,7 @@ func NewPortManager(netType string, bindAddr string, allowPorts map[int]struct{}
 		freePorts:     make(map[int]struct{}),
 		bindAddr:      bindAddr,
 		netType:       netType,
+		stop:          make(chan bool),
 	}
 	if len(allowPorts) > 0 {
 		for port, _ := range allowPorts {
@@ -176,14 +178,29 @@ func (pm *PortManager) Release(port int) {
 
 // Release reserved port if it isn't used in last 24 hours.
 func (pm *PortManager) cleanReservedPortsWorker() {
+	ticket := time.NewTicker(CleanReservedPortsInterval)
+FOR:
 	for {
-		time.Sleep(CleanReservedPortsInterval)
-		pm.mu.Lock()
-		for name, ctx := range pm.reservedPorts {
-			if ctx.Closed && time.Since(ctx.UpdateTime) > MaxPortReservedDuration {
-				delete(pm.reservedPorts, name)
+		select {
+		case <-ticket.C:
+			pm.mu.Lock()
+			for name, ctx := range pm.reservedPorts {
+				if ctx.Closed && time.Since(ctx.UpdateTime) > MaxPortReservedDuration {
+					delete(pm.reservedPorts, name)
+				}
+			}
+			pm.mu.Unlock()
+		case s := <-pm.stop:
+			if s {
+				ticket.Stop()
+				close(pm.stop)
+				break FOR
 			}
 		}
-		pm.mu.Unlock()
 	}
 }
+
+// Stop Stop
+func (pm *PortManager) Stop() {
+	pm.stop <- true
+}
diff --git a/server/service.go b/server/service.go
index 9e360cb1..453fe4a7 100644
--- a/server/service.go
+++ b/server/service.go
@@ -250,14 +250,24 @@ func (svr *Service) Run() {
 // Stop 停止服务
 func (svr *Service) Stop() error {
 	err := svr.muxer.Close()
-	<-svr.closedCh
+	if svr.listener != nil {
+		_ = svr.listener.Close()
+	}
+	if svr.websocketListener != nil {
+		_ = svr.websocketListener.Close()
+	}
+	if svr.kcpListener != nil {
+		_ = svr.kcpListener.Close()
+	}
+	close(svr.closedCh)
 	svr.Closed = true
+	svr.rc.TcpPortManager.Stop()
+	svr.rc.UdpPortManager.Stop()
 	return err
 }
 
 func (svr *Service) HandleListener(l frpNet.Listener) {
 	defer func() {
-		close(svr.closedCh)
 		log.Info("Frps is Closed")
 	}()
 	// Listen for incoming connections from client.
diff --git a/utils/net/websocket.go b/utils/net/websocket.go
index 890af5f1..4addd65f 100644
--- a/utils/net/websocket.go
+++ b/utils/net/websocket.go
@@ -76,6 +76,14 @@ func (p *WebsocketListener) Accept() (Conn, error) {
 }
 
 func (p *WebsocketListener) Close() error {
+	if p.accept != nil {
+		close(p.accept)
+	}
+	if p.ln != nil {
+		if err := p.ln.Close(); err != nil {
+			return err
+		}
+	}
 	return p.server.Close()
 }