frp/client/proxy/proxy_wrapper.go

245 lines
5.7 KiB
Go
Raw Normal View History

2018-12-09 22:06:22 +08:00
package proxy
2018-12-07 17:05:36 +08:00
import (
"fmt"
"sync"
"sync/atomic"
"time"
2018-12-09 22:06:22 +08:00
"github.com/fatedier/frp/client/event"
"github.com/fatedier/frp/client/health"
2018-12-07 17:05:36 +08:00
"github.com/fatedier/frp/models/config"
"github.com/fatedier/frp/models/msg"
"github.com/fatedier/frp/utils/log"
frpNet "github.com/fatedier/frp/utils/net"
2018-12-09 21:56:46 +08:00
"github.com/fatedier/golib/errors"
2018-12-07 17:05:36 +08:00
)
const (
ProxyStatusNew = "new"
ProxyStatusWaitStart = "wait start"
ProxyStatusStartErr = "start error"
ProxyStatusRunning = "running"
ProxyStatusCheckFailed = "check failed"
ProxyStatusClosed = "closed"
)
var (
statusCheckInterval time.Duration = 3 * time.Second
waitResponseTimeout = 20 * time.Second
startErrTimeout = 30 * time.Second
)
type ProxyStatus struct {
Name string `json:"name"`
Type string `json:"type"`
Status string `json:"status"`
Err string `json:"err"`
Cfg config.ProxyConf `json:"cfg"`
// Got from server.
RemoteAddr string `json:"remote_addr"`
}
type ProxyWrapper struct {
ProxyStatus
// underlying proxy
pxy Proxy
// if ProxyConf has healcheck config
// monitor will watch if it is alive
2018-12-09 22:06:22 +08:00
monitor *health.HealthCheckMonitor
2018-12-07 17:05:36 +08:00
// event handler
2018-12-09 22:06:22 +08:00
handler event.EventHandler
2018-12-07 17:05:36 +08:00
health uint32
lastSendStartMsg time.Time
lastStartErr time.Time
closeCh chan struct{}
2018-12-09 21:56:46 +08:00
healthNotifyCh chan struct{}
2018-12-07 17:05:36 +08:00
mu sync.RWMutex
log.Logger
}
2018-12-09 22:06:22 +08:00
func NewProxyWrapper(cfg config.ProxyConf, eventHandler event.EventHandler, logPrefix string) *ProxyWrapper {
2018-12-07 17:05:36 +08:00
baseInfo := cfg.GetBaseInfo()
pw := &ProxyWrapper{
ProxyStatus: ProxyStatus{
Name: baseInfo.ProxyName,
Type: baseInfo.ProxyType,
Status: ProxyStatusNew,
Cfg: cfg,
},
2018-12-09 21:56:46 +08:00
closeCh: make(chan struct{}),
healthNotifyCh: make(chan struct{}),
handler: eventHandler,
Logger: log.NewPrefixLogger(logPrefix),
2018-12-07 17:05:36 +08:00
}
pw.AddLogPrefix(pw.Name)
if baseInfo.HealthCheckType != "" {
pw.health = 1 // means failed
2018-12-09 22:06:22 +08:00
pw.monitor = health.NewHealthCheckMonitor(baseInfo.HealthCheckType, baseInfo.HealthCheckIntervalS,
2018-12-07 17:05:36 +08:00
baseInfo.HealthCheckTimeoutS, baseInfo.HealthCheckMaxFailed, baseInfo.HealthCheckAddr,
baseInfo.HealthCheckUrl, pw.statusNormalCallback, pw.statusFailedCallback)
pw.monitor.SetLogger(pw.Logger)
pw.Trace("enable health check monitor")
}
pw.pxy = NewProxy(pw.Cfg)
return pw
}
func (pw *ProxyWrapper) SetRunningStatus(remoteAddr string, respErr string) error {
pw.mu.Lock()
defer pw.mu.Unlock()
if pw.Status != ProxyStatusWaitStart {
return fmt.Errorf("status not wait start, ignore start message")
}
pw.RemoteAddr = remoteAddr
if respErr != "" {
pw.Status = ProxyStatusStartErr
pw.Err = respErr
pw.lastStartErr = time.Now()
return fmt.Errorf(pw.Err)
}
if err := pw.pxy.Run(); err != nil {
pw.Status = ProxyStatusStartErr
pw.Err = err.Error()
pw.lastStartErr = time.Now()
return err
}
pw.Status = ProxyStatusRunning
pw.Err = ""
return nil
}
func (pw *ProxyWrapper) Start() {
go pw.checkWorker()
if pw.monitor != nil {
go pw.monitor.Start()
}
}
func (pw *ProxyWrapper) Stop() {
pw.mu.Lock()
defer pw.mu.Unlock()
2018-12-09 21:56:46 +08:00
close(pw.closeCh)
close(pw.healthNotifyCh)
2018-12-07 17:05:36 +08:00
pw.pxy.Close()
if pw.monitor != nil {
pw.monitor.Stop()
}
pw.Status = ProxyStatusClosed
2018-12-09 22:06:22 +08:00
pw.handler(event.EvCloseProxy, &event.CloseProxyPayload{
2018-12-07 17:05:36 +08:00
CloseProxyMsg: &msg.CloseProxy{
ProxyName: pw.Name,
},
})
}
func (pw *ProxyWrapper) checkWorker() {
2018-12-09 21:56:46 +08:00
if pw.monitor != nil {
// let monitor do check request first
time.Sleep(500 * time.Millisecond)
}
2018-12-07 17:05:36 +08:00
for {
// check proxy status
now := time.Now()
if atomic.LoadUint32(&pw.health) == 0 {
pw.mu.Lock()
if pw.Status == ProxyStatusNew ||
pw.Status == ProxyStatusCheckFailed ||
(pw.Status == ProxyStatusWaitStart && now.After(pw.lastSendStartMsg.Add(waitResponseTimeout))) ||
(pw.Status == ProxyStatusStartErr && now.After(pw.lastStartErr.Add(startErrTimeout))) {
pw.Trace("change status from [%s] to [%s]", pw.Status, ProxyStatusWaitStart)
pw.Status = ProxyStatusWaitStart
var newProxyMsg msg.NewProxy
pw.Cfg.MarshalToMsg(&newProxyMsg)
pw.lastSendStartMsg = now
2018-12-09 22:06:22 +08:00
pw.handler(event.EvStartProxy, &event.StartProxyPayload{
2018-12-07 17:05:36 +08:00
NewProxyMsg: &newProxyMsg,
})
}
pw.mu.Unlock()
} else {
pw.mu.Lock()
if pw.Status == ProxyStatusRunning || pw.Status == ProxyStatusWaitStart {
2018-12-09 22:06:22 +08:00
pw.handler(event.EvCloseProxy, &event.CloseProxyPayload{
2018-12-07 17:05:36 +08:00
CloseProxyMsg: &msg.CloseProxy{
ProxyName: pw.Name,
},
})
pw.Trace("change status from [%s] to [%s]", pw.Status, ProxyStatusCheckFailed)
pw.Status = ProxyStatusCheckFailed
}
pw.mu.Unlock()
}
select {
case <-pw.closeCh:
return
case <-time.After(statusCheckInterval):
2018-12-09 21:56:46 +08:00
case <-pw.healthNotifyCh:
2018-12-07 17:05:36 +08:00
}
}
}
func (pw *ProxyWrapper) statusNormalCallback() {
atomic.StoreUint32(&pw.health, 0)
2018-12-09 21:56:46 +08:00
errors.PanicToError(func() {
select {
case pw.healthNotifyCh <- struct{}{}:
default:
}
})
2018-12-07 17:05:36 +08:00
pw.Info("health check success")
}
func (pw *ProxyWrapper) statusFailedCallback() {
atomic.StoreUint32(&pw.health, 1)
2018-12-09 21:56:46 +08:00
errors.PanicToError(func() {
select {
case pw.healthNotifyCh <- struct{}{}:
default:
}
})
2018-12-07 17:05:36 +08:00
pw.Info("health check failed")
}
func (pw *ProxyWrapper) InWorkConn(workConn frpNet.Conn) {
pw.mu.RLock()
pxy := pw.pxy
pw.mu.RUnlock()
if pxy != nil {
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 (pw *ProxyWrapper) GetStatus() *ProxyStatus {
pw.mu.RLock()
defer pw.mu.RUnlock()
ps := &ProxyStatus{
Name: pw.Name,
Type: pw.Type,
Status: pw.Status,
Err: pw.Err,
Cfg: pw.Cfg,
RemoteAddr: pw.RemoteAddr,
}
return ps
}