mirror of
https://github.com/fatedier/frp.git
synced 2024-12-16 03:31:07 +01:00
623 lines
15 KiB
Go
623 lines
15 KiB
Go
// 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"
|
|
"runtime"
|
|
"sync"
|
|
"time"
|
|
|
|
"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"
|
|
"github.com/fatedier/frp/utils/version"
|
|
"github.com/xtaci/smux"
|
|
)
|
|
|
|
const (
|
|
connReadTimeout time.Duration = 10 * time.Second
|
|
)
|
|
|
|
type Control struct {
|
|
// frpc service
|
|
svr *Service
|
|
|
|
// login message to server
|
|
loginMsg *msg.Login
|
|
|
|
// proxy configures
|
|
pxyCfgs map[string]config.ProxyConf
|
|
|
|
// proxies
|
|
proxies map[string]Proxy
|
|
|
|
// visitor configures
|
|
visitorCfgs map[string]config.ProxyConf
|
|
|
|
// visitors
|
|
visitors map[string]Visitor
|
|
|
|
// control connection
|
|
conn frpNet.Conn
|
|
|
|
// tcp stream multiplexing, if enabled
|
|
session *smux.Session
|
|
|
|
// 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
|
|
|
|
// 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
|
|
|
|
// last time got the Pong message
|
|
lastPong time.Time
|
|
|
|
mu sync.RWMutex
|
|
|
|
log.Logger
|
|
}
|
|
|
|
func NewControl(svr *Service, pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) *Control {
|
|
loginMsg := &msg.Login{
|
|
Arch: runtime.GOARCH,
|
|
Os: runtime.GOOS,
|
|
PoolCount: config.ClientCommonCfg.PoolCount,
|
|
User: config.ClientCommonCfg.User,
|
|
Version: version.Full(),
|
|
}
|
|
return &Control{
|
|
svr: svr,
|
|
loginMsg: loginMsg,
|
|
pxyCfgs: pxyCfgs,
|
|
visitorCfgs: visitorCfgs,
|
|
proxies: make(map[string]Proxy),
|
|
visitors: make(map[string]Visitor),
|
|
sendCh: make(chan msg.Message, 10),
|
|
readCh: make(chan msg.Message, 10),
|
|
closedCh: make(chan int),
|
|
Logger: log.NewPrefixLogger(""),
|
|
}
|
|
}
|
|
|
|
// 1. login
|
|
// 2. start reader() writer() manager()
|
|
// 3. connection closed
|
|
// 4. In reader(): close closedCh and exit, controler() get it
|
|
// 5. In controler(): close readCh and sendCh, manager() and writer() will exit
|
|
// 6. In controler(): ini readCh, sendCh, closedCh
|
|
// 7. In controler(): start new reader(), writer(), manager()
|
|
// controler() will keep running
|
|
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 config.ClientCommonCfg.LoginFailExit {
|
|
return
|
|
} else {
|
|
time.Sleep(30 * time.Second)
|
|
}
|
|
} else {
|
|
break
|
|
}
|
|
}
|
|
|
|
go ctl.controler()
|
|
go ctl.manager()
|
|
go ctl.writer()
|
|
go ctl.reader()
|
|
|
|
// start all local visitors
|
|
for _, cfg := range ctl.visitorCfgs {
|
|
visitor := NewVisitor(ctl, cfg)
|
|
err = visitor.Run()
|
|
if err != nil {
|
|
visitor.Warn("start error: %v", err)
|
|
continue
|
|
}
|
|
ctl.visitors[cfg.GetName()] = visitor
|
|
visitor.Info("start visitor success")
|
|
}
|
|
|
|
// send NewProxy message for all configured proxies
|
|
for _, cfg := range ctl.pxyCfgs {
|
|
var newProxyMsg msg.NewProxy
|
|
cfg.UnMarshalToMsg(&newProxyMsg)
|
|
ctl.sendCh <- &newProxyMsg
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (ctl *Control) NewWorkConn() {
|
|
workConn, err := ctl.connectServer()
|
|
if err != nil {
|
|
return
|
|
}
|
|
|
|
m := &msg.NewWorkConn{
|
|
RunId: ctl.getRunId(),
|
|
}
|
|
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 {
|
|
ctl.Error("work connection closed, %v", err)
|
|
workConn.Close()
|
|
return
|
|
}
|
|
workConn.AddLogPrefix(startMsg.ProxyName)
|
|
|
|
// dispatch this work connection to related proxy
|
|
pxy, ok := ctl.getProxy(startMsg.ProxyName)
|
|
if ok {
|
|
workConn.Debug("start a new work connection, localAddr: %s remoteAddr: %s", workConn.LocalAddr().String(), workConn.RemoteAddr().String())
|
|
go pxy.InWorkConn(workConn)
|
|
} else {
|
|
workConn.Close()
|
|
}
|
|
}
|
|
|
|
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)
|
|
ctl.closedCh = make(chan int)
|
|
}
|
|
|
|
// 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.ConnectServerByHttpProxy(config.ClientCommonCfg.HttpProxy, config.ClientCommonCfg.Protocol,
|
|
fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
defer func() {
|
|
if err != nil {
|
|
conn.Close()
|
|
}
|
|
}()
|
|
|
|
if config.ClientCommonCfg.TcpMux {
|
|
session, errRet := smux.Client(conn, nil)
|
|
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(config.ClientCommonCfg.PrivilegeToken, now)
|
|
ctl.loginMsg.Timestamp = now
|
|
ctl.loginMsg.RunId = ctl.getRunId()
|
|
|
|
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.setRunId(loginRespMsg.RunId)
|
|
config.ClientCommonCfg.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)
|
|
|
|
// login success, so we let closedCh available again
|
|
ctl.closedCh = make(chan int)
|
|
ctl.lastPong = time.Now()
|
|
|
|
return nil
|
|
}
|
|
|
|
func (ctl *Control) connectServer() (conn frpNet.Conn, err error) {
|
|
if config.ClientCommonCfg.TcpMux {
|
|
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.ConnectServerByHttpProxy(config.ClientCommonCfg.HttpProxy, config.ClientCommonCfg.Protocol,
|
|
fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
|
|
if err != nil {
|
|
ctl.Warn("start new connection to server error: %v", err)
|
|
return
|
|
}
|
|
}
|
|
return
|
|
}
|
|
|
|
func (ctl *Control) reader() {
|
|
defer func() {
|
|
if err := recover(); err != nil {
|
|
ctl.Error("panic error: %v", err)
|
|
}
|
|
}()
|
|
defer close(ctl.closedCh)
|
|
|
|
encReader := crypto.NewReader(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken))
|
|
for {
|
|
if m, err := msg.ReadMsg(encReader); err != nil {
|
|
if err == io.EOF {
|
|
ctl.Debug("read from control connection EOF")
|
|
return
|
|
} else {
|
|
ctl.Warn("read error: %v", err)
|
|
return
|
|
}
|
|
} else {
|
|
ctl.readCh <- m
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ctl *Control) writer() {
|
|
encWriter, err := crypto.NewWriter(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken))
|
|
if err != nil {
|
|
ctl.conn.Error("crypto new writer error: %v", err)
|
|
ctl.conn.Close()
|
|
return
|
|
}
|
|
for {
|
|
if m, ok := <-ctl.sendCh; !ok {
|
|
ctl.Info("control writer is closing")
|
|
return
|
|
} else {
|
|
if err := msg.WriteMsg(encWriter, m); err != nil {
|
|
ctl.Warn("write message to control connection error: %v", err)
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// manager handles all channel events and do corresponding process
|
|
func (ctl *Control) manager() {
|
|
defer func() {
|
|
if err := recover(); err != nil {
|
|
ctl.Error("panic error: %v", err)
|
|
}
|
|
}()
|
|
|
|
hbSend := time.NewTicker(time.Duration(config.ClientCommonCfg.HeartBeatInterval) * time.Second)
|
|
defer hbSend.Stop()
|
|
hbCheck := time.NewTicker(time.Second)
|
|
defer hbCheck.Stop()
|
|
|
|
for {
|
|
select {
|
|
case <-hbSend.C:
|
|
// send heartbeat to server
|
|
ctl.Debug("send heartbeat to server")
|
|
ctl.sendCh <- &msg.Ping{}
|
|
case <-hbCheck.C:
|
|
if time.Since(ctl.lastPong) > time.Duration(config.ClientCommonCfg.HeartBeatTimeout)*time.Second {
|
|
ctl.Warn("heartbeat timeout")
|
|
// let reader() stop
|
|
ctl.conn.Close()
|
|
return
|
|
}
|
|
case rawMsg, ok := <-ctl.readCh:
|
|
if !ok {
|
|
return
|
|
}
|
|
|
|
switch m := rawMsg.(type) {
|
|
case *msg.ReqWorkConn:
|
|
go ctl.NewWorkConn()
|
|
case *msg.NewProxyResp:
|
|
// Server will return NewProxyResp message to each NewProxy message.
|
|
// Start a new proxy handler if no error got
|
|
if m.Error != "" {
|
|
ctl.Warn("[%s] start error: %s", m.ProxyName, m.Error)
|
|
continue
|
|
}
|
|
cfg, ok := ctl.getProxyConf(m.ProxyName)
|
|
if !ok {
|
|
// it will never go to this branch now
|
|
ctl.Warn("[%s] no proxy conf found", m.ProxyName)
|
|
continue
|
|
}
|
|
|
|
oldPxy, ok := ctl.getProxy(m.ProxyName)
|
|
if ok {
|
|
oldPxy.Close()
|
|
}
|
|
pxy := NewProxy(ctl, cfg)
|
|
if err := pxy.Run(); err != nil {
|
|
ctl.Warn("[%s] proxy start running error: %v", m.ProxyName, err)
|
|
ctl.sendCh <- &msg.CloseProxy{
|
|
ProxyName: m.ProxyName,
|
|
}
|
|
continue
|
|
}
|
|
ctl.addProxy(m.ProxyName, pxy)
|
|
ctl.Info("[%s] start proxy success", m.ProxyName)
|
|
case *msg.Pong:
|
|
ctl.lastPong = time.Now()
|
|
ctl.Debug("receive heartbeat from server")
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// controler keep watching closedCh, start a new connection if previous control connection is closed.
|
|
// If controler is notified by closedCh, reader and writer and manager will exit, then recall these functions.
|
|
func (ctl *Control) controler() {
|
|
var err error
|
|
maxDelayTime := 30 * time.Second
|
|
delayTime := time.Second
|
|
|
|
checkInterval := 10 * time.Second
|
|
checkProxyTicker := time.NewTicker(checkInterval)
|
|
for {
|
|
select {
|
|
case <-checkProxyTicker.C:
|
|
// Every 10 seconds, check which proxy registered failed and reregister it to server.
|
|
ctl.mu.RLock()
|
|
for _, cfg := range ctl.pxyCfgs {
|
|
if _, exist := ctl.proxies[cfg.GetName()]; !exist {
|
|
ctl.Info("try to register proxy [%s]", cfg.GetName())
|
|
var newProxyMsg msg.NewProxy
|
|
cfg.UnMarshalToMsg(&newProxyMsg)
|
|
ctl.sendCh <- &newProxyMsg
|
|
}
|
|
}
|
|
|
|
for _, cfg := range ctl.visitorCfgs {
|
|
if _, exist := ctl.visitors[cfg.GetName()]; !exist {
|
|
ctl.Info("try to start visitor [%s]", cfg.GetName())
|
|
visitor := NewVisitor(ctl, cfg)
|
|
err = visitor.Run()
|
|
if err != nil {
|
|
visitor.Warn("start error: %v", err)
|
|
continue
|
|
}
|
|
ctl.visitors[cfg.GetName()] = visitor
|
|
visitor.Info("start visitor success")
|
|
}
|
|
}
|
|
ctl.mu.RUnlock()
|
|
case _, ok := <-ctl.closedCh:
|
|
// we won't get any variable from this channel
|
|
if !ok {
|
|
// close related channels
|
|
close(ctl.readCh)
|
|
close(ctl.sendCh)
|
|
|
|
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
|
|
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 the delayTime
|
|
delayTime = time.Second
|
|
break
|
|
}
|
|
|
|
// init related channels and variables
|
|
ctl.init()
|
|
|
|
// previous work goroutines should be closed and start them here
|
|
go ctl.manager()
|
|
go ctl.writer()
|
|
go ctl.reader()
|
|
|
|
// send NewProxy message for all configured proxies
|
|
ctl.mu.RLock()
|
|
for _, cfg := range ctl.pxyCfgs {
|
|
var newProxyMsg msg.NewProxy
|
|
cfg.UnMarshalToMsg(&newProxyMsg)
|
|
ctl.sendCh <- &newProxyMsg
|
|
}
|
|
ctl.mu.RUnlock()
|
|
|
|
checkProxyTicker.Stop()
|
|
checkProxyTicker = time.NewTicker(checkInterval)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ctl *Control) setRunId(runId string) {
|
|
ctl.mu.Lock()
|
|
defer ctl.mu.Unlock()
|
|
ctl.runId = runId
|
|
}
|
|
|
|
func (ctl *Control) getRunId() string {
|
|
ctl.mu.RLock()
|
|
defer ctl.mu.RUnlock()
|
|
return ctl.runId
|
|
}
|
|
|
|
func (ctl *Control) getProxy(name string) (pxy Proxy, ok bool) {
|
|
ctl.mu.RLock()
|
|
defer ctl.mu.RUnlock()
|
|
pxy, ok = ctl.proxies[name]
|
|
return
|
|
}
|
|
|
|
func (ctl *Control) addProxy(name string, pxy Proxy) {
|
|
ctl.mu.Lock()
|
|
defer ctl.mu.Unlock()
|
|
ctl.proxies[name] = pxy
|
|
}
|
|
|
|
func (ctl *Control) getProxyConf(name string) (conf config.ProxyConf, ok bool) {
|
|
ctl.mu.RLock()
|
|
defer ctl.mu.RUnlock()
|
|
conf, ok = ctl.pxyCfgs[name]
|
|
return
|
|
}
|
|
|
|
func (ctl *Control) reloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) {
|
|
ctl.mu.Lock()
|
|
defer ctl.mu.Unlock()
|
|
|
|
removedPxyNames := make([]string, 0)
|
|
for name, oldCfg := range ctl.pxyCfgs {
|
|
del := false
|
|
cfg, ok := pxyCfgs[name]
|
|
if !ok {
|
|
del = true
|
|
} else {
|
|
if !oldCfg.Compare(cfg) {
|
|
del = true
|
|
}
|
|
}
|
|
|
|
if del {
|
|
removedPxyNames = append(removedPxyNames, name)
|
|
delete(ctl.pxyCfgs, name)
|
|
if pxy, ok := ctl.proxies[name]; ok {
|
|
pxy.Close()
|
|
}
|
|
delete(ctl.proxies, name)
|
|
ctl.sendCh <- &msg.CloseProxy{
|
|
ProxyName: name,
|
|
}
|
|
}
|
|
}
|
|
ctl.Info("proxy removed: %v", removedPxyNames)
|
|
|
|
addedPxyNames := make([]string, 0)
|
|
for name, cfg := range pxyCfgs {
|
|
if _, ok := ctl.pxyCfgs[name]; !ok {
|
|
ctl.pxyCfgs[name] = cfg
|
|
addedPxyNames = append(addedPxyNames, name)
|
|
}
|
|
}
|
|
ctl.Info("proxy added: %v", addedPxyNames)
|
|
|
|
removedVisitorName := make([]string, 0)
|
|
for name, oldVisitorCfg := range ctl.visitorCfgs {
|
|
del := false
|
|
cfg, ok := visitorCfgs[name]
|
|
if !ok {
|
|
del = true
|
|
} else {
|
|
if !oldVisitorCfg.Compare(cfg) {
|
|
del = true
|
|
}
|
|
}
|
|
|
|
if del {
|
|
removedVisitorName = append(removedVisitorName, name)
|
|
delete(ctl.visitorCfgs, name)
|
|
if visitor, ok := ctl.visitors[name]; ok {
|
|
visitor.Close()
|
|
}
|
|
delete(ctl.visitors, name)
|
|
}
|
|
}
|
|
ctl.Info("visitor removed: %v", removedVisitorName)
|
|
|
|
addedVisitorName := make([]string, 0)
|
|
for name, visitorCfg := range visitorCfgs {
|
|
if _, ok := ctl.visitorCfgs[name]; !ok {
|
|
ctl.visitorCfgs[name] = visitorCfg
|
|
addedVisitorName = append(addedVisitorName, name)
|
|
}
|
|
}
|
|
ctl.Info("visitor added: %v", addedVisitorName)
|
|
}
|