proxy_wrapper.go 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284
  1. // Copyright 2023 The frp Authors
  2. //
  3. // Licensed under the Apache License, Version 2.0 (the "License");
  4. // you may not use this file except in compliance with the License.
  5. // You may obtain a copy of the License at
  6. //
  7. // http://www.apache.org/licenses/LICENSE-2.0
  8. //
  9. // Unless required by applicable law or agreed to in writing, software
  10. // distributed under the License is distributed on an "AS IS" BASIS,
  11. // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. // See the License for the specific language governing permissions and
  13. // limitations under the License.
  14. package proxy
  15. import (
  16. "context"
  17. "fmt"
  18. "net"
  19. "strconv"
  20. "sync"
  21. "sync/atomic"
  22. "time"
  23. "github.com/fatedier/golib/errors"
  24. "github.com/fatedier/frp/client/event"
  25. "github.com/fatedier/frp/client/health"
  26. v1 "github.com/fatedier/frp/pkg/config/v1"
  27. "github.com/fatedier/frp/pkg/msg"
  28. "github.com/fatedier/frp/pkg/transport"
  29. "github.com/fatedier/frp/pkg/util/xlog"
  30. "github.com/fatedier/frp/pkg/vnet"
  31. )
  32. const (
  33. ProxyPhaseNew = "new"
  34. ProxyPhaseWaitStart = "wait start"
  35. ProxyPhaseStartErr = "start error"
  36. ProxyPhaseRunning = "running"
  37. ProxyPhaseCheckFailed = "check failed"
  38. ProxyPhaseClosed = "closed"
  39. )
  40. var (
  41. statusCheckInterval = 3 * time.Second
  42. waitResponseTimeout = 20 * time.Second
  43. startErrTimeout = 30 * time.Second
  44. )
  45. type WorkingStatus struct {
  46. Name string `json:"name"`
  47. Type string `json:"type"`
  48. Phase string `json:"status"`
  49. Err string `json:"err"`
  50. Cfg v1.ProxyConfigurer `json:"cfg"`
  51. // Got from server.
  52. RemoteAddr string `json:"remote_addr"`
  53. }
  54. type Wrapper struct {
  55. WorkingStatus
  56. // underlying proxy
  57. pxy Proxy
  58. // if ProxyConf has healcheck config
  59. // monitor will watch if it is alive
  60. monitor *health.Monitor
  61. // event handler
  62. handler event.Handler
  63. msgTransporter transport.MessageTransporter
  64. // vnet controller
  65. vnetController *vnet.Controller
  66. health uint32
  67. lastSendStartMsg time.Time
  68. lastStartErr time.Time
  69. closeCh chan struct{}
  70. healthNotifyCh chan struct{}
  71. mu sync.RWMutex
  72. xl *xlog.Logger
  73. ctx context.Context
  74. }
  75. func NewWrapper(
  76. ctx context.Context,
  77. cfg v1.ProxyConfigurer,
  78. clientCfg *v1.ClientCommonConfig,
  79. eventHandler event.Handler,
  80. msgTransporter transport.MessageTransporter,
  81. vnetController *vnet.Controller,
  82. ) *Wrapper {
  83. baseInfo := cfg.GetBaseConfig()
  84. xl := xlog.FromContextSafe(ctx).Spawn().AppendPrefix(baseInfo.Name)
  85. pw := &Wrapper{
  86. WorkingStatus: WorkingStatus{
  87. Name: baseInfo.Name,
  88. Type: baseInfo.Type,
  89. Phase: ProxyPhaseNew,
  90. Cfg: cfg,
  91. },
  92. closeCh: make(chan struct{}),
  93. healthNotifyCh: make(chan struct{}),
  94. handler: eventHandler,
  95. msgTransporter: msgTransporter,
  96. vnetController: vnetController,
  97. xl: xl,
  98. ctx: xlog.NewContext(ctx, xl),
  99. }
  100. if baseInfo.HealthCheck.Type != "" && baseInfo.LocalPort > 0 {
  101. pw.health = 1 // means failed
  102. addr := net.JoinHostPort(baseInfo.LocalIP, strconv.Itoa(baseInfo.LocalPort))
  103. pw.monitor = health.NewMonitor(pw.ctx, baseInfo.HealthCheck, addr,
  104. pw.statusNormalCallback, pw.statusFailedCallback)
  105. xl.Tracef("enable health check monitor")
  106. }
  107. pw.pxy = NewProxy(pw.ctx, pw.Cfg, clientCfg, pw.msgTransporter, pw.vnetController)
  108. return pw
  109. }
  110. func (pw *Wrapper) SetInWorkConnCallback(cb func(*v1.ProxyBaseConfig, net.Conn, *msg.StartWorkConn) bool) {
  111. pw.pxy.SetInWorkConnCallback(cb)
  112. }
  113. func (pw *Wrapper) SetRunningStatus(remoteAddr string, respErr string) error {
  114. pw.mu.Lock()
  115. defer pw.mu.Unlock()
  116. if pw.Phase != ProxyPhaseWaitStart {
  117. return fmt.Errorf("status not wait start, ignore start message")
  118. }
  119. pw.RemoteAddr = remoteAddr
  120. if respErr != "" {
  121. pw.Phase = ProxyPhaseStartErr
  122. pw.Err = respErr
  123. pw.lastStartErr = time.Now()
  124. return fmt.Errorf("%s", pw.Err)
  125. }
  126. if err := pw.pxy.Run(); err != nil {
  127. pw.close()
  128. pw.Phase = ProxyPhaseStartErr
  129. pw.Err = err.Error()
  130. pw.lastStartErr = time.Now()
  131. return err
  132. }
  133. pw.Phase = ProxyPhaseRunning
  134. pw.Err = ""
  135. return nil
  136. }
  137. func (pw *Wrapper) Start() {
  138. go pw.checkWorker()
  139. if pw.monitor != nil {
  140. go pw.monitor.Start()
  141. }
  142. }
  143. func (pw *Wrapper) Stop() {
  144. pw.mu.Lock()
  145. defer pw.mu.Unlock()
  146. close(pw.closeCh)
  147. close(pw.healthNotifyCh)
  148. pw.pxy.Close()
  149. if pw.monitor != nil {
  150. pw.monitor.Stop()
  151. }
  152. pw.Phase = ProxyPhaseClosed
  153. pw.close()
  154. }
  155. func (pw *Wrapper) close() {
  156. _ = pw.handler(&event.CloseProxyPayload{
  157. CloseProxyMsg: &msg.CloseProxy{
  158. ProxyName: pw.Name,
  159. },
  160. })
  161. }
  162. func (pw *Wrapper) checkWorker() {
  163. xl := pw.xl
  164. if pw.monitor != nil {
  165. // let monitor do check request first
  166. time.Sleep(500 * time.Millisecond)
  167. }
  168. for {
  169. // check proxy status
  170. now := time.Now()
  171. if atomic.LoadUint32(&pw.health) == 0 {
  172. pw.mu.Lock()
  173. if pw.Phase == ProxyPhaseNew ||
  174. pw.Phase == ProxyPhaseCheckFailed ||
  175. (pw.Phase == ProxyPhaseWaitStart && now.After(pw.lastSendStartMsg.Add(waitResponseTimeout))) ||
  176. (pw.Phase == ProxyPhaseStartErr && now.After(pw.lastStartErr.Add(startErrTimeout))) {
  177. xl.Tracef("change status from [%s] to [%s]", pw.Phase, ProxyPhaseWaitStart)
  178. pw.Phase = ProxyPhaseWaitStart
  179. var newProxyMsg msg.NewProxy
  180. pw.Cfg.MarshalToMsg(&newProxyMsg)
  181. pw.lastSendStartMsg = now
  182. _ = pw.handler(&event.StartProxyPayload{
  183. NewProxyMsg: &newProxyMsg,
  184. })
  185. }
  186. pw.mu.Unlock()
  187. } else {
  188. pw.mu.Lock()
  189. if pw.Phase == ProxyPhaseRunning || pw.Phase == ProxyPhaseWaitStart {
  190. pw.close()
  191. xl.Tracef("change status from [%s] to [%s]", pw.Phase, ProxyPhaseCheckFailed)
  192. pw.Phase = ProxyPhaseCheckFailed
  193. }
  194. pw.mu.Unlock()
  195. }
  196. select {
  197. case <-pw.closeCh:
  198. return
  199. case <-time.After(statusCheckInterval):
  200. case <-pw.healthNotifyCh:
  201. }
  202. }
  203. }
  204. func (pw *Wrapper) statusNormalCallback() {
  205. xl := pw.xl
  206. atomic.StoreUint32(&pw.health, 0)
  207. _ = errors.PanicToError(func() {
  208. select {
  209. case pw.healthNotifyCh <- struct{}{}:
  210. default:
  211. }
  212. })
  213. xl.Infof("health check success")
  214. }
  215. func (pw *Wrapper) statusFailedCallback() {
  216. xl := pw.xl
  217. atomic.StoreUint32(&pw.health, 1)
  218. _ = errors.PanicToError(func() {
  219. select {
  220. case pw.healthNotifyCh <- struct{}{}:
  221. default:
  222. }
  223. })
  224. xl.Infof("health check failed")
  225. }
  226. func (pw *Wrapper) InWorkConn(workConn net.Conn, m *msg.StartWorkConn) {
  227. xl := pw.xl
  228. pw.mu.RLock()
  229. pxy := pw.pxy
  230. pw.mu.RUnlock()
  231. if pxy != nil && pw.Phase == ProxyPhaseRunning {
  232. xl.Debugf("start a new work connection, localAddr: %s remoteAddr: %s", workConn.LocalAddr().String(), workConn.RemoteAddr().String())
  233. go pxy.InWorkConn(workConn, m)
  234. } else {
  235. workConn.Close()
  236. }
  237. }
  238. func (pw *Wrapper) GetStatus() *WorkingStatus {
  239. pw.mu.RLock()
  240. defer pw.mu.RUnlock()
  241. ps := &WorkingStatus{
  242. Name: pw.Name,
  243. Type: pw.Type,
  244. Phase: pw.Phase,
  245. Err: pw.Err,
  246. Cfg: pw.Cfg,
  247. RemoteAddr: pw.RemoteAddr,
  248. }
  249. return ps
  250. }