proxy_manager.go 7.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349
  1. package client
  2. import (
  3. "fmt"
  4. "sync"
  5. "github.com/fatedier/frp/models/config"
  6. "github.com/fatedier/frp/models/msg"
  7. "github.com/fatedier/frp/utils/errors"
  8. "github.com/fatedier/frp/utils/log"
  9. frpNet "github.com/fatedier/frp/utils/net"
  10. )
  11. const (
  12. ProxyStatusNew = "new"
  13. ProxyStatusStartErr = "start error"
  14. ProxyStatusRunning = "running"
  15. ProxyStatusClosed = "closed"
  16. )
  17. type ProxyManager struct {
  18. ctl *Control
  19. proxies map[string]*ProxyWrapper
  20. visitorCfgs map[string]config.ProxyConf
  21. visitors map[string]Visitor
  22. sendCh chan (msg.Message)
  23. closed bool
  24. mu sync.RWMutex
  25. log.Logger
  26. }
  27. type ProxyWrapper struct {
  28. Name string
  29. Type string
  30. Status string
  31. Err string
  32. Cfg config.ProxyConf
  33. RemoteAddr string
  34. pxy Proxy
  35. mu sync.RWMutex
  36. }
  37. type ProxyStatus struct {
  38. Name string `json:"name"`
  39. Type string `json:"type"`
  40. Status string `json:"status"`
  41. Err string `json:"err"`
  42. Cfg config.ProxyConf `json:"cfg"`
  43. // Got from server.
  44. RemoteAddr string `json:"remote_addr"`
  45. }
  46. func NewProxyWrapper(cfg config.ProxyConf) *ProxyWrapper {
  47. return &ProxyWrapper{
  48. Name: cfg.GetName(),
  49. Type: cfg.GetType(),
  50. Status: ProxyStatusNew,
  51. Cfg: cfg,
  52. pxy: nil,
  53. }
  54. }
  55. func (pw *ProxyWrapper) IsRunning() bool {
  56. pw.mu.RLock()
  57. defer pw.mu.RUnlock()
  58. if pw.Status == ProxyStatusRunning {
  59. return true
  60. } else {
  61. return false
  62. }
  63. }
  64. func (pw *ProxyWrapper) GetStatus() *ProxyStatus {
  65. pw.mu.RLock()
  66. defer pw.mu.RUnlock()
  67. ps := &ProxyStatus{
  68. Name: pw.Name,
  69. Type: pw.Type,
  70. Status: pw.Status,
  71. Err: pw.Err,
  72. Cfg: pw.Cfg,
  73. RemoteAddr: pw.RemoteAddr,
  74. }
  75. return ps
  76. }
  77. func (pw *ProxyWrapper) Start(remoteAddr string, serverRespErr string) error {
  78. if pw.pxy != nil {
  79. pw.pxy.Close()
  80. pw.pxy = nil
  81. }
  82. if serverRespErr != "" {
  83. pw.mu.Lock()
  84. pw.Status = ProxyStatusStartErr
  85. pw.RemoteAddr = remoteAddr
  86. pw.Err = serverRespErr
  87. pw.mu.Unlock()
  88. return fmt.Errorf(serverRespErr)
  89. }
  90. pxy := NewProxy(pw.Cfg)
  91. pw.mu.Lock()
  92. defer pw.mu.Unlock()
  93. pw.RemoteAddr = remoteAddr
  94. if err := pxy.Run(); err != nil {
  95. pw.Status = ProxyStatusStartErr
  96. pw.Err = err.Error()
  97. return err
  98. }
  99. pw.Status = ProxyStatusRunning
  100. pw.Err = ""
  101. pw.pxy = pxy
  102. return nil
  103. }
  104. func (pw *ProxyWrapper) InWorkConn(workConn frpNet.Conn) {
  105. pw.mu.RLock()
  106. pxy := pw.pxy
  107. pw.mu.RUnlock()
  108. if pxy != nil {
  109. workConn.Debug("start a new work connection, localAddr: %s remoteAddr: %s", workConn.LocalAddr().String(), workConn.RemoteAddr().String())
  110. go pxy.InWorkConn(workConn)
  111. } else {
  112. workConn.Close()
  113. }
  114. }
  115. func (pw *ProxyWrapper) Close() {
  116. pw.mu.Lock()
  117. defer pw.mu.Unlock()
  118. if pw.pxy != nil {
  119. pw.pxy.Close()
  120. pw.pxy = nil
  121. }
  122. pw.Status = ProxyStatusClosed
  123. }
  124. func NewProxyManager(ctl *Control, msgSendCh chan (msg.Message), logPrefix string) *ProxyManager {
  125. return &ProxyManager{
  126. ctl: ctl,
  127. proxies: make(map[string]*ProxyWrapper),
  128. visitorCfgs: make(map[string]config.ProxyConf),
  129. visitors: make(map[string]Visitor),
  130. sendCh: msgSendCh,
  131. closed: false,
  132. Logger: log.NewPrefixLogger(logPrefix),
  133. }
  134. }
  135. func (pm *ProxyManager) Reset(msgSendCh chan (msg.Message), logPrefix string) {
  136. pm.mu.Lock()
  137. defer pm.mu.Unlock()
  138. pm.closed = false
  139. pm.sendCh = msgSendCh
  140. pm.ClearLogPrefix()
  141. pm.AddLogPrefix(logPrefix)
  142. }
  143. // Must hold the lock before calling this function.
  144. func (pm *ProxyManager) sendMsg(m msg.Message) error {
  145. err := errors.PanicToError(func() {
  146. pm.sendCh <- m
  147. })
  148. if err != nil {
  149. pm.closed = true
  150. }
  151. return err
  152. }
  153. func (pm *ProxyManager) StartProxy(name string, remoteAddr string, serverRespErr string) error {
  154. pm.mu.Lock()
  155. defer pm.mu.Unlock()
  156. if pm.closed {
  157. return fmt.Errorf("ProxyManager is closed now")
  158. }
  159. pxy, ok := pm.proxies[name]
  160. if !ok {
  161. return fmt.Errorf("no proxy found")
  162. }
  163. if err := pxy.Start(remoteAddr, serverRespErr); err != nil {
  164. errRet := err
  165. err = pm.sendMsg(&msg.CloseProxy{
  166. ProxyName: name,
  167. })
  168. if err != nil {
  169. errRet = fmt.Errorf("send CloseProxy message error")
  170. }
  171. return errRet
  172. }
  173. return nil
  174. }
  175. func (pm *ProxyManager) CloseProxies() {
  176. pm.mu.RLock()
  177. defer pm.mu.RUnlock()
  178. for _, pxy := range pm.proxies {
  179. pxy.Close()
  180. }
  181. }
  182. func (pm *ProxyManager) CheckAndStartProxy() {
  183. pm.mu.RLock()
  184. defer pm.mu.RUnlock()
  185. if pm.closed {
  186. pm.Warn("CheckAndStartProxy error: ProxyManager is closed now")
  187. return
  188. }
  189. for _, pxy := range pm.proxies {
  190. if !pxy.IsRunning() {
  191. var newProxyMsg msg.NewProxy
  192. pxy.Cfg.UnMarshalToMsg(&newProxyMsg)
  193. err := pm.sendMsg(&newProxyMsg)
  194. if err != nil {
  195. pm.Warn("[%s] proxy send NewProxy message error")
  196. return
  197. }
  198. }
  199. }
  200. for _, cfg := range pm.visitorCfgs {
  201. if _, exist := pm.visitors[cfg.GetName()]; !exist {
  202. pm.Info("try to start visitor [%s]", cfg.GetName())
  203. visitor := NewVisitor(pm.ctl, cfg)
  204. err := visitor.Run()
  205. if err != nil {
  206. visitor.Warn("start error: %v", err)
  207. continue
  208. }
  209. pm.visitors[cfg.GetName()] = visitor
  210. visitor.Info("start visitor success")
  211. }
  212. }
  213. }
  214. func (pm *ProxyManager) Reload(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) error {
  215. pm.mu.Lock()
  216. defer pm.mu.Unlock()
  217. if pm.closed {
  218. err := fmt.Errorf("Reload error: ProxyManager is closed now")
  219. pm.Warn(err.Error())
  220. return err
  221. }
  222. delPxyNames := make([]string, 0)
  223. for name, pxy := range pm.proxies {
  224. del := false
  225. cfg, ok := pxyCfgs[name]
  226. if !ok {
  227. del = true
  228. } else {
  229. if !pxy.Cfg.Compare(cfg) {
  230. del = true
  231. }
  232. }
  233. if del {
  234. delPxyNames = append(delPxyNames, name)
  235. delete(pm.proxies, name)
  236. pxy.Close()
  237. err := pm.sendMsg(&msg.CloseProxy{
  238. ProxyName: name,
  239. })
  240. if err != nil {
  241. err = fmt.Errorf("Reload error: ProxyManager is closed now")
  242. pm.Warn(err.Error())
  243. return err
  244. }
  245. }
  246. }
  247. pm.Info("proxy removed: %v", delPxyNames)
  248. addPxyNames := make([]string, 0)
  249. for name, cfg := range pxyCfgs {
  250. if _, ok := pm.proxies[name]; !ok {
  251. pxy := NewProxyWrapper(cfg)
  252. pm.proxies[name] = pxy
  253. addPxyNames = append(addPxyNames, name)
  254. }
  255. }
  256. pm.Info("proxy added: %v", addPxyNames)
  257. delVisitorName := make([]string, 0)
  258. for name, oldVisitorCfg := range pm.visitorCfgs {
  259. del := false
  260. cfg, ok := visitorCfgs[name]
  261. if !ok {
  262. del = true
  263. } else {
  264. if !oldVisitorCfg.Compare(cfg) {
  265. del = true
  266. }
  267. }
  268. if del {
  269. delVisitorName = append(delVisitorName, name)
  270. delete(pm.visitorCfgs, name)
  271. if visitor, ok := pm.visitors[name]; ok {
  272. visitor.Close()
  273. }
  274. delete(pm.visitors, name)
  275. }
  276. }
  277. pm.Info("visitor removed: %v", delVisitorName)
  278. addVisitorName := make([]string, 0)
  279. for name, visitorCfg := range visitorCfgs {
  280. if _, ok := pm.visitorCfgs[name]; !ok {
  281. pm.visitorCfgs[name] = visitorCfg
  282. addVisitorName = append(addVisitorName, name)
  283. }
  284. }
  285. pm.Info("visitor added: %v", addVisitorName)
  286. return nil
  287. }
  288. func (pm *ProxyManager) HandleWorkConn(name string, workConn frpNet.Conn) {
  289. pm.mu.RLock()
  290. pw, ok := pm.proxies[name]
  291. pm.mu.RUnlock()
  292. if ok {
  293. pw.InWorkConn(workConn)
  294. } else {
  295. workConn.Close()
  296. }
  297. }
  298. func (pm *ProxyManager) GetAllProxyStatus() []*ProxyStatus {
  299. ps := make([]*ProxyStatus, 0)
  300. pm.mu.RLock()
  301. defer pm.mu.RUnlock()
  302. for _, pxy := range pm.proxies {
  303. ps = append(ps, pxy.GetStatus())
  304. }
  305. return ps
  306. }