control.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442
  1. // Copyright 2017 fatedier, fatedier@gmail.com
  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 client
  15. import (
  16. "fmt"
  17. "io"
  18. "runtime"
  19. "sync"
  20. "time"
  21. "github.com/fatedier/frp/models/config"
  22. "github.com/fatedier/frp/models/msg"
  23. "github.com/fatedier/frp/utils/crypto"
  24. "github.com/fatedier/frp/utils/log"
  25. frpNet "github.com/fatedier/frp/utils/net"
  26. "github.com/fatedier/frp/utils/shutdown"
  27. "github.com/fatedier/frp/utils/util"
  28. "github.com/fatedier/frp/utils/version"
  29. "github.com/xtaci/smux"
  30. )
  31. const (
  32. connReadTimeout time.Duration = 10 * time.Second
  33. )
  34. type Control struct {
  35. // frpc service
  36. svr *Service
  37. // login message to server, only used
  38. loginMsg *msg.Login
  39. pm *ProxyManager
  40. // control connection
  41. conn frpNet.Conn
  42. // tcp stream multiplexing, if enabled
  43. session *smux.Session
  44. // put a message in this channel to send it over control connection to server
  45. sendCh chan (msg.Message)
  46. // read from this channel to get the next message sent by server
  47. readCh chan (msg.Message)
  48. // run id got from server
  49. runId string
  50. // if we call close() in control, do not reconnect to server
  51. exit bool
  52. // goroutines can block by reading from this channel, it will be closed only in reader() when control connection is closed
  53. closedCh chan int
  54. // last time got the Pong message
  55. lastPong time.Time
  56. readerShutdown *shutdown.Shutdown
  57. writerShutdown *shutdown.Shutdown
  58. msgHandlerShutdown *shutdown.Shutdown
  59. mu sync.RWMutex
  60. log.Logger
  61. }
  62. func NewControl(svr *Service, pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) *Control {
  63. loginMsg := &msg.Login{
  64. Arch: runtime.GOARCH,
  65. Os: runtime.GOOS,
  66. PoolCount: config.ClientCommonCfg.PoolCount,
  67. User: config.ClientCommonCfg.User,
  68. Version: version.Full(),
  69. }
  70. ctl := &Control{
  71. svr: svr,
  72. loginMsg: loginMsg,
  73. sendCh: make(chan msg.Message, 100),
  74. readCh: make(chan msg.Message, 100),
  75. closedCh: make(chan int),
  76. readerShutdown: shutdown.New(),
  77. writerShutdown: shutdown.New(),
  78. msgHandlerShutdown: shutdown.New(),
  79. Logger: log.NewPrefixLogger(""),
  80. }
  81. ctl.pm = NewProxyManager(ctl, ctl.sendCh, "")
  82. ctl.pm.Reload(pxyCfgs, visitorCfgs, false)
  83. return ctl
  84. }
  85. func (ctl *Control) Run() (err error) {
  86. for {
  87. err = ctl.login()
  88. if err != nil {
  89. ctl.Warn("login to server failed: %v", err)
  90. // if login_fail_exit is true, just exit this program
  91. // otherwise sleep a while and continues relogin to server
  92. if config.ClientCommonCfg.LoginFailExit {
  93. return
  94. } else {
  95. time.Sleep(10 * time.Second)
  96. }
  97. } else {
  98. break
  99. }
  100. }
  101. go ctl.worker()
  102. // start all local visitors and send NewProxy message for all configured proxies
  103. ctl.pm.Reset(ctl.sendCh, ctl.runId)
  104. ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew})
  105. return nil
  106. }
  107. func (ctl *Control) HandleReqWorkConn(inMsg *msg.ReqWorkConn) {
  108. workConn, err := ctl.connectServer()
  109. if err != nil {
  110. return
  111. }
  112. m := &msg.NewWorkConn{
  113. RunId: ctl.runId,
  114. }
  115. if err = msg.WriteMsg(workConn, m); err != nil {
  116. ctl.Warn("work connection write to server error: %v", err)
  117. workConn.Close()
  118. return
  119. }
  120. var startMsg msg.StartWorkConn
  121. if err = msg.ReadMsgInto(workConn, &startMsg); err != nil {
  122. ctl.Error("work connection closed, %v", err)
  123. workConn.Close()
  124. return
  125. }
  126. workConn.AddLogPrefix(startMsg.ProxyName)
  127. // dispatch this work connection to related proxy
  128. ctl.pm.HandleWorkConn(startMsg.ProxyName, workConn)
  129. }
  130. func (ctl *Control) HandleNewProxyResp(inMsg *msg.NewProxyResp) {
  131. // Server will return NewProxyResp message to each NewProxy message.
  132. // Start a new proxy handler if no error got
  133. err := ctl.pm.StartProxy(inMsg.ProxyName, inMsg.RemoteAddr, inMsg.Error)
  134. if err != nil {
  135. ctl.Warn("[%s] start error: %v", inMsg.ProxyName, err)
  136. } else {
  137. ctl.Info("[%s] start proxy success", inMsg.ProxyName)
  138. }
  139. }
  140. func (ctl *Control) Close() error {
  141. ctl.mu.Lock()
  142. defer ctl.mu.Unlock()
  143. ctl.exit = true
  144. ctl.pm.CloseProxies()
  145. return nil
  146. }
  147. // login send a login message to server and wait for a loginResp message.
  148. func (ctl *Control) login() (err error) {
  149. if ctl.conn != nil {
  150. ctl.conn.Close()
  151. }
  152. if ctl.session != nil {
  153. ctl.session.Close()
  154. }
  155. conn, err := frpNet.ConnectServerByHttpProxy(config.ClientCommonCfg.HttpProxy, config.ClientCommonCfg.Protocol,
  156. fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
  157. if err != nil {
  158. return err
  159. }
  160. defer func() {
  161. if err != nil {
  162. conn.Close()
  163. }
  164. }()
  165. if config.ClientCommonCfg.TcpMux {
  166. session, errRet := smux.Client(conn, nil)
  167. if errRet != nil {
  168. return errRet
  169. }
  170. stream, errRet := session.OpenStream()
  171. if errRet != nil {
  172. session.Close()
  173. return errRet
  174. }
  175. conn = frpNet.WrapConn(stream)
  176. ctl.session = session
  177. }
  178. now := time.Now().Unix()
  179. ctl.loginMsg.PrivilegeKey = util.GetAuthKey(config.ClientCommonCfg.PrivilegeToken, now)
  180. ctl.loginMsg.Timestamp = now
  181. ctl.loginMsg.RunId = ctl.runId
  182. if err = msg.WriteMsg(conn, ctl.loginMsg); err != nil {
  183. return err
  184. }
  185. var loginRespMsg msg.LoginResp
  186. conn.SetReadDeadline(time.Now().Add(connReadTimeout))
  187. if err = msg.ReadMsgInto(conn, &loginRespMsg); err != nil {
  188. return err
  189. }
  190. conn.SetReadDeadline(time.Time{})
  191. if loginRespMsg.Error != "" {
  192. err = fmt.Errorf("%s", loginRespMsg.Error)
  193. ctl.Error("%s", loginRespMsg.Error)
  194. return err
  195. }
  196. ctl.conn = conn
  197. // update runId got from server
  198. ctl.runId = loginRespMsg.RunId
  199. config.ClientCommonCfg.ServerUdpPort = loginRespMsg.ServerUdpPort
  200. ctl.ClearLogPrefix()
  201. ctl.AddLogPrefix(loginRespMsg.RunId)
  202. ctl.Info("login to server success, get run id [%s], server udp port [%d]", loginRespMsg.RunId, loginRespMsg.ServerUdpPort)
  203. return nil
  204. }
  205. func (ctl *Control) connectServer() (conn frpNet.Conn, err error) {
  206. if config.ClientCommonCfg.TcpMux {
  207. stream, errRet := ctl.session.OpenStream()
  208. if errRet != nil {
  209. err = errRet
  210. ctl.Warn("start new connection to server error: %v", err)
  211. return
  212. }
  213. conn = frpNet.WrapConn(stream)
  214. } else {
  215. conn, err = frpNet.ConnectServerByHttpProxy(config.ClientCommonCfg.HttpProxy, config.ClientCommonCfg.Protocol,
  216. fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerPort))
  217. if err != nil {
  218. ctl.Warn("start new connection to server error: %v", err)
  219. return
  220. }
  221. }
  222. return
  223. }
  224. // reader read all messages from frps and send to readCh
  225. func (ctl *Control) reader() {
  226. defer func() {
  227. if err := recover(); err != nil {
  228. ctl.Error("panic error: %v", err)
  229. }
  230. }()
  231. defer ctl.readerShutdown.Done()
  232. defer close(ctl.closedCh)
  233. encReader := crypto.NewReader(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken))
  234. for {
  235. if m, err := msg.ReadMsg(encReader); err != nil {
  236. if err == io.EOF {
  237. ctl.Debug("read from control connection EOF")
  238. return
  239. } else {
  240. ctl.Warn("read error: %v", err)
  241. return
  242. }
  243. } else {
  244. ctl.readCh <- m
  245. }
  246. }
  247. }
  248. // writer writes messages got from sendCh to frps
  249. func (ctl *Control) writer() {
  250. defer ctl.writerShutdown.Done()
  251. encWriter, err := crypto.NewWriter(ctl.conn, []byte(config.ClientCommonCfg.PrivilegeToken))
  252. if err != nil {
  253. ctl.conn.Error("crypto new writer error: %v", err)
  254. ctl.conn.Close()
  255. return
  256. }
  257. for {
  258. if m, ok := <-ctl.sendCh; !ok {
  259. ctl.Info("control writer is closing")
  260. return
  261. } else {
  262. if err := msg.WriteMsg(encWriter, m); err != nil {
  263. ctl.Warn("write message to control connection error: %v", err)
  264. return
  265. }
  266. }
  267. }
  268. }
  269. // msgHandler handles all channel events and do corresponding operations.
  270. func (ctl *Control) msgHandler() {
  271. defer func() {
  272. if err := recover(); err != nil {
  273. ctl.Error("panic error: %v", err)
  274. }
  275. }()
  276. defer ctl.msgHandlerShutdown.Done()
  277. hbSend := time.NewTicker(time.Duration(config.ClientCommonCfg.HeartBeatInterval) * time.Second)
  278. defer hbSend.Stop()
  279. hbCheck := time.NewTicker(time.Second)
  280. defer hbCheck.Stop()
  281. ctl.lastPong = time.Now()
  282. for {
  283. select {
  284. case <-hbSend.C:
  285. // send heartbeat to server
  286. ctl.Debug("send heartbeat to server")
  287. ctl.sendCh <- &msg.Ping{}
  288. case <-hbCheck.C:
  289. if time.Since(ctl.lastPong) > time.Duration(config.ClientCommonCfg.HeartBeatTimeout)*time.Second {
  290. ctl.Warn("heartbeat timeout")
  291. // let reader() stop
  292. ctl.conn.Close()
  293. return
  294. }
  295. case rawMsg, ok := <-ctl.readCh:
  296. if !ok {
  297. return
  298. }
  299. switch m := rawMsg.(type) {
  300. case *msg.ReqWorkConn:
  301. go ctl.HandleReqWorkConn(m)
  302. case *msg.NewProxyResp:
  303. ctl.HandleNewProxyResp(m)
  304. case *msg.Pong:
  305. ctl.lastPong = time.Now()
  306. ctl.Debug("receive heartbeat from server")
  307. }
  308. }
  309. }
  310. }
  311. // controler keep watching closedCh, start a new connection if previous control connection is closed.
  312. // If controler is notified by closedCh, reader and writer and handler will exit, then recall these functions.
  313. func (ctl *Control) worker() {
  314. go ctl.msgHandler()
  315. go ctl.reader()
  316. go ctl.writer()
  317. var err error
  318. maxDelayTime := 20 * time.Second
  319. delayTime := time.Second
  320. checkInterval := 60 * time.Second
  321. checkProxyTicker := time.NewTicker(checkInterval)
  322. for {
  323. select {
  324. case <-checkProxyTicker.C:
  325. // check which proxy registered failed and reregister it to server
  326. ctl.pm.CheckAndStartProxy([]string{ProxyStatusStartErr, ProxyStatusClosed})
  327. case _, ok := <-ctl.closedCh:
  328. // we won't get any variable from this channel
  329. if !ok {
  330. // close related channels and wait until other goroutines done
  331. close(ctl.readCh)
  332. ctl.readerShutdown.WaitDone()
  333. ctl.msgHandlerShutdown.WaitDone()
  334. close(ctl.sendCh)
  335. ctl.writerShutdown.WaitDone()
  336. ctl.pm.CloseProxies()
  337. // if ctl.exit is true, just exit
  338. ctl.mu.RLock()
  339. exit := ctl.exit
  340. ctl.mu.RUnlock()
  341. if exit {
  342. return
  343. }
  344. // loop util reconnecting to server success
  345. for {
  346. ctl.Info("try to reconnect to server...")
  347. err = ctl.login()
  348. if err != nil {
  349. ctl.Warn("reconnect to server error: %v", err)
  350. time.Sleep(delayTime)
  351. delayTime = delayTime * 2
  352. if delayTime > maxDelayTime {
  353. delayTime = maxDelayTime
  354. }
  355. continue
  356. }
  357. // reconnect success, init delayTime
  358. delayTime = time.Second
  359. break
  360. }
  361. // init related channels and variables
  362. ctl.sendCh = make(chan msg.Message, 100)
  363. ctl.readCh = make(chan msg.Message, 100)
  364. ctl.closedCh = make(chan int)
  365. ctl.readerShutdown = shutdown.New()
  366. ctl.writerShutdown = shutdown.New()
  367. ctl.msgHandlerShutdown = shutdown.New()
  368. ctl.pm.Reset(ctl.sendCh, ctl.runId)
  369. // previous work goroutines should be closed and start them here
  370. go ctl.msgHandler()
  371. go ctl.writer()
  372. go ctl.reader()
  373. // start all configured proxies
  374. ctl.pm.CheckAndStartProxy([]string{ProxyStatusNew, ProxyStatusClosed})
  375. checkProxyTicker.Stop()
  376. checkProxyTicker = time.NewTicker(checkInterval)
  377. }
  378. }
  379. }
  380. }
  381. func (ctl *Control) reloadConf(pxyCfgs map[string]config.ProxyConf, visitorCfgs map[string]config.ProxyConf) error {
  382. err := ctl.pm.Reload(pxyCfgs, visitorCfgs, true)
  383. return err
  384. }