1
0

control.go 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555
  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 server
  15. import (
  16. "context"
  17. "fmt"
  18. "io"
  19. "net"
  20. "runtime/debug"
  21. "sync"
  22. "time"
  23. "github.com/fatedier/frp/models/auth"
  24. "github.com/fatedier/frp/models/config"
  25. "github.com/fatedier/frp/models/consts"
  26. frpErr "github.com/fatedier/frp/models/errors"
  27. "github.com/fatedier/frp/models/msg"
  28. plugin "github.com/fatedier/frp/models/plugin/server"
  29. "github.com/fatedier/frp/server/controller"
  30. "github.com/fatedier/frp/server/metrics"
  31. "github.com/fatedier/frp/server/proxy"
  32. "github.com/fatedier/frp/utils/util"
  33. "github.com/fatedier/frp/utils/version"
  34. "github.com/fatedier/frp/utils/xlog"
  35. "github.com/fatedier/golib/control/shutdown"
  36. "github.com/fatedier/golib/crypto"
  37. "github.com/fatedier/golib/errors"
  38. )
  39. type ControlManager struct {
  40. // controls indexed by run id
  41. ctlsByRunId map[string]*Control
  42. mu sync.RWMutex
  43. }
  44. func NewControlManager() *ControlManager {
  45. return &ControlManager{
  46. ctlsByRunId: make(map[string]*Control),
  47. }
  48. }
  49. func (cm *ControlManager) Add(runId string, ctl *Control) (oldCtl *Control) {
  50. cm.mu.Lock()
  51. defer cm.mu.Unlock()
  52. oldCtl, ok := cm.ctlsByRunId[runId]
  53. if ok {
  54. oldCtl.Replaced(ctl)
  55. }
  56. cm.ctlsByRunId[runId] = ctl
  57. return
  58. }
  59. // we should make sure if it's the same control to prevent delete a new one
  60. func (cm *ControlManager) Del(runId string, ctl *Control) {
  61. cm.mu.Lock()
  62. defer cm.mu.Unlock()
  63. if c, ok := cm.ctlsByRunId[runId]; ok && c == ctl {
  64. delete(cm.ctlsByRunId, runId)
  65. }
  66. }
  67. func (cm *ControlManager) GetById(runId string) (ctl *Control, ok bool) {
  68. cm.mu.RLock()
  69. defer cm.mu.RUnlock()
  70. ctl, ok = cm.ctlsByRunId[runId]
  71. return
  72. }
  73. type Control struct {
  74. // all resource managers and controllers
  75. rc *controller.ResourceController
  76. // proxy manager
  77. pxyManager *proxy.ProxyManager
  78. // plugin manager
  79. pluginManager *plugin.Manager
  80. // verifies authentication based on selected method
  81. authVerifier auth.Verifier
  82. // login message
  83. loginMsg *msg.Login
  84. // control connection
  85. conn net.Conn
  86. // put a message in this channel to send it over control connection to client
  87. sendCh chan (msg.Message)
  88. // read from this channel to get the next message sent by client
  89. readCh chan (msg.Message)
  90. // work connections
  91. workConnCh chan net.Conn
  92. // proxies in one client
  93. proxies map[string]proxy.Proxy
  94. // pool count
  95. poolCount int
  96. // ports used, for limitations
  97. portsUsedNum int
  98. // last time got the Ping message
  99. lastPing time.Time
  100. // A new run id will be generated when a new client login.
  101. // If run id got from login message has same run id, it means it's the same client, so we can
  102. // replace old controller instantly.
  103. runId string
  104. // control status
  105. status string
  106. readerShutdown *shutdown.Shutdown
  107. writerShutdown *shutdown.Shutdown
  108. managerShutdown *shutdown.Shutdown
  109. allShutdown *shutdown.Shutdown
  110. mu sync.RWMutex
  111. // Server configuration information
  112. serverCfg config.ServerCommonConf
  113. xl *xlog.Logger
  114. ctx context.Context
  115. }
  116. func NewControl(
  117. ctx context.Context,
  118. rc *controller.ResourceController,
  119. pxyManager *proxy.ProxyManager,
  120. pluginManager *plugin.Manager,
  121. authVerifier auth.Verifier,
  122. ctlConn net.Conn,
  123. loginMsg *msg.Login,
  124. serverCfg config.ServerCommonConf,
  125. ) *Control {
  126. poolCount := loginMsg.PoolCount
  127. if poolCount > int(serverCfg.MaxPoolCount) {
  128. poolCount = int(serverCfg.MaxPoolCount)
  129. }
  130. return &Control{
  131. rc: rc,
  132. pxyManager: pxyManager,
  133. pluginManager: pluginManager,
  134. authVerifier: authVerifier,
  135. conn: ctlConn,
  136. loginMsg: loginMsg,
  137. sendCh: make(chan msg.Message, 10),
  138. readCh: make(chan msg.Message, 10),
  139. workConnCh: make(chan net.Conn, poolCount+10),
  140. proxies: make(map[string]proxy.Proxy),
  141. poolCount: poolCount,
  142. portsUsedNum: 0,
  143. lastPing: time.Now(),
  144. runId: loginMsg.RunId,
  145. status: consts.Working,
  146. readerShutdown: shutdown.New(),
  147. writerShutdown: shutdown.New(),
  148. managerShutdown: shutdown.New(),
  149. allShutdown: shutdown.New(),
  150. serverCfg: serverCfg,
  151. xl: xlog.FromContextSafe(ctx),
  152. ctx: ctx,
  153. }
  154. }
  155. // Start send a login success message to client and start working.
  156. func (ctl *Control) Start() {
  157. loginRespMsg := &msg.LoginResp{
  158. Version: version.Full(),
  159. RunId: ctl.runId,
  160. ServerUdpPort: ctl.serverCfg.BindUdpPort,
  161. Error: "",
  162. }
  163. msg.WriteMsg(ctl.conn, loginRespMsg)
  164. go ctl.writer()
  165. for i := 0; i < ctl.poolCount; i++ {
  166. ctl.sendCh <- &msg.ReqWorkConn{}
  167. }
  168. go ctl.manager()
  169. go ctl.reader()
  170. go ctl.stoper()
  171. }
  172. func (ctl *Control) RegisterWorkConn(conn net.Conn) error {
  173. xl := ctl.xl
  174. defer func() {
  175. if err := recover(); err != nil {
  176. xl.Error("panic error: %v", err)
  177. xl.Error(string(debug.Stack()))
  178. }
  179. }()
  180. select {
  181. case ctl.workConnCh <- conn:
  182. xl.Debug("new work connection registered")
  183. return nil
  184. default:
  185. xl.Debug("work connection pool is full, discarding")
  186. return fmt.Errorf("work connection pool is full, discarding")
  187. }
  188. }
  189. // When frps get one user connection, we get one work connection from the pool and return it.
  190. // If no workConn available in the pool, send message to frpc to get one or more
  191. // and wait until it is available.
  192. // return an error if wait timeout
  193. func (ctl *Control) GetWorkConn() (workConn net.Conn, err error) {
  194. xl := ctl.xl
  195. defer func() {
  196. if err := recover(); err != nil {
  197. xl.Error("panic error: %v", err)
  198. xl.Error(string(debug.Stack()))
  199. }
  200. }()
  201. var ok bool
  202. // get a work connection from the pool
  203. select {
  204. case workConn, ok = <-ctl.workConnCh:
  205. if !ok {
  206. err = frpErr.ErrCtlClosed
  207. return
  208. }
  209. xl.Debug("get work connection from pool")
  210. default:
  211. // no work connections available in the poll, send message to frpc to get more
  212. err = errors.PanicToError(func() {
  213. ctl.sendCh <- &msg.ReqWorkConn{}
  214. })
  215. if err != nil {
  216. xl.Error("%v", err)
  217. return
  218. }
  219. select {
  220. case workConn, ok = <-ctl.workConnCh:
  221. if !ok {
  222. err = frpErr.ErrCtlClosed
  223. xl.Warn("no work connections avaiable, %v", err)
  224. return
  225. }
  226. case <-time.After(time.Duration(ctl.serverCfg.UserConnTimeout) * time.Second):
  227. err = fmt.Errorf("timeout trying to get work connection")
  228. xl.Warn("%v", err)
  229. return
  230. }
  231. }
  232. // When we get a work connection from pool, replace it with a new one.
  233. errors.PanicToError(func() {
  234. ctl.sendCh <- &msg.ReqWorkConn{}
  235. })
  236. return
  237. }
  238. func (ctl *Control) Replaced(newCtl *Control) {
  239. xl := ctl.xl
  240. xl.Info("Replaced by client [%s]", newCtl.runId)
  241. ctl.runId = ""
  242. ctl.allShutdown.Start()
  243. }
  244. func (ctl *Control) writer() {
  245. xl := ctl.xl
  246. defer func() {
  247. if err := recover(); err != nil {
  248. xl.Error("panic error: %v", err)
  249. xl.Error(string(debug.Stack()))
  250. }
  251. }()
  252. defer ctl.allShutdown.Start()
  253. defer ctl.writerShutdown.Done()
  254. encWriter, err := crypto.NewWriter(ctl.conn, []byte(ctl.serverCfg.Token))
  255. if err != nil {
  256. xl.Error("crypto new writer error: %v", err)
  257. ctl.allShutdown.Start()
  258. return
  259. }
  260. for {
  261. if m, ok := <-ctl.sendCh; !ok {
  262. xl.Info("control writer is closing")
  263. return
  264. } else {
  265. if err := msg.WriteMsg(encWriter, m); err != nil {
  266. xl.Warn("write message to control connection error: %v", err)
  267. return
  268. }
  269. }
  270. }
  271. }
  272. func (ctl *Control) reader() {
  273. xl := ctl.xl
  274. defer func() {
  275. if err := recover(); err != nil {
  276. xl.Error("panic error: %v", err)
  277. xl.Error(string(debug.Stack()))
  278. }
  279. }()
  280. defer ctl.allShutdown.Start()
  281. defer ctl.readerShutdown.Done()
  282. encReader := crypto.NewReader(ctl.conn, []byte(ctl.serverCfg.Token))
  283. for {
  284. if m, err := msg.ReadMsg(encReader); err != nil {
  285. if err == io.EOF {
  286. xl.Debug("control connection closed")
  287. return
  288. } else {
  289. xl.Warn("read error: %v", err)
  290. ctl.conn.Close()
  291. return
  292. }
  293. } else {
  294. ctl.readCh <- m
  295. }
  296. }
  297. }
  298. func (ctl *Control) stoper() {
  299. xl := ctl.xl
  300. defer func() {
  301. if err := recover(); err != nil {
  302. xl.Error("panic error: %v", err)
  303. xl.Error(string(debug.Stack()))
  304. }
  305. }()
  306. ctl.allShutdown.WaitStart()
  307. close(ctl.readCh)
  308. ctl.managerShutdown.WaitDone()
  309. close(ctl.sendCh)
  310. ctl.writerShutdown.WaitDone()
  311. ctl.conn.Close()
  312. ctl.readerShutdown.WaitDone()
  313. ctl.mu.Lock()
  314. defer ctl.mu.Unlock()
  315. close(ctl.workConnCh)
  316. for workConn := range ctl.workConnCh {
  317. workConn.Close()
  318. }
  319. for _, pxy := range ctl.proxies {
  320. pxy.Close()
  321. ctl.pxyManager.Del(pxy.GetName())
  322. metrics.Server.CloseProxy(pxy.GetName(), pxy.GetConf().GetBaseInfo().ProxyType)
  323. }
  324. ctl.allShutdown.Done()
  325. xl.Info("client exit success")
  326. metrics.Server.CloseClient()
  327. }
  328. // block until Control closed
  329. func (ctl *Control) WaitClosed() {
  330. ctl.allShutdown.WaitDone()
  331. }
  332. func (ctl *Control) manager() {
  333. xl := ctl.xl
  334. defer func() {
  335. if err := recover(); err != nil {
  336. xl.Error("panic error: %v", err)
  337. xl.Error(string(debug.Stack()))
  338. }
  339. }()
  340. defer ctl.allShutdown.Start()
  341. defer ctl.managerShutdown.Done()
  342. heartbeat := time.NewTicker(time.Second)
  343. defer heartbeat.Stop()
  344. for {
  345. select {
  346. case <-heartbeat.C:
  347. if time.Since(ctl.lastPing) > time.Duration(ctl.serverCfg.HeartBeatTimeout)*time.Second {
  348. xl.Warn("heartbeat timeout")
  349. return
  350. }
  351. case rawMsg, ok := <-ctl.readCh:
  352. if !ok {
  353. return
  354. }
  355. switch m := rawMsg.(type) {
  356. case *msg.NewProxy:
  357. content := &plugin.NewProxyContent{
  358. User: plugin.UserInfo{
  359. User: ctl.loginMsg.User,
  360. Metas: ctl.loginMsg.Metas,
  361. RunId: ctl.loginMsg.RunId,
  362. },
  363. NewProxy: *m,
  364. }
  365. var remoteAddr string
  366. retContent, err := ctl.pluginManager.NewProxy(content)
  367. if err == nil {
  368. m = &retContent.NewProxy
  369. remoteAddr, err = ctl.RegisterProxy(m)
  370. }
  371. // register proxy in this control
  372. resp := &msg.NewProxyResp{
  373. ProxyName: m.ProxyName,
  374. }
  375. if err != nil {
  376. xl.Warn("new proxy [%s] error: %v", m.ProxyName, err)
  377. resp.Error = util.GenerateResponseErrorString(fmt.Sprintf("new proxy [%s] error", m.ProxyName), err, ctl.serverCfg.DetailedErrorsToClient)
  378. } else {
  379. resp.RemoteAddr = remoteAddr
  380. xl.Info("new proxy [%s] success", m.ProxyName)
  381. metrics.Server.NewProxy(m.ProxyName, m.ProxyType)
  382. }
  383. ctl.sendCh <- resp
  384. case *msg.CloseProxy:
  385. ctl.CloseProxy(m)
  386. xl.Info("close proxy [%s] success", m.ProxyName)
  387. case *msg.Ping:
  388. content := &plugin.PingContent{
  389. User: plugin.UserInfo{
  390. User: ctl.loginMsg.User,
  391. Metas: ctl.loginMsg.Metas,
  392. RunId: ctl.loginMsg.RunId,
  393. },
  394. Ping: *m,
  395. }
  396. retContent, err := ctl.pluginManager.Ping(content)
  397. if err == nil {
  398. m = &retContent.Ping
  399. err = ctl.authVerifier.VerifyPing(m)
  400. }
  401. if err != nil {
  402. xl.Warn("received invalid ping: %v", err)
  403. ctl.sendCh <- &msg.Pong{
  404. Error: util.GenerateResponseErrorString("invalid ping", err, ctl.serverCfg.DetailedErrorsToClient),
  405. }
  406. return
  407. }
  408. ctl.lastPing = time.Now()
  409. xl.Debug("receive heartbeat")
  410. ctl.sendCh <- &msg.Pong{}
  411. }
  412. }
  413. }
  414. }
  415. func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err error) {
  416. var pxyConf config.ProxyConf
  417. // Load configures from NewProxy message and check.
  418. pxyConf, err = config.NewProxyConfFromMsg(pxyMsg, ctl.serverCfg)
  419. if err != nil {
  420. return
  421. }
  422. // NewProxy will return a interface Proxy.
  423. // In fact it create different proxies by different proxy type, we just call run() here.
  424. pxy, err := proxy.NewProxy(ctl.ctx, ctl.runId, ctl.rc, ctl.poolCount, ctl.GetWorkConn, pxyConf, ctl.serverCfg)
  425. if err != nil {
  426. return remoteAddr, err
  427. }
  428. // Check ports used number in each client
  429. if ctl.serverCfg.MaxPortsPerClient > 0 {
  430. ctl.mu.Lock()
  431. if ctl.portsUsedNum+pxy.GetUsedPortsNum() > int(ctl.serverCfg.MaxPortsPerClient) {
  432. ctl.mu.Unlock()
  433. err = fmt.Errorf("exceed the max_ports_per_client")
  434. return
  435. }
  436. ctl.portsUsedNum = ctl.portsUsedNum + pxy.GetUsedPortsNum()
  437. ctl.mu.Unlock()
  438. defer func() {
  439. if err != nil {
  440. ctl.mu.Lock()
  441. ctl.portsUsedNum = ctl.portsUsedNum - pxy.GetUsedPortsNum()
  442. ctl.mu.Unlock()
  443. }
  444. }()
  445. }
  446. remoteAddr, err = pxy.Run()
  447. if err != nil {
  448. return
  449. }
  450. defer func() {
  451. if err != nil {
  452. pxy.Close()
  453. }
  454. }()
  455. err = ctl.pxyManager.Add(pxyMsg.ProxyName, pxy)
  456. if err != nil {
  457. return
  458. }
  459. ctl.mu.Lock()
  460. ctl.proxies[pxy.GetName()] = pxy
  461. ctl.mu.Unlock()
  462. return
  463. }
  464. func (ctl *Control) CloseProxy(closeMsg *msg.CloseProxy) (err error) {
  465. ctl.mu.Lock()
  466. pxy, ok := ctl.proxies[closeMsg.ProxyName]
  467. if !ok {
  468. ctl.mu.Unlock()
  469. return
  470. }
  471. if ctl.serverCfg.MaxPortsPerClient > 0 {
  472. ctl.portsUsedNum = ctl.portsUsedNum - pxy.GetUsedPortsNum()
  473. }
  474. pxy.Close()
  475. ctl.pxyManager.Del(pxy.GetName())
  476. delete(ctl.proxies, closeMsg.ProxyName)
  477. ctl.mu.Unlock()
  478. metrics.Server.CloseProxy(pxy.GetName(), pxy.GetConf().GetBaseInfo().ProxyType)
  479. return
  480. }