From 5febee62013d568bdccad83c6e4efce5a52ac391 Mon Sep 17 00:00:00 2001 From: fatedier Date: Fri, 3 Jun 2016 17:51:45 +0800 Subject: [PATCH] all: add "--reload" command for frps, reload ini file without kill old program --- Makefile.cross-compiles | 2 +- conf/frps.ini | 2 +- src/frp/cmd/frps/control.go | 5 +- src/frp/cmd/frps/main.go | 55 ++++++++++++--- src/frp/models/server/config.go | 93 +++++++++++++++++++++----- src/frp/models/server/dashboard.go | 2 +- src/frp/models/server/dashboard_api.go | 28 ++++++-- src/frp/models/server/server.go | 40 ++++++++++- src/frp/utils/conn/conn.go | 12 +++- 9 files changed, 200 insertions(+), 39 deletions(-) diff --git a/Makefile.cross-compiles b/Makefile.cross-compiles index e8038cd..76c287a 100644 --- a/Makefile.cross-compiles +++ b/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/... diff --git a/conf/frps.ini b/conf/frps.ini index e252b53..e3087d8 100644 --- a/conf/frps.ini +++ b/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 diff --git a/src/frp/cmd/frps/control.go b/src/frp/cmd/frps/control.go index 08da7fe..0784955 100644 --- a/src/frp/cmd/frps/control.go +++ b/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) diff --git a/src/frp/cmd/frps/main.go b/src/frp/cmd/frps/main.go index 97ba737..62ee934 100644 --- a/src/frp/cmd/frps/main.go +++ b/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=] [--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= set log level: debug, info, warn, error --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 { diff --git a/src/frp/models/server/config.go b/src/frp/models/server/config.go index b568411..14b9753 100644 --- a/src/frp/models/server/config.go +++ b/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 } diff --git a/src/frp/models/server/dashboard.go b/src/frp/models/server/dashboard.go index e8dd832..5ff20f9 100644 --- a/src/frp/models/server/dashboard.go +++ b/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 diff --git a/src/frp/models/server/dashboard_api.go b/src/frp/models/server/dashboard_api.go index 8332ada..85a3274 100644 --- a/src/frp/models/server/dashboard_api.go +++ b/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) } diff --git a/src/frp/models/server/server.go b/src/frp/models/server/server.go index 1e8f9bd..b9affdf 100644 --- a/src/frp/models/server/server.go +++ b/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() } diff --git a/src/frp/utils/conn/conn.go b/src/frp/utils/conn/conn.go index 29b4d22..ae79571 100644 --- a/src/frp/utils/conn/conn.go +++ b/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