From aede4e54f8911a0496049dc3de1edc8f7dbebdde Mon Sep 17 00:00:00 2001
From: fatedier <fatedier@gmail.com>
Date: Tue, 27 Jun 2017 01:59:30 +0800
Subject: [PATCH] close all proxies if protocol = kcp

---
 client/control.go | 37 ++++++++++++++++++++++++++++++++++---
 client/service.go |  4 ++++
 cmd/frpc/main.go  | 18 ++++++++++++++++++
 server/control.go |  1 +
 4 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/client/control.go b/client/control.go
index 45f617b..825d024 100644
--- a/client/control.go
+++ b/client/control.go
@@ -24,6 +24,7 @@ import (
 	"github.com/fatedier/frp/models/config"
 	"github.com/fatedier/frp/models/msg"
 	"github.com/fatedier/frp/utils/crypto"
+	"github.com/fatedier/frp/utils/errors"
 	"github.com/fatedier/frp/utils/log"
 	frpNet "github.com/fatedier/frp/utils/net"
 	"github.com/fatedier/frp/utils/util"
@@ -69,8 +70,8 @@ type Control struct {
 	// run id got from server
 	runId string
 
-	// connection or other error happens , control will try to reconnect to server
-	closed int32
+	// if we call close() in control, do not reconnect to server
+	exit bool
 
 	// goroutines can block by reading from this channel, it will be closed only in reader() when control connection is closed
 	closedCh chan int
@@ -181,7 +182,10 @@ func (ctl *Control) NewWorkConn() {
 	workConn.AddLogPrefix(startMsg.ProxyName)
 
 	// dispatch this work connection to related proxy
-	if pxy, ok := ctl.proxies[startMsg.ProxyName]; ok {
+	ctl.mu.RLock()
+	pxy, ok := ctl.proxies[startMsg.ProxyName]
+	ctl.mu.RUnlock()
+	if ok {
 		workConn.Debug("start a new work connection, localAddr: %s remoteAddr: %s", workConn.LocalAddr().String(), workConn.RemoteAddr().String())
 		go pxy.InWorkConn(workConn)
 	} else {
@@ -189,6 +193,20 @@ func (ctl *Control) NewWorkConn() {
 	}
 }
 
+func (ctl *Control) Close() error {
+	ctl.mu.Lock()
+	ctl.exit = true
+	err := errors.PanicToError(func() {
+		for name, _ := range ctl.proxies {
+			ctl.sendCh <- &msg.CloseProxy{
+				ProxyName: name,
+			}
+		}
+	})
+	ctl.mu.Unlock()
+	return err
+}
+
 func (ctl *Control) init() {
 	ctl.sendCh = make(chan msg.Message, 10)
 	ctl.readCh = make(chan msg.Message, 10)
@@ -377,7 +395,10 @@ func (ctl *Control) manager() {
 					ctl.Warn("[%s] no proxy conf found", m.ProxyName)
 					continue
 				}
+
+				ctl.mu.RLock()
 				oldPxy, ok := ctl.proxies[m.ProxyName]
+				ctl.mu.RUnlock()
 				if ok {
 					oldPxy.Close()
 				}
@@ -389,7 +410,9 @@ func (ctl *Control) manager() {
 					}
 					continue
 				}
+				ctl.mu.Lock()
 				ctl.proxies[m.ProxyName] = pxy
+				ctl.mu.Unlock()
 				ctl.Info("[%s] start proxy success", m.ProxyName)
 			case *msg.Pong:
 				ctl.lastPong = time.Now()
@@ -429,6 +452,14 @@ func (ctl *Control) controler() {
 				for _, pxy := range ctl.proxies {
 					pxy.Close()
 				}
+				// if ctl.exit is true, just exit
+				ctl.mu.RLock()
+				exit := ctl.exit
+				ctl.mu.RUnlock()
+				if exit {
+					return
+				}
+
 				time.Sleep(time.Second)
 
 				// loop util reconnect to server success
diff --git a/client/service.go b/client/service.go
index 776951a..ff28cc9 100644
--- a/client/service.go
+++ b/client/service.go
@@ -41,3 +41,7 @@ func (svr *Service) Run() error {
 	<-svr.closedCh
 	return nil
 }
+
+func (svr *Service) Close() error {
+	return svr.ctl.Close()
+}
diff --git a/cmd/frpc/main.go b/cmd/frpc/main.go
index 870eeae..88fc59b 100644
--- a/cmd/frpc/main.go
+++ b/cmd/frpc/main.go
@@ -17,8 +17,11 @@ package main
 import (
 	"fmt"
 	"os"
+	"os/signal"
 	"strconv"
 	"strings"
+	"syscall"
+	"time"
 
 	docopt "github.com/docopt/docopt-go"
 	ini "github.com/vaughan0/go-ini"
@@ -116,9 +119,24 @@ func main() {
 		config.ClientCommonCfg.LogLevel, config.ClientCommonCfg.LogMaxDays)
 
 	svr := client.NewService(pxyCfgs, vistorCfgs)
+
+	// Capture the exit signal if we use kcp.
+	if config.ClientCommonCfg.Protocol == "kcp" {
+		go HandleSignal(svr)
+	}
+
 	err = svr.Run()
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(1)
 	}
 }
+
+func HandleSignal(svr *client.Service) {
+	ch := make(chan os.Signal)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+	<-ch
+	svr.Close()
+	time.Sleep(250 * time.Millisecond)
+	os.Exit(0)
+}
diff --git a/server/control.go b/server/control.go
index d6b2c2c..5a84394 100644
--- a/server/control.go
+++ b/server/control.go
@@ -378,6 +378,7 @@ func (ctl *Control) CloseProxy(closeMsg *msg.CloseProxy) (err error) {
 
 	pxy.Close()
 	ctl.svr.DelProxy(pxy.GetName())
+	delete(ctl.proxies, closeMsg.ProxyName)
 	StatsCloseProxy(pxy.GetName(), pxy.GetConf().GetBaseInfo().ProxyType)
 	return
 }