frp/client/proxy_manager.go

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"

	"github.com/fatedier/golib/errors"
)

const (
	ProxyStatusNew         = "new"
	ProxyStatusStartErr    = "start error"
	ProxyStatusWaitStart   = "wait start"
	ProxyStatusRunning     = "running"
	ProxyStatusCheckFailed = "check failed"
	ProxyStatusClosed      = "closed"
)

type ProxyManager struct {
	sendCh  chan (msg.Message)
	proxies map[string]*ProxyWrapper
	closed  bool
	mu      sync.RWMutex

	log.Logger
}

func NewProxyManager(msgSendCh chan (msg.Message), logPrefix string) *ProxyManager {
	return &ProxyManager{
		proxies: make(map[string]*ProxyWrapper),
		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, remoteAddr 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(remoteAddr, 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()
	}
}

// pxyStatus: check and start proxies in which status
func (pm *ProxyManager) CheckAndStartProxy(pxyStatus []string) {
	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
				pxy.Cfg.MarshalToMsg(&newProxyMsg)
				err := pm.sendMsg(&newProxyMsg)
				if err != nil {
					pm.Warn("[%s] proxy send NewProxy message error")
					return
				}
				pxy.WaitStart()
				break
			}
		}
	}
}

func (pm *ProxyManager) Reload(pxyCfgs map[string]config.ProxyConf, startNow bool) error {
	pm.mu.Lock()
	defer func() {
		pm.mu.Unlock()
		if startNow {
			go pm.CheckAndStartProxy([]string{ProxyStatusNew})
		}
	}()
	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
}

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
}