frp/client/proxy_manager.go

312 lines
6.3 KiB
Go
Raw Normal View History

2018-01-17 01:09:33 +08:00
package client
import (
"fmt"
"sync"
"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-05-08 02:13:30 +08:00
"github.com/fatedier/golib/errors"
2018-01-17 01:09:33 +08:00
)
const (
2018-10-18 13:55:51 +08:00
ProxyStatusNew = "new"
ProxyStatusStartErr = "start error"
ProxyStatusWaitStart = "wait start"
ProxyStatusRunning = "running"
ProxyStatusCheckFailed = "check failed"
ProxyStatusClosed = "closed"
2018-01-17 01:09:33 +08:00
)
type ProxyManager struct {
sendCh chan (msg.Message)
2018-01-17 01:09:33 +08:00
proxies map[string]*ProxyWrapper
closed bool
mu sync.RWMutex
2018-01-17 01:09:33 +08:00
log.Logger
}
2018-11-06 18:35:05 +08:00
func NewProxyManager(msgSendCh chan (msg.Message), logPrefix string) *ProxyManager {
2018-01-17 01:09:33 +08:00
return &ProxyManager{
proxies: make(map[string]*ProxyWrapper),
sendCh: msgSendCh,
closed: false,
Logger: log.NewPrefixLogger(logPrefix),
2018-01-17 01:09:33 +08:00
}
}
func (pm *ProxyManager) Reset(msgSendCh chan (msg.Message), logPrefix string) {
pm.mu.Lock()
defer pm.mu.Unlock()
pm.closed = false
pm.sendCh = msgSendCh
pm.ClearLogPrefix()
pm.AddLogPrefix(logPrefix)
}
// Must hold the lock before calling this function.
func (pm *ProxyManager) sendMsg(m msg.Message) error {
err := errors.PanicToError(func() {
pm.sendCh <- m
})
if err != nil {
pm.closed = true
}
return err
}
2018-01-17 14:40:08 +08:00
func (pm *ProxyManager) StartProxy(name string, remoteAddr string, serverRespErr string) error {
2018-01-17 01:09:33 +08:00
pm.mu.Lock()
defer pm.mu.Unlock()
if pm.closed {
return fmt.Errorf("ProxyManager is closed now")
}
pxy, ok := pm.proxies[name]
if !ok {
return fmt.Errorf("no proxy found")
}
2018-01-17 14:40:08 +08:00
if err := pxy.Start(remoteAddr, serverRespErr); err != nil {
2018-01-17 01:09:33 +08:00
errRet := err
err = pm.sendMsg(&msg.CloseProxy{
ProxyName: name,
})
if err != nil {
errRet = fmt.Errorf("send CloseProxy message error")
}
return errRet
}
return nil
}
func (pm *ProxyManager) CloseProxies() {
pm.mu.RLock()
defer pm.mu.RUnlock()
for _, pxy := range pm.proxies {
pxy.Close()
}
}
// pxyStatus: check and start proxies in which status
func (pm *ProxyManager) CheckAndStartProxy(pxyStatus []string) {
2018-01-17 01:09:33 +08:00
pm.mu.RLock()
defer pm.mu.RUnlock()
if pm.closed {
pm.Warn("CheckAndStartProxy error: ProxyManager is closed now")
return
}
for _, pxy := range pm.proxies {
status := pxy.GetStatusStr()
for _, s := range pxyStatus {
if status == s {
var newProxyMsg msg.NewProxy
2018-04-10 17:46:49 +08:00
pxy.Cfg.MarshalToMsg(&newProxyMsg)
err := pm.sendMsg(&newProxyMsg)
if err != nil {
pm.Warn("[%s] proxy send NewProxy message error")
return
}
2018-01-29 23:51:46 +08:00
pxy.WaitStart()
break
2018-01-17 01:09:33 +08:00
}
}
}
}
func (pm *ProxyManager) Reload(pxyCfgs map[string]config.ProxyConf, startNow bool) error {
2018-01-17 01:09:33 +08:00
pm.mu.Lock()
defer func() {
pm.mu.Unlock()
if startNow {
go pm.CheckAndStartProxy([]string{ProxyStatusNew})
}
}()
2018-01-17 01:09:33 +08:00
if pm.closed {
err := fmt.Errorf("Reload error: ProxyManager is closed now")
pm.Warn(err.Error())
return err
}
delPxyNames := make([]string, 0)
for name, pxy := range pm.proxies {
del := false
cfg, ok := pxyCfgs[name]
if !ok {
del = true
} else {
if !pxy.Cfg.Compare(cfg) {
del = true
}
}
if del {
delPxyNames = append(delPxyNames, name)
delete(pm.proxies, name)
pxy.Close()
err := pm.sendMsg(&msg.CloseProxy{
ProxyName: name,
})
if err != nil {
err = fmt.Errorf("Reload error: ProxyManager is closed now")
pm.Warn(err.Error())
return err
}
}
}
pm.Info("proxy removed: %v", delPxyNames)
addPxyNames := make([]string, 0)
for name, cfg := range pxyCfgs {
if _, ok := pm.proxies[name]; !ok {
pxy := NewProxyWrapper(cfg)
pm.proxies[name] = pxy
addPxyNames = append(addPxyNames, name)
}
}
pm.Info("proxy added: %v", addPxyNames)
return nil
}
func (pm *ProxyManager) HandleWorkConn(name string, workConn frpNet.Conn) {
pm.mu.RLock()
pw, ok := pm.proxies[name]
pm.mu.RUnlock()
if ok {
pw.InWorkConn(workConn)
} else {
workConn.Close()
}
}
func (pm *ProxyManager) GetAllProxyStatus() []*ProxyStatus {
ps := make([]*ProxyStatus, 0)
pm.mu.RLock()
defer pm.mu.RUnlock()
for _, pxy := range pm.proxies {
ps = append(ps, pxy.GetStatus())
}
return ps
}
2018-11-06 18:35:05 +08:00
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"`
}
// ProxyWrapper is a wrapper of Proxy interface only used in ProxyManager
// Add additional proxy status info
type ProxyWrapper struct {
Name string
Type string
Status string
Err string
Cfg config.ProxyConf
RemoteAddr string
pxy Proxy
mu sync.RWMutex
}
func NewProxyWrapper(cfg config.ProxyConf) *ProxyWrapper {
return &ProxyWrapper{
Name: cfg.GetBaseInfo().ProxyName,
Type: cfg.GetBaseInfo().ProxyType,
Status: ProxyStatusNew,
Cfg: cfg,
pxy: nil,
}
}
func (pw *ProxyWrapper) GetStatusStr() string {
pw.mu.RLock()
defer pw.mu.RUnlock()
return pw.Status
}
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
}
func (pw *ProxyWrapper) WaitStart() {
pw.mu.Lock()
defer pw.mu.Unlock()
pw.Status = ProxyStatusWaitStart
}
func (pw *ProxyWrapper) Start(remoteAddr string, serverRespErr string) error {
if pw.pxy != nil {
pw.pxy.Close()
pw.pxy = nil
}
if serverRespErr != "" {
pw.mu.Lock()
pw.Status = ProxyStatusStartErr
pw.RemoteAddr = remoteAddr
pw.Err = serverRespErr
pw.mu.Unlock()
return fmt.Errorf(serverRespErr)
}
pxy := NewProxy(pw.Cfg)
pw.mu.Lock()
defer pw.mu.Unlock()
pw.RemoteAddr = remoteAddr
if err := pxy.Run(); err != nil {
pw.Status = ProxyStatusStartErr
pw.Err = err.Error()
return err
}
pw.Status = ProxyStatusRunning
pw.Err = ""
pw.pxy = pxy
return nil
}
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) Close() {
pw.mu.Lock()
defer pw.mu.Unlock()
if pw.pxy != nil {
pw.pxy.Close()
pw.pxy = nil
}
pw.Status = ProxyStatusClosed
}