mirrors/frp
1
0
Fork 0
mirror of https://github.com/fatedier/frp.git synced 2025-07-13 09:19:46 +07:00
Code Issues Packages Projects Releases Wiki Activity

client/control: refactor code

Browse Source
This commit is contained in:
fatedier
2018-11-06 18:35:05 +08:00
parent 1a8ac148ca
commit b33ea9274c
6 changed files with 323 additions and 313 deletions
Download Patch File Download Diff File Expand all files Collapse all files

213
client/control.go
View File

@ -17,8 +17,6 @@ package client
import (
"fmt"
"io"
"io/ioutil"
"runtime"
"runtime/debug"
"sync"
"time"
@ -28,24 +26,15 @@ import (
"github.com/fatedier/frp/models/msg"
"github.com/fatedier/frp/utils/log"
frpNet "github.com/fatedier/frp/utils/net"
"github.com/fatedier/frp/utils/util"
"github.com/fatedier/frp/utils/version"
"github.com/fatedier/golib/control/shutdown"
"github.com/fatedier/golib/crypto"
fmux "github.com/hashicorp/yamux"
)
const (
connReadTimeout time.Duration = 10 * time.Second
)
type Control struct {
// frpc service
svr *Service
// login message to server, only used
loginMsg *msg.Login
// uniq id got from frps, attach it in loginMsg
runId string
// manage all proxies
pm *ProxyManager
@ -65,14 +54,10 @@ type Control struct {
// read from this channel to get the next message sent by server
readCh chan (msg.Message)
// run id got from server
runId string
// 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
closedCh chan struct{}
closedDoneCh chan struct{}
// last time got the Pong message
lastPong time.Time
@ -86,50 +71,28 @@ type Control struct {
log.Logger
}
func NewControl(svr *Service, pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) *Control {
loginMsg := &msg.Login{
Arch: runtime.GOARCH,
Os: runtime.GOOS,
PoolCount: g.GlbClientCfg.PoolCount,
User: g.GlbClientCfg.User,
Version: version.Full(),
}
func NewControl(runId string, conn frpNet.Conn, session *fmux.Session, pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) *Control {
ctl := &Control{
svr: svr,
loginMsg: loginMsg,
runId: runId,
conn: conn,
session: session,
sendCh: make(chan msg.Message, 100),
readCh: make(chan msg.Message, 100),
closedCh: make(chan int),
closedCh: make(chan struct{}),
closedDoneCh: make(chan struct{}),
readerShutdown: shutdown.New(),
writerShutdown: shutdown.New(),
msgHandlerShutdown: shutdown.New(),
Logger: log.NewPrefixLogger(""),
}
ctl.pm = NewProxyManager(ctl, ctl.sendCh, "")
ctl.pm = NewProxyManager(ctl.sendCh, "")
ctl.pm.Reload(pxyCfgs, false)
ctl.vm = NewVisitorManager(ctl)
ctl.vm.Reload(visitorCfgs)
return ctl
}
func (ctl *Control) Run() (err error) {
for {
err = ctl.login()
if err != nil {
ctl.Warn("login to server failed: %v", err)
// if login_fail_exit is true, just exit this program
// otherwise sleep a while and continues relogin to server
if g.GlbClientCfg.LoginFailExit {
return
} else {
time.Sleep(10 * time.Second)
}
} else {
break
}
}
func (ctl *Control) Run() {
go ctl.worker()
// start all local visitors and send NewProxy message for all configured proxies
@ -137,7 +100,7 @@ func (ctl *Control) Run() (err error) {
ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew})
go ctl.vm.Run()
return nil
return
}
func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
@ -179,80 +142,13 @@ func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
}
func (ctl *Control) Close() error {
ctl.mu.Lock()
defer ctl.mu.Unlock()
ctl.exit = true
ctl.pm.CloseProxies()
return nil
}
// login send a login message to server and wait for a loginResp message.
func (ctl *Control) login() (err error) {
if ctl.conn != nil {
ctl.conn.Close()
}
if ctl.session != nil {
ctl.session.Close()
}
conn, err := frpNet.ConnectServerByProxy(g.GlbClientCfg.HttpProxy, g.GlbClientCfg.Protocol,
fmt.Sprintf("%s:%d", g.GlbClientCfg.ServerAddr, g.GlbClientCfg.ServerPort))
if err != nil {
return err
}
defer func() {
if err != nil {
conn.Close()
}
}()
if g.GlbClientCfg.TcpMux {
fmuxCfg := fmux.DefaultConfig()
fmuxCfg.LogOutput = ioutil.Discard
session, errRet := fmux.Client(conn, fmuxCfg)
if errRet != nil {
return errRet
}
stream, errRet := session.OpenStream()
if errRet != nil {
session.Close()
return errRet
}
conn = frpNet.WrapConn(stream)
ctl.session = session
}
now := time.Now().Unix()
ctl.loginMsg.PrivilegeKey = util.GetAuthKey(g.GlbClientCfg.Token, now)
ctl.loginMsg.Timestamp = now
ctl.loginMsg.RunId = ctl.runId
if err = msg.WriteMsg(conn, ctl.loginMsg); err != nil {
return err
}
var loginRespMsg msg.LoginResp
conn.SetReadDeadline(time.Now().Add(connReadTimeout))
if err = msg.ReadMsgInto(conn, &loginRespMsg); err != nil {
return err
}
conn.SetReadDeadline(time.Time{})
if loginRespMsg.Error != "" {
err = fmt.Errorf("%s", loginRespMsg.Error)
ctl.Error("%s", loginRespMsg.Error)
return err
}
ctl.conn = conn
// update runId got from server
ctl.runId = loginRespMsg.RunId
g.GlbClientCfg.ServerUdpPort = loginRespMsg.ServerUdpPort
ctl.ClearLogPrefix()
ctl.AddLogPrefix(loginRespMsg.RunId)
ctl.Info("login to server success, get run id [%s], server udp port [%d]", loginRespMsg.RunId, loginRespMsg.ServerUdpPort)
return nil
// ClosedDoneCh returns a channel which will be closed after all resources are released
func (ctl *Control) ClosedDoneCh() <-chan struct{} {
return ctl.closedDoneCh
}
// connectServer return a new connection to frps
@ -373,87 +269,38 @@ func (ctl *Control) msgHandler() {
}
}
// controler keep watching closedCh, start a new connection if previous control connection is closed.
// If controler is notified by closedCh, reader and writer and handler will exit, then recall these functions.
// If controler is notified by closedCh, reader and writer and handler will exit
func (ctl *Control) worker() {
go ctl.msgHandler()
go ctl.reader()
go ctl.writer()
var err error
maxDelayTime := 20 * time.Second
delayTime := time.Second
checkInterval := 60 * time.Second
checkProxyTicker := time.NewTicker(checkInterval)
for {
select {
case <-checkProxyTicker.C:
// check which proxy registered failed and reregister it to server
ctl.pm.CheckAndStartProxy([]string{ProxyStatusStartErr, ProxyStatusClosed})
case _, ok := <-ctl.closedCh:
// we won't get any variable from this channel
if !ok {
// close related channels and wait until other goroutines done
close(ctl.readCh)
ctl.readerShutdown.WaitDone()
ctl.msgHandlerShutdown.WaitDone()
case <-ctl.closedCh:
// close related channels and wait until other goroutines done
close(ctl.readCh)
ctl.readerShutdown.WaitDone()
ctl.msgHandlerShutdown.WaitDone()
close(ctl.sendCh)
ctl.writerShutdown.WaitDone()
close(ctl.sendCh)
ctl.writerShutdown.WaitDone()
ctl.pm.CloseProxies()
// if ctl.exit is true, just exit
ctl.mu.RLock()
exit := ctl.exit
ctl.mu.RUnlock()
if exit {
return
}
ctl.pm.CloseProxies()
// loop util reconnecting to server success
for {
ctl.Info("try to reconnect to server...")
err = ctl.login()
if err != nil {
ctl.Warn("reconnect to server error: %v", err)
time.Sleep(delayTime)
delayTime = delayTime * 2
if delayTime > maxDelayTime {
delayTime = maxDelayTime
}
continue
}
// reconnect success, init delayTime
delayTime = time.Second
break
}
// init related channels and variables
ctl.sendCh = make(chan msg.Message, 100)
ctl.readCh = make(chan msg.Message, 100)
ctl.closedCh = make(chan int)
ctl.readerShutdown = shutdown.New()
ctl.writerShutdown = shutdown.New()
ctl.msgHandlerShutdown = shutdown.New()
ctl.pm.Reset(ctl.sendCh, ctl.runId)
// previous work goroutines should be closed and start them here
go ctl.msgHandler()
go ctl.writer()
go ctl.reader()
// start all configured proxies
ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew, ProxyStatusClosed})
checkProxyTicker.Stop()
checkProxyTicker = time.NewTicker(checkInterval)
}
close(ctl.closedDoneCh)
return
}
}
}
func (ctl *Control) reloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {
func (ctl *Control) ReloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.VisitorConf) error {
ctl.vm.Reload(visitorCfgs)
err := ctl.pm.Reload(pxyCfgs, true)
return err