frp/client/proxy/proxy_manager.go

140 lines
2.8 KiB
Go
Raw Normal View History

2018-12-09 22:06:22 +08:00
package proxy
2018-01-17 01:09:33 +08:00
import (
"fmt"
"sync"
2018-12-09 22:06:22 +08:00
"github.com/fatedier/frp/client/event"
2018-01-17 01:09:33 +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-05-08 02:13:30 +08:00
"github.com/fatedier/golib/errors"
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
2018-12-07 17:05:36 +08:00
closed bool
mu sync.RWMutex
logPrefix string
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{
2018-12-07 17:05:36 +08:00
proxies: make(map[string]*ProxyWrapper),
sendCh: msgSendCh,
closed: false,
logPrefix: logPrefix,
Logger: log.NewPrefixLogger(logPrefix),
2018-01-17 01:09:33 +08:00
}
}
2018-01-17 14:40:08 +08:00
func (pm *ProxyManager) StartProxy(name string, remoteAddr string, serverRespErr string) error {
2018-12-07 17:05:36 +08:00
pm.mu.RLock()
2018-01-17 01:09:33 +08:00
pxy, ok := pm.proxies[name]
2018-12-07 17:05:36 +08:00
pm.mu.RUnlock()
2018-01-17 01:09:33 +08:00
if !ok {
2018-12-07 17:05:36 +08:00
return fmt.Errorf("proxy [%s] not found", name)
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
err := pxy.SetRunningStatus(remoteAddr, serverRespErr)
if err != nil {
return err
2018-01-17 01:09:33 +08:00
}
return nil
}
2018-12-07 17:05:36 +08:00
func (pm *ProxyManager) Close() {
2018-12-11 15:06:54 +08:00
pm.mu.Lock()
defer pm.mu.Unlock()
2018-01-17 01:09:33 +08:00
for _, pxy := range pm.proxies {
2018-12-07 17:05:36 +08:00
pxy.Stop()
2018-01-17 01:09:33 +08:00
}
2018-12-11 15:06:54 +08:00
pm.proxies = make(map[string]*ProxyWrapper)
2018-01-17 01:09:33 +08:00
}
2019-03-29 19:01:18 +08:00
func (pm *ProxyManager) HandleWorkConn(name string, workConn frpNet.Conn, m *msg.StartWorkConn) {
2018-01-17 01:09:33 +08:00
pm.mu.RLock()
2018-12-07 17:05:36 +08:00
pw, ok := pm.proxies[name]
pm.mu.RUnlock()
if ok {
2019-03-29 19:01:18 +08:00
pw.InWorkConn(workConn, m)
2018-12-07 17:05:36 +08:00
} else {
workConn.Close()
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
}
2018-12-09 22:06:22 +08:00
func (pm *ProxyManager) HandleEvent(evType event.EventType, payload interface{}) error {
2018-12-07 17:05:36 +08:00
var m msg.Message
2018-12-09 22:06:22 +08:00
switch e := payload.(type) {
case *event.StartProxyPayload:
m = e.NewProxyMsg
case *event.CloseProxyPayload:
m = e.CloseProxyMsg
2018-12-07 17:05:36 +08:00
default:
2018-12-09 22:06:22 +08:00
return event.ErrPayloadType
2018-12-07 17:05:36 +08:00
}
err := errors.PanicToError(func() {
pm.sendCh <- m
})
return err
}
2018-01-17 01:09:33 +08:00
2018-12-07 17:05:36 +08:00
func (pm *ProxyManager) GetAllProxyStatus() []*ProxyStatus {
ps := make([]*ProxyStatus, 0)
pm.mu.RLock()
defer pm.mu.RUnlock()
2018-01-17 01:09:33 +08:00
for _, pxy := range pm.proxies {
2018-12-07 17:05:36 +08:00
ps = append(ps, pxy.GetStatus())
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
return ps
2018-01-17 01:09:33 +08:00
}
2018-12-07 17:05:36 +08:00
func (pm *ProxyManager) Reload(pxyCfgs map[string]config.ProxyConf) {
2018-01-17 01:09:33 +08:00
pm.mu.Lock()
2018-12-07 17:05:36 +08:00
defer pm.mu.Unlock()
2018-01-17 01:09:33 +08:00
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)
2018-12-07 17:05:36 +08:00
pxy.Stop()
2018-01-17 01:09:33 +08:00
}
}
2018-12-07 17:05:36 +08:00
if len(delPxyNames) > 0 {
pm.Info("proxy removed: %v", delPxyNames)
}
2018-01-17 01:09:33 +08:00
addPxyNames := make([]string, 0)
for name, cfg := range pxyCfgs {
if _, ok := pm.proxies[name]; !ok {
2018-12-07 17:05:36 +08:00
pxy := NewProxyWrapper(cfg, pm.HandleEvent, pm.logPrefix)
2018-01-17 01:09:33 +08:00
pm.proxies[name] = pxy
addPxyNames = append(addPxyNames, name)
2018-11-06 18:35:05 +08:00
2018-12-07 17:05:36 +08:00
pxy.Start()
}
2018-11-06 18:35:05 +08:00
}
2018-12-07 17:05:36 +08:00
if len(addPxyNames) > 0 {
pm.Info("proxy added: %v", addPxyNames)
2018-11-06 18:35:05 +08:00
}
}