Browse Source

all: add "--reload" command for frps, reload ini file without kill old program

fatedier 8 years ago
parent
commit
5febee6201

+ 1 - 1
Makefile.cross-compiles

@@ -10,4 +10,4 @@ godep:
 	GOPATH=$(OLDGOPATH) go get github.com/mitchellh/gox
 
 app:
-	gox -osarch "darwin/386 darwin/amd64 linux/386 linux/amd64 linux/arm windows/386 windows/amd64" ./...
+	gox -osarch "darwin/386 darwin/amd64 linux/386 linux/amd64 linux/arm windows/386 windows/amd64" ./src/...

+ 1 - 1
conf/frps.ini

@@ -4,7 +4,7 @@ bind_addr = 0.0.0.0
 bind_port = 7000
 # if you want to support virtual host, you must set the http port for listening (optional)
 vhost_http_port = 80
-# if you want to configure or reload frps by dashboard, dashboard_port is needed 
+# if you want to configure or reload frps by dashboard, dashboard_port must be set
 dashboard_port = 7500
 # console or real logFile path like ./frps.log
 log_file = ./frps.log

+ 2 - 3
src/frp/cmd/frps/control.go

@@ -67,7 +67,7 @@ func controlWorker(c *conn.Conn) {
 		return
 	}
 
-	// do login when type is NewCtlConn or NewWorkConn
+	// login when type is NewCtlConn or NewWorkConn
 	ret, info := doLogin(cliReq, c)
 	s, ok := server.ProxyServers[cliReq.ProxyName]
 	if !ok {
@@ -134,7 +134,6 @@ func msgReader(s *server.ProxyServer, c *conn.Conn, msgSendChan chan interface{}
 	timer := time.AfterFunc(time.Duration(server.HeartBeatTimeout)*time.Second, func() {
 		heartbeatTimeout = true
 		s.Close()
-		c.Close()
 		log.Error("ProxyName [%s], client heartbeat timeout", s.Name)
 	})
 	defer timer.Stop()
@@ -229,7 +228,7 @@ func doLogin(req *msg.ControlReq, c *conn.Conn) (ret int64, info string) {
 		s.UseEncryption = req.UseEncryption
 
 		// start proxy and listen for user connections, no block
-		err := s.Start()
+		err := s.Start(c)
 		if err != nil {
 			info = fmt.Sprintf("ProxyName [%s], start proxy error: %v", req.ProxyName, err)
 			log.Warn(info)

+ 46 - 9
src/frp/cmd/frps/main.go

@@ -15,7 +15,10 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
+	"io/ioutil"
+	"net/http"
 	"os"
 	"strconv"
 	"strings"
@@ -30,15 +33,13 @@ import (
 	"frp/utils/vhost"
 )
 
-var (
-	configFile string = "./frps.ini"
-)
-
 var usage string = `frps is the server of frp
 
 Usage: 
 	frps [-c config_file] [-L log_file] [--log-level=<log_level>] [--addr=<bind_addr>]
-	frps -h | --help | --version
+	frps --reload
+	frps -h | --help
+	frps -v | --version
 
 Options:
 	-c config_file            set config file
@@ -46,7 +47,7 @@ Options:
 	--log-level=<log_level>   set log level: debug, info, warn, error
 	--addr=<bind_addr>        listen addr for client, example: 0.0.0.0:7000
 	-h --help                 show this screen
-	--version                 show version
+	-v --version              show version
 `
 
 func main() {
@@ -54,14 +55,43 @@ func main() {
 	args, err := docopt.Parse(usage, nil, true, version.Full(), false)
 
 	if args["-c"] != nil {
-		configFile = args["-c"].(string)
+		server.ConfigFile = args["-c"].(string)
 	}
-	err = server.LoadConf(configFile)
+	err = server.LoadConf(server.ConfigFile)
 	if err != nil {
 		fmt.Println(err)
 		os.Exit(-1)
 	}
 
+	// reload check
+	if args["--reload"] != nil {
+		if args["--reload"].(bool) {
+			resp, err := http.Get("http://" + server.BindAddr + ":" + fmt.Sprintf("%d", server.DashboardPort) + "/api/reload")
+			if err != nil {
+				fmt.Printf("frps reload error: %v\n", err)
+				os.Exit(1)
+			} else {
+				defer resp.Body.Close()
+				body, err := ioutil.ReadAll(resp.Body)
+				if err != nil {
+					fmt.Printf("frps reload error: %v\n", err)
+					os.Exit(1)
+				}
+				res := &server.GeneralResponse{}
+				err = json.Unmarshal(body, &res)
+				if err != nil {
+					fmt.Printf("http response error: %v\n", err)
+					os.Exit(1)
+				} else if res.Code != 0 {
+					fmt.Printf("reload error: %s\n", res.Msg)
+					os.Exit(1)
+				}
+				fmt.Printf("reload success\n")
+				os.Exit(0)
+			}
+		}
+	}
+
 	if args["-L"] != nil {
 		if args["-L"].(string) == "console" {
 			server.LogWay = "console"
@@ -90,6 +120,13 @@ func main() {
 		server.BindPort = bindPort
 	}
 
+	if args["-v"] != nil {
+		if args["-v"].(bool) {
+			fmt.Println(version.Full())
+			os.Exit(0)
+		}
+	}
+
 	log.InitLog(server.LogWay, server.LogFile, server.LogLevel, server.LogMaxDays)
 
 	l, err := conn.Listen(server.BindAddr, server.BindPort)
@@ -111,7 +148,7 @@ func main() {
 		}
 	}
 
-	// create dashboard web server if DashboardPort != 0
+	// create dashboard web server if DashboardPort is set, so it won't be 0
 	if server.DashboardPort != 0 {
 		err := server.RunDashboardServer(server.BindAddr, server.DashboardPort)
 		if err != nil {

+ 78 - 15
src/frp/models/server/config.go

@@ -18,14 +18,17 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"sync"
 
 	ini "github.com/vaughan0/go-ini"
 
+	"frp/utils/log"
 	"frp/utils/vhost"
 )
 
 // common config
 var (
+	ConfigFile       string = "./frps.ini"
 	BindAddr         string = "0.0.0.0"
 	BindPort         int64  = 7000
 	VhostHttpPort    int64  = 0 // if VhostHttpPort equals 0, don't listen a public port for http
@@ -37,20 +40,39 @@ var (
 	HeartBeatTimeout int64  = 90
 	UserConnTimeout  int64  = 10
 
-	VhostMuxer *vhost.HttpMuxer
+	VhostMuxer        *vhost.HttpMuxer
+	ProxyServers      map[string]*ProxyServer = make(map[string]*ProxyServer) // all proxy servers info and resources
+	ProxyServersMutex sync.RWMutex
 )
 
-var ProxyServers map[string]*ProxyServer = make(map[string]*ProxyServer)
-
 func LoadConf(confFile string) (err error) {
+	err = loadCommonConf(confFile)
+	if err != nil {
+		return err
+	}
+
+	// load all proxy server's configure and initialize
+	// and set ProxyServers map
+	newProxyServers, err := loadProxyConf(confFile)
+	if err != nil {
+		return err
+	}
+	for _, proxyServer := range newProxyServers {
+		proxyServer.Init()
+	}
+	ProxyServersMutex.Lock()
+	ProxyServers = newProxyServers
+	ProxyServersMutex.Unlock()
+	return nil
+}
+
+func loadCommonConf(confFile string) error {
 	var tmpStr string
 	var ok bool
-
 	conf, err := ini.LoadFile(confFile)
 	if err != nil {
 		return err
 	}
-
 	// common
 	tmpStr, ok = conf.Get("common", "bind_addr")
 	if ok {
@@ -95,18 +117,26 @@ func LoadConf(confFile string) (err error) {
 	if ok {
 		LogMaxDays, _ = strconv.ParseInt(tmpStr, 10, 64)
 	}
+	return nil
+}
 
+func loadProxyConf(confFile string) (proxyServers map[string]*ProxyServer, err error) {
+	var ok bool
+	proxyServers = make(map[string]*ProxyServer)
+	conf, err := ini.LoadFile(confFile)
+	if err != nil {
+		return proxyServers, err
+	}
 	// servers
 	for name, section := range conf {
 		if name != "common" {
-			proxyServer := &ProxyServer{}
-			proxyServer.CustomDomains = make([]string, 0)
+			proxyServer := NewProxyServer()
 			proxyServer.Name = name
 
 			proxyServer.Type, ok = section["type"]
 			if ok {
 				if proxyServer.Type != "tcp" && proxyServer.Type != "http" {
-					return fmt.Errorf("Parse ini file error: proxy [%s] type error", proxyServer.Name)
+					return proxyServers, fmt.Errorf("Parse conf error: proxy [%s] type error", proxyServer.Name)
 				}
 			} else {
 				proxyServer.Type = "tcp"
@@ -114,7 +144,7 @@ func LoadConf(confFile string) (err error) {
 
 			proxyServer.AuthToken, ok = section["auth_token"]
 			if !ok {
-				return fmt.Errorf("Parse ini file error: proxy [%s] no auth_token found", proxyServer.Name)
+				return proxyServers, fmt.Errorf("Parse conf error: proxy [%s] no auth_token found", proxyServer.Name)
 			}
 
 			// for tcp
@@ -128,10 +158,10 @@ func LoadConf(confFile string) (err error) {
 				if ok {
 					proxyServer.ListenPort, err = strconv.ParseInt(portStr, 10, 64)
 					if err != nil {
-						return fmt.Errorf("Parse ini file error: proxy [%s] listen_port error", proxyServer.Name)
+						return proxyServers, fmt.Errorf("Parse conf error: proxy [%s] listen_port error", proxyServer.Name)
 					}
 				} else {
-					return fmt.Errorf("Parse ini file error: proxy [%s] listen_port not found", proxyServer.Name)
+					return proxyServers, fmt.Errorf("Parse conf error: proxy [%s] listen_port not found", proxyServer.Name)
 				}
 			} else if proxyServer.Type == "http" {
 				// for http
@@ -142,20 +172,53 @@ func LoadConf(confFile string) (err error) {
 						suffix = fmt.Sprintf(":%d", VhostHttpPort)
 					}
 					proxyServer.CustomDomains = strings.Split(domainStr, ",")
+					if len(proxyServer.CustomDomains) == 0 {
+						return proxyServers, fmt.Errorf("Parse conf error: proxy [%s] custom_domains must be set when type equals http", proxyServer.Name)
+					}
 					for i, domain := range proxyServer.CustomDomains {
 						proxyServer.CustomDomains[i] = strings.ToLower(strings.TrimSpace(domain)) + suffix
 					}
 				}
 			}
+			proxyServers[proxyServer.Name] = proxyServer
+		}
+	}
+	return proxyServers, nil
+}
 
+// the function can only reload proxy configures
+// common section won't be changed
+func ReloadConf(confFile string) (err error) {
+	loadProxyServers, err := loadProxyConf(confFile)
+	if err != nil {
+		return err
+	}
+
+	ProxyServersMutex.Lock()
+	for name, proxyServer := range loadProxyServers {
+		oldProxyServer, ok := ProxyServers[name]
+		if ok {
+			if !oldProxyServer.Compare(proxyServer) {
+				oldProxyServer.Close()
+				proxyServer.Init()
+				ProxyServers[name] = proxyServer
+				log.Info("ProxyName [%s] configure change, restart", name)
+			}
+		} else {
 			proxyServer.Init()
-			ProxyServers[proxyServer.Name] = proxyServer
+			ProxyServers[name] = proxyServer
+			log.Info("ProxyName [%s] is new, init it", name)
 		}
 	}
 
-	if len(ProxyServers) == 0 {
-		return fmt.Errorf("Parse ini file error: no proxy config found")
+	for name, oldProxyServer := range ProxyServers {
+		_, ok := loadProxyServers[name]
+		if !ok {
+			oldProxyServer.Close()
+			delete(ProxyServers, name)
+			log.Info("ProxyName [%s] deleted, close it", name)
+		}
 	}
-
+	ProxyServersMutex.Unlock()
 	return nil
 }

+ 1 - 1
src/frp/models/server/dashboard.go

@@ -28,7 +28,7 @@ func RunDashboardServer(addr string, port int64) (err error) {
 	}()
 	gin.SetMode(gin.ReleaseMode)
 	router := gin.New()
-	router.LoadHTMLGlob("assets/*")
+	//router.LoadHTMLGlob("assets/*")
 	router.GET("/api/reload", apiReload)
 	go router.Run(fmt.Sprintf("%s:%d", addr, port))
 	return

+ 24 - 4
src/frp/models/server/dashboard_api.go

@@ -15,12 +15,32 @@
 package server
 
 import (
+	"encoding/json"
+	"fmt"
+
 	"github.com/gin-gonic/gin"
+
+	"frp/utils/log"
 )
 
+type GeneralResponse struct {
+	Code int64  `json:"code"`
+	Msg  string `json:"msg"`
+}
+
 func apiReload(c *gin.Context) {
-	c.JSON(200, gin.H{
-		"code": 0,
-		"msg":  "ok",
-	})
+	res := &GeneralResponse{}
+	defer func() {
+		buf, _ := json.Marshal(res)
+		log.Info("Http response [/api/reload]: %s", string(buf))
+	}()
+
+	log.Info("Http request: [/api/reload]")
+	err := ReloadConf(ConfigFile)
+	if err != nil {
+		res.Code = 2
+		res.Msg = fmt.Sprintf("%v", err)
+		log.Error("frps reload error: %v", err)
+	}
+	c.JSON(200, res)
 }

+ 37 - 3
src/frp/models/server/server.go

@@ -35,21 +35,49 @@ type ProxyServer struct {
 	Type          string
 	BindAddr      string
 	ListenPort    int64
-	UseEncryption bool
 	CustomDomains []string
 
+	// configure in frpc.ini
+	UseEncryption bool
+
 	Status       int64
+	CtlConn      *conn.Conn      // control connection with frpc
 	listeners    []Listener      // accept new connection from remote users
 	ctlMsgChan   chan int64      // every time accept a new user conn, put "1" to the channel
 	workConnChan chan *conn.Conn // get new work conns from control goroutine
 	mutex        sync.Mutex
 }
 
+func NewProxyServer() (p *ProxyServer) {
+	p = &ProxyServer{
+		CustomDomains: make([]string, 0),
+	}
+	return p
+}
+
 func (p *ProxyServer) Init() {
+	p.Lock()
 	p.Status = consts.Idle
 	p.workConnChan = make(chan *conn.Conn, 100)
 	p.ctlMsgChan = make(chan int64)
 	p.listeners = make([]Listener, 0)
+	p.Unlock()
+}
+
+func (p *ProxyServer) Compare(p2 *ProxyServer) bool {
+	if p.Name != p2.Name || p.AuthToken != p2.AuthToken || p.Type != p2.Type ||
+		p.BindAddr != p2.BindAddr || p.ListenPort != p2.ListenPort {
+		return false
+	}
+	if len(p.CustomDomains) != len(p2.CustomDomains) {
+		return false
+	}
+	for i, _ := range p.CustomDomains {
+		if p.CustomDomains[i] != p2.CustomDomains[i] {
+			return false
+		}
+	}
+	return true
 }
 
 func (p *ProxyServer) Lock() {
@@ -61,7 +89,8 @@ func (p *ProxyServer) Unlock() {
 }
 
 // start listening for user conns
-func (p *ProxyServer) Start() (err error) {
+func (p *ProxyServer) Start(c *conn.Conn) (err error) {
+	p.CtlConn = c
 	p.Init()
 	if p.Type == "tcp" {
 		l, err := conn.Listen(p.BindAddr, p.ListenPort)
@@ -79,9 +108,11 @@ func (p *ProxyServer) Start() (err error) {
 		}
 	}
 
+	p.Lock()
 	p.Status = consts.Working
+	p.Unlock()
 
-	// start a goroutine for listener to accept user connection
+	// start a goroutine for every listener to accept user connection
 	for _, listener := range p.listeners {
 		go func(l Listener) {
 			for {
@@ -138,6 +169,9 @@ func (p *ProxyServer) Close() {
 		}
 		close(p.ctlMsgChan)
 		close(p.workConnChan)
+		if p.CtlConn != nil {
+			p.CtlConn.Close()
+		}
 	}
 	p.Unlock()
 }

+ 10 - 2
src/frp/utils/conn/conn.go

@@ -92,6 +92,7 @@ type Conn struct {
 	TcpConn   net.Conn
 	Reader    *bufio.Reader
 	closeFlag bool
+	mutex     sync.RWMutex
 }
 
 func NewConn(conn net.Conn) (c *Conn) {
@@ -129,7 +130,9 @@ func (c *Conn) GetLocalAddr() (addr string) {
 func (c *Conn) ReadLine() (buff string, err error) {
 	buff, err = c.Reader.ReadString('\n')
 	if err == io.EOF {
+		c.mutex.Lock()
 		c.closeFlag = true
+		c.mutex.Unlock()
 	}
 	return buff, err
 }
@@ -146,14 +149,19 @@ func (c *Conn) SetDeadline(t time.Time) error {
 }
 
 func (c *Conn) Close() {
+	c.mutex.Lock()
 	if c.TcpConn != nil && c.closeFlag == false {
 		c.closeFlag = true
 		c.TcpConn.Close()
 	}
+	c.mutex.Unlock()
 }
 
-func (c *Conn) IsClosed() bool {
-	return c.closeFlag
+func (c *Conn) IsClosed() (closeFlag bool) {
+	c.mutex.RLock()
+	closeFlag = c.closeFlag
+	c.mutex.RUnlock()
+	return
 }
 
 // will block until connection close