frp/client/control.go

448 lines
11 KiB
Go
Raw Normal View History

2017-03-09 02:03:47 +08:00
// Copyright 2017 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package client
import (
"fmt"
"io"
2018-05-05 00:09:39 +08:00
"io/ioutil"
2017-03-09 02:03:47 +08:00
"runtime"
"sync"
"time"
2018-04-10 17:46:49 +08:00
"github.com/fatedier/frp/g"
2017-03-09 02:03:47 +08:00
"github.com/fatedier/frp/models/config"
"github.com/fatedier/frp/models/msg"
2017-03-10 01:42:06 +08:00
"github.com/fatedier/frp/utils/crypto"
2017-03-09 02:03:47 +08:00
"github.com/fatedier/frp/utils/log"
2017-06-26 03:02:33 +08:00
frpNet "github.com/fatedier/frp/utils/net"
2018-01-17 01:09:33 +08:00
"github.com/fatedier/frp/utils/shutdown"
2017-03-09 02:03:47 +08:00
"github.com/fatedier/frp/utils/util"
"github.com/fatedier/frp/utils/version"
2018-04-25 02:34:07 +08:00
fmux "github.com/hashicorp/yamux"
2017-03-09 02:03:47 +08:00
)
2017-05-10 00:46:42 +08:00
const (
connReadTimeout time.Duration = 10 * time.Second
)
2017-03-09 02:03:47 +08:00
type Control struct {
// frpc service
svr *Service
2018-01-17 01:09:33 +08:00
// login message to server, only used
2017-03-09 02:03:47 +08:00
loginMsg *msg.Login
2018-01-17 01:09:33 +08:00
pm *ProxyManager
2017-06-26 03:02:33 +08:00
2017-03-09 02:03:47 +08:00
// control connection
2017-06-26 03:02:33 +08:00
conn frpNet.Conn
2017-03-09 02:03:47 +08:00
// tcp stream multiplexing, if enabled
2018-04-25 02:34:07 +08:00
session *fmux.Session
2017-03-09 02:03:47 +08:00
// put a message in this channel to send it over control connection to server
sendCh chan (msg.Message)
// read from this channel to get the next message sent by server
readCh chan (msg.Message)
// run id got from server
runId string
2017-06-27 01:59:30 +08:00
// if we call close() in control, do not reconnect to server
exit bool
2017-03-09 02:03:47 +08:00
// goroutines can block by reading from this channel, it will be closed only in reader() when control connection is closed
closedCh chan int
// last time got the Pong message
lastPong time.Time
2018-01-17 01:09:33 +08:00
readerShutdown *shutdown.Shutdown
writerShutdown *shutdown.Shutdown
msgHandlerShutdown *shutdown.Shutdown
2017-03-09 02:03:47 +08:00
mu sync.RWMutex
log.Logger
}
2017-12-05 01:34:33 +08:00
func NewControl(svr *Service, pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) *Control {
2017-03-09 02:03:47 +08:00
loginMsg := &msg.Login{
Arch: runtime.GOARCH,
Os: runtime.GOOS,
2018-04-10 17:46:49 +08:00
PoolCount: g.GlbClientCfg.PoolCount,
User: g.GlbClientCfg.User,
2017-03-09 02:03:47 +08:00
Version: version.Full(),
}
2018-01-17 01:09:33 +08:00
ctl := &Control{
svr: svr,
loginMsg: loginMsg,
sendCh: make(chan msg.Message, 100),
readCh: make(chan msg.Message, 100),
2018-01-17 01:09:33 +08:00
closedCh: make(chan int),
readerShutdown: shutdown.New(),
writerShutdown: shutdown.New(),
msgHandlerShutdown: shutdown.New(),
Logger: log.NewPrefixLogger(""),
2017-03-09 02:03:47 +08:00
}
2018-01-17 01:09:33 +08:00
ctl.pm = NewProxyManager(ctl, ctl.sendCh, "")
ctl.pm.Reload(pxyCfgs, visitorCfgs, false)
2018-01-17 01:09:33 +08:00
return ctl
2017-03-09 02:03:47 +08:00
}
2017-06-26 03:02:33 +08:00
func (ctl *Control) Run() (err error) {
for {
2017-06-26 03:02:33 +08:00
err = ctl.login()
if err != nil {
2017-06-26 03:02:33 +08:00
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
2018-04-10 17:46:49 +08:00
if g.GlbClientCfg.LoginFailExit {
2017-06-26 03:02:33 +08:00
return
} else {
2018-01-17 01:09:33 +08:00
time.Sleep(10 * time.Second)
}
} else {
break
}
2017-03-09 02:03:47 +08:00
}
2018-01-17 01:09:33 +08:00
go ctl.worker()
2017-06-27 23:30:09 +08:00
2018-01-17 01:09:33 +08:00
// start all local visitors and send NewProxy message for all configured proxies
ctl.pm.Reset(ctl.sendCh, ctl.runId)
ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew})
2017-06-26 03:02:33 +08:00
return nil
}
2018-01-17 01:09:33 +08:00
func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
2017-06-26 03:02:33 +08:00
workConn, err := ctl.connectServer()
if err != nil {
return
2017-03-09 02:03:47 +08:00
}
m := &msg.NewWorkConn{
2018-01-17 01:09:33 +08:00
RunId: ctl.runId,
2017-03-09 02:03:47 +08:00
}
if err = msg.WriteMsg(workConn, m); err != nil {
ctl.Warn("work connection write to server error: %v", err)
workConn.Close()
return
}
var startMsg msg.StartWorkConn
if err = msg.ReadMsgInto(workConn, &startMsg); err != nil {
2017-03-10 01:42:06 +08:00
ctl.Error("work connection closed, %v", err)
2017-03-09 02:03:47 +08:00
workConn.Close()
return
}
workConn.AddLogPrefix(startMsg.ProxyName)
// dispatch this work connection to related proxy
2018-01-17 01:09:33 +08:00
ctl.pm.HandleWorkConn(startMsg.ProxyName, workConn)
}
func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
// Server will return NewProxyResp message to each NewProxy message.
// Start a new proxy handler if no error got
2018-01-17 14:40:08 +08:00
err := ctl.pm.StartProxy(inMsg.ProxyName, inMsg.RemoteAddr, inMsg.Error)
2018-01-17 01:09:33 +08:00
if err != nil {
ctl.Warn("[%s] start error: %v", inMsg.ProxyName, err)
2017-03-13 02:44:47 +08:00
} else {
2018-01-17 01:09:33 +08:00
ctl.Info("[%s] start proxy success", inMsg.ProxyName)
2017-03-09 02:03:47 +08:00
}
}
2017-06-27 01:59:30 +08:00
func (ctl *Control) Close() error {
ctl.mu.Lock()
2018-01-17 01:09:33 +08:00
defer ctl.mu.Unlock()
2017-06-27 01:59:30 +08:00
ctl.exit = true
2018-01-17 01:09:33 +08:00
ctl.pm.CloseProxies()
return nil
2017-03-09 02:03:47 +08:00
}
// 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,
2018-04-10 17:46:49 +08:00
fmt.Sprintf("%s:%d", g.GlbClientCfg.ServerAddr, g.GlbClientCfg.ServerPort))
2017-03-09 02:03:47 +08:00
if err != nil {
return err
}
2017-05-10 00:46:42 +08:00
defer func() {
if err != nil {
conn.Close()
}
}()
2018-04-10 17:46:49 +08:00
if g.GlbClientCfg.TcpMux {
2018-05-05 00:09:39 +08:00
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
}
2017-06-26 03:02:33 +08:00
conn = frpNet.WrapConn(stream)
ctl.session = session
}
2017-03-09 02:03:47 +08:00
now := time.Now().Unix()
2018-04-10 17:46:49 +08:00
ctl.loginMsg.PrivilegeKey = util.GetAuthKey(g.GlbClientCfg.Token, now)
2017-03-09 02:03:47 +08:00
ctl.loginMsg.Timestamp = now
2018-01-17 01:09:33 +08:00
ctl.loginMsg.RunId = ctl.runId
2017-03-09 02:03:47 +08:00
if err = msg.WriteMsg(conn, ctl.loginMsg); err != nil {
return err
}
var loginRespMsg msg.LoginResp
2017-05-10 00:46:42 +08:00
conn.SetReadDeadline(time.Now().Add(connReadTimeout))
2017-03-09 02:03:47 +08:00
if err = msg.ReadMsgInto(conn, &loginRespMsg); err != nil {
return err
}
2017-05-10 00:46:42 +08:00
conn.SetReadDeadline(time.Time{})
2017-03-09 02:03:47 +08:00
if loginRespMsg.Error != "" {
err = fmt.Errorf("%s", loginRespMsg.Error)
ctl.Error("%s", loginRespMsg.Error)
return err
}
ctl.conn = conn
// update runId got from server
2018-01-17 01:09:33 +08:00
ctl.runId = loginRespMsg.RunId
2018-04-10 17:46:49 +08:00
g.GlbClientCfg.ServerUdpPort = loginRespMsg.ServerUdpPort
2017-03-09 02:03:47 +08:00
ctl.ClearLogPrefix()
ctl.AddLogPrefix(loginRespMsg.RunId)
2017-10-24 18:20:07 +08:00
ctl.Info("login to server success, get run id [%s], server udp port [%d]", loginRespMsg.RunId, loginRespMsg.ServerUdpPort)
2017-03-09 02:03:47 +08:00
return nil
}
2017-06-26 03:02:33 +08:00
func (ctl *Control) connectServer() (conn frpNet.Conn, err error) {
2018-04-10 17:46:49 +08:00
if g.GlbClientCfg.TcpMux {
2017-06-26 03:02:33 +08:00
stream, errRet := ctl.session.OpenStream()
if errRet != nil {
err = errRet
ctl.Warn("start new connection to server error: %v", err)
return
}
conn = frpNet.WrapConn(stream)
} else {
conn, err = frpNet.ConnectServerByProxy(g.GlbClientCfg.HttpProxy, g.GlbClientCfg.Protocol,
2018-04-10 17:46:49 +08:00
fmt.Sprintf("%s:%d", g.GlbClientCfg.ServerAddr, g.GlbClientCfg.ServerPort))
2017-06-26 03:02:33 +08:00
if err != nil {
ctl.Warn("start new connection to server error: %v", err)
return
}
}
return
}
2018-01-17 01:09:33 +08:00
// reader read all messages from frps and send to readCh
2017-03-09 02:03:47 +08:00
func (ctl *Control) reader() {
defer func() {
if err := recover(); err != nil {
ctl.Error("panic error: %v", err)
}
}()
2018-01-17 01:09:33 +08:00
defer ctl.readerShutdown.Done()
2017-03-10 01:42:06 +08:00
defer close(ctl.closedCh)
2017-03-09 02:03:47 +08:00
2018-04-10 17:46:49 +08:00
encReader := crypto.NewReader(ctl.conn, []byte(g.GlbClientCfg.Token))
2017-03-09 02:03:47 +08:00
for {
2017-03-10 01:42:06 +08:00
if m, err := msg.ReadMsg(encReader); err != nil {
2017-03-09 02:03:47 +08:00
if err == io.EOF {
ctl.Debug("read from control connection EOF")
return
} else {
ctl.Warn("read error: %v", err)
2017-03-27 17:25:25 +08:00
return
2017-03-09 02:03:47 +08:00
}
} else {
ctl.readCh <- m
}
}
}
2018-01-17 01:09:33 +08:00
// writer writes messages got from sendCh to frps
2017-03-09 02:03:47 +08:00
func (ctl *Control) writer() {
2018-01-17 01:09:33 +08:00
defer ctl.writerShutdown.Done()
2018-04-10 17:46:49 +08:00
encWriter, err := crypto.NewWriter(ctl.conn, []byte(g.GlbClientCfg.Token))
2017-03-10 01:42:06 +08:00
if err != nil {
ctl.conn.Error("crypto new writer error: %v", err)
ctl.conn.Close()
return
}
2017-03-09 02:03:47 +08:00
for {
if m, ok := <-ctl.sendCh; !ok {
ctl.Info("control writer is closing")
return
} else {
2017-03-10 01:42:06 +08:00
if err := msg.WriteMsg(encWriter, m); err != nil {
2017-03-09 02:03:47 +08:00
ctl.Warn("write message to control connection error: %v", err)
return
}
}
}
}
2018-01-17 01:09:33 +08:00
// msgHandler handles all channel events and do corresponding operations.
func (ctl *Control) msgHandler() {
2017-03-09 02:03:47 +08:00
defer func() {
if err := recover(); err != nil {
ctl.Error("panic error: %v", err)
}
}()
2018-01-17 01:09:33 +08:00
defer ctl.msgHandlerShutdown.Done()
2017-03-09 02:03:47 +08:00
2018-04-10 17:46:49 +08:00
hbSend := time.NewTicker(time.Duration(g.GlbClientCfg.HeartBeatInterval) * time.Second)
2017-03-09 02:03:47 +08:00
defer hbSend.Stop()
hbCheck := time.NewTicker(time.Second)
defer hbCheck.Stop()
2018-01-17 01:09:33 +08:00
ctl.lastPong = time.Now()
2017-03-09 02:03:47 +08:00
for {
select {
case <-hbSend.C:
// send heartbeat to server
2017-03-27 17:25:25 +08:00
ctl.Debug("send heartbeat to server")
2017-03-09 02:03:47 +08:00
ctl.sendCh <- &msg.Ping{}
case <-hbCheck.C:
2018-04-10 17:46:49 +08:00
if time.Since(ctl.lastPong) > time.Duration(g.GlbClientCfg.HeartBeatTimeout)*time.Second {
2017-03-09 02:03:47 +08:00
ctl.Warn("heartbeat timeout")
2017-03-27 17:25:25 +08:00
// let reader() stop
ctl.conn.Close()
2017-03-09 02:03:47 +08:00
return
}
case rawMsg, ok := <-ctl.readCh:
if !ok {
return
}
switch m := rawMsg.(type) {
case *msg.ReqWorkConn:
2018-01-17 01:09:33 +08:00
go ctl.HandleReqWorkConn(m)
2017-03-09 02:03:47 +08:00
case *msg.NewProxyResp:
2018-01-17 01:09:33 +08:00
ctl.HandleNewProxyResp(m)
2017-03-09 02:03:47 +08:00
case *msg.Pong:
ctl.lastPong = time.Now()
2017-03-27 17:25:25 +08:00
ctl.Debug("receive heartbeat from server")
2017-03-09 02:03:47 +08:00
}
}
}
}
2017-06-27 23:30:09 +08:00
// controler keep watching closedCh, start a new connection if previous control connection is closed.
2018-01-17 01:09:33 +08:00
// If controler is notified by closedCh, reader and writer and handler will exit, then recall these functions.
func (ctl *Control) worker() {
go ctl.msgHandler()
go ctl.reader()
go ctl.writer()
2018-01-17 01:09:33 +08:00
2017-03-09 02:03:47 +08:00
var err error
2018-01-17 01:09:33 +08:00
maxDelayTime := 20 * time.Second
2017-03-09 02:03:47 +08:00
delayTime := time.Second
2017-03-12 02:03:24 +08:00
checkInterval := 60 * time.Second
2017-03-12 02:03:24 +08:00
checkProxyTicker := time.NewTicker(checkInterval)
2017-03-09 02:03:47 +08:00
for {
2017-03-12 02:03:24 +08:00
select {
case <-checkProxyTicker.C:
// check which proxy registered failed and reregister it to server
ctl.pm.CheckAndStartProxy([]string{ProxyStatusStartErr, ProxyStatusClosed})
2017-03-12 02:03:24 +08:00
case _, ok := <-ctl.closedCh:
// we won't get any variable from this channel
if !ok {
2018-01-17 01:09:33 +08:00
// close related channels and wait until other goroutines done
2017-03-12 02:03:24 +08:00
close(ctl.readCh)
2018-01-17 01:09:33 +08:00
ctl.readerShutdown.WaitDone()
ctl.msgHandlerShutdown.WaitDone()
2017-03-12 02:03:24 +08:00
close(ctl.sendCh)
2018-01-17 01:09:33 +08:00
ctl.writerShutdown.WaitDone()
2017-04-25 00:34:14 +08:00
2018-01-17 01:09:33 +08:00
ctl.pm.CloseProxies()
2017-06-27 01:59:30 +08:00
// if ctl.exit is true, just exit
ctl.mu.RLock()
exit := ctl.exit
ctl.mu.RUnlock()
if exit {
return
}
2018-01-17 01:09:33 +08:00
// loop util reconnecting to server success
2017-03-12 02:03:24 +08:00
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
}
2018-01-17 01:09:33 +08:00
// reconnect success, init delayTime
2017-03-12 02:03:24 +08:00
delayTime = time.Second
break
}
2017-03-09 02:03:47 +08:00
2017-03-12 02:03:24 +08:00
// init related channels and variables
ctl.sendCh = make(chan msg.Message, 100)
ctl.readCh = make(chan msg.Message, 100)
2018-01-17 01:09:33 +08:00
ctl.closedCh = make(chan int)
ctl.readerShutdown = shutdown.New()
ctl.writerShutdown = shutdown.New()
ctl.msgHandlerShutdown = shutdown.New()
ctl.pm.Reset(ctl.sendCh, ctl.runId)
2017-03-09 02:03:47 +08:00
2017-03-12 02:03:24 +08:00
// previous work goroutines should be closed and start them here
2018-01-17 01:09:33 +08:00
go ctl.msgHandler()
2017-03-12 02:03:24 +08:00
go ctl.writer()
go ctl.reader()
2017-03-09 02:03:47 +08:00
2018-01-17 01:09:33 +08:00
// start all configured proxies
2018-01-26 14:56:55 +08:00
ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew, ProxyStatusClosed})
2017-03-12 02:03:24 +08:00
checkProxyTicker.Stop()
checkProxyTicker = time.NewTicker(checkInterval)
2017-03-09 02:03:47 +08:00
}
}
}
}
2017-06-27 23:30:09 +08:00
2018-01-17 01:09:33 +08:00
func (ctl *Control) reloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) error {
err := ctl.pm.Reload(pxyCfgs, visitorCfgs, true)
2018-01-17 01:09:33 +08:00
return err
}