package client import ( "fmt" "sync" "github.com/fatedier/frp/models/config" "github.com/fatedier/frp/models/msg" "github.com/fatedier/frp/utils/errors" "github.com/fatedier/frp/utils/log" frpNet "github.com/fatedier/frp/utils/net" ) const ( ProxyStatusNew = "new" ProxyStatusStartErr = "start error" ProxyStatusRunning = "running" ProxyStatusClosed = "closed" ) type ProxyManager struct { ctl *Control proxies map[string]*ProxyWrapper visitorCfgs map[string]config.ProxyConf visitors map[string]Visitor sendCh chan (msg.Message) closed bool mu sync.RWMutex log.Logger } type ProxyWrapper struct { Name string Type string Status string Err string Cfg config.ProxyConf pxy Proxy mu sync.RWMutex } type ProxyStatus struct { Name string `json:"name"` Type string `json:"type"` Status string `json:"status"` Err string `json:"err"` Cfg config.ProxyConf `json:"cfg"` } func NewProxyWrapper(cfg config.ProxyConf) *ProxyWrapper { return &ProxyWrapper{ Name: cfg.GetName(), Type: cfg.GetType(), Status: ProxyStatusNew, Cfg: cfg, pxy: nil, } } func (pw *ProxyWrapper) IsRunning() bool { pw.mu.RLock() defer pw.mu.RUnlock() if pw.Status == ProxyStatusRunning { return true } else { return false } } 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, } return ps } func (pw *ProxyWrapper) Start(serverRespErr string) error { if pw.pxy != nil { pw.pxy.Close() pw.pxy = nil } if serverRespErr != "" { pw.mu.Lock() pw.Status = ProxyStatusStartErr pw.Err = serverRespErr pw.mu.Unlock() return fmt.Errorf(serverRespErr) } pxy := NewProxy(pw.Cfg) pw.mu.Lock() defer pw.mu.Unlock() 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 } func NewProxyManager(ctl *Control, msgSendCh chan (msg.Message), logPrefix string) *ProxyManager { return &ProxyManager{ proxies: make(map[string]*ProxyWrapper), visitorCfgs: make(map[string]config.ProxyConf), visitors: make(map[string]Visitor), sendCh: msgSendCh, closed: false, Logger: log.NewPrefixLogger(logPrefix), } } 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 } func (pm *ProxyManager) StartProxy(name string, serverRespErr string) error { 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") } if err := pxy.Start(serverRespErr); err != nil { 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() } } func (pm *ProxyManager) CheckAndStartProxy() { pm.mu.RLock() defer pm.mu.RUnlock() if pm.closed { pm.Warn("CheckAndStartProxy error: ProxyManager is closed now") return } for _, pxy := range pm.proxies { if !pxy.IsRunning() { var newProxyMsg msg.NewProxy pxy.Cfg.UnMarshalToMsg(&newProxyMsg) err := pm.sendMsg(&newProxyMsg) if err != nil { pm.Warn("[%s] proxy send NewProxy message error") return } } } for _, cfg := range pm.visitorCfgs { if _, exist := pm.visitors[cfg.GetName()]; !exist { pm.Info("try to start visitor [%s]", cfg.GetName()) visitor := NewVisitor(pm.ctl, cfg) err := visitor.Run() if err != nil { visitor.Warn("start error: %v", err) continue } pm.visitors[cfg.GetName()] = visitor visitor.Info("start visitor success") } } } func (pm *ProxyManager) Reload(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) error { pm.mu.Lock() defer pm.mu.Unlock() 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) delVisitorName := make([]string, 0) for name, oldVisitorCfg := range pm.visitorCfgs { del := false cfg, ok := visitorCfgs[name] if !ok { del = true } else { if !oldVisitorCfg.Compare(cfg) { del = true } } if del { delVisitorName = append(delVisitorName, name) delete(pm.visitorCfgs, name) if visitor, ok := pm.visitors[name]; ok { visitor.Close() } delete(pm.visitors, name) } } pm.Info("visitor removed: %v", delVisitorName) addVisitorName := make([]string, 0) for name, visitorCfg := range visitorCfgs { if _, ok := pm.visitorCfgs[name]; !ok { pm.visitorCfgs[name] = visitorCfg addVisitorName = append(addVisitorName, name) } } pm.Info("visitor added: %v", addVisitorName) 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 }