service.go 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651
  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. "bytes"
  17. "context"
  18. "crypto/tls"
  19. "fmt"
  20. "io"
  21. "net"
  22. "net/http"
  23. "os"
  24. "strconv"
  25. "time"
  26. "github.com/fatedier/golib/net/mux"
  27. fmux "github.com/hashicorp/yamux"
  28. quic "github.com/quic-go/quic-go"
  29. "github.com/samber/lo"
  30. "github.com/fatedier/frp/pkg/auth"
  31. v1 "github.com/fatedier/frp/pkg/config/v1"
  32. modelmetrics "github.com/fatedier/frp/pkg/metrics"
  33. "github.com/fatedier/frp/pkg/msg"
  34. "github.com/fatedier/frp/pkg/nathole"
  35. plugin "github.com/fatedier/frp/pkg/plugin/server"
  36. "github.com/fatedier/frp/pkg/ssh"
  37. "github.com/fatedier/frp/pkg/transport"
  38. httppkg "github.com/fatedier/frp/pkg/util/http"
  39. "github.com/fatedier/frp/pkg/util/log"
  40. netpkg "github.com/fatedier/frp/pkg/util/net"
  41. "github.com/fatedier/frp/pkg/util/tcpmux"
  42. "github.com/fatedier/frp/pkg/util/util"
  43. "github.com/fatedier/frp/pkg/util/version"
  44. "github.com/fatedier/frp/pkg/util/vhost"
  45. "github.com/fatedier/frp/pkg/util/xlog"
  46. "github.com/fatedier/frp/server/controller"
  47. "github.com/fatedier/frp/server/group"
  48. "github.com/fatedier/frp/server/metrics"
  49. "github.com/fatedier/frp/server/ports"
  50. "github.com/fatedier/frp/server/proxy"
  51. "github.com/fatedier/frp/server/visitor"
  52. )
  53. const (
  54. connReadTimeout time.Duration = 10 * time.Second
  55. vhostReadWriteTimeout time.Duration = 30 * time.Second
  56. )
  57. func init() {
  58. // Disable quic-go's receive buffer warning.
  59. os.Setenv("QUIC_GO_DISABLE_RECEIVE_BUFFER_WARNING", "true")
  60. }
  61. // Server service
  62. type Service struct {
  63. // Dispatch connections to different handlers listen on same port
  64. muxer *mux.Mux
  65. // Accept connections from client
  66. listener net.Listener
  67. // Accept connections using kcp
  68. kcpListener net.Listener
  69. // Accept connections using quic
  70. quicListener *quic.Listener
  71. // Accept connections using websocket
  72. websocketListener net.Listener
  73. // Accept frp tls connections
  74. tlsListener net.Listener
  75. // Accept pipe connections from ssh tunnel gateway
  76. sshTunnelListener *netpkg.InternalListener
  77. // Manage all controllers
  78. ctlManager *ControlManager
  79. // Manage all proxies
  80. pxyManager *proxy.Manager
  81. // Manage all plugins
  82. pluginManager *plugin.Manager
  83. // HTTP vhost router
  84. httpVhostRouter *vhost.Routers
  85. // All resource managers and controllers
  86. rc *controller.ResourceController
  87. // web server for dashboard UI and apis
  88. webServer *httppkg.Server
  89. sshTunnelGateway *ssh.Gateway
  90. // Verifies authentication based on selected method
  91. authVerifier auth.Verifier
  92. tlsConfig *tls.Config
  93. cfg *v1.ServerConfig
  94. // service context
  95. ctx context.Context
  96. // call cancel to stop service
  97. cancel context.CancelFunc
  98. }
  99. func NewService(cfg *v1.ServerConfig) (*Service, error) {
  100. tlsConfig, err := transport.NewServerTLSConfig(
  101. cfg.Transport.TLS.CertFile,
  102. cfg.Transport.TLS.KeyFile,
  103. cfg.Transport.TLS.TrustedCaFile)
  104. if err != nil {
  105. return nil, err
  106. }
  107. var webServer *httppkg.Server
  108. if cfg.WebServer.Port > 0 {
  109. ws, err := httppkg.NewServer(cfg.WebServer)
  110. if err != nil {
  111. return nil, err
  112. }
  113. webServer = ws
  114. modelmetrics.EnableMem()
  115. if cfg.EnablePrometheus {
  116. modelmetrics.EnablePrometheus()
  117. }
  118. }
  119. svr := &Service{
  120. ctlManager: NewControlManager(),
  121. pxyManager: proxy.NewManager(),
  122. pluginManager: plugin.NewManager(),
  123. rc: &controller.ResourceController{
  124. VisitorManager: visitor.NewManager(),
  125. TCPPortManager: ports.NewManager("tcp", cfg.ProxyBindAddr, cfg.AllowPorts),
  126. UDPPortManager: ports.NewManager("udp", cfg.ProxyBindAddr, cfg.AllowPorts),
  127. },
  128. sshTunnelListener: netpkg.NewInternalListener(),
  129. httpVhostRouter: vhost.NewRouters(),
  130. authVerifier: auth.NewAuthVerifier(cfg.Auth),
  131. webServer: webServer,
  132. tlsConfig: tlsConfig,
  133. cfg: cfg,
  134. ctx: context.Background(),
  135. }
  136. if webServer != nil {
  137. webServer.RouteRegister(svr.registerRouteHandlers)
  138. }
  139. // Create tcpmux httpconnect multiplexer.
  140. if cfg.TCPMuxHTTPConnectPort > 0 {
  141. var l net.Listener
  142. address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.TCPMuxHTTPConnectPort))
  143. l, err = net.Listen("tcp", address)
  144. if err != nil {
  145. return nil, fmt.Errorf("create server listener error, %v", err)
  146. }
  147. svr.rc.TCPMuxHTTPConnectMuxer, err = tcpmux.NewHTTPConnectTCPMuxer(l, cfg.TCPMuxPassthrough, vhostReadWriteTimeout)
  148. if err != nil {
  149. return nil, fmt.Errorf("create vhost tcpMuxer error, %v", err)
  150. }
  151. log.Infof("tcpmux httpconnect multiplexer listen on %s, passthough: %v", address, cfg.TCPMuxPassthrough)
  152. }
  153. // Init all plugins
  154. for _, p := range cfg.HTTPPlugins {
  155. svr.pluginManager.Register(plugin.NewHTTPPluginOptions(p))
  156. log.Infof("plugin [%s] has been registered", p.Name)
  157. }
  158. svr.rc.PluginManager = svr.pluginManager
  159. // Init group controller
  160. svr.rc.TCPGroupCtl = group.NewTCPGroupCtl(svr.rc.TCPPortManager)
  161. // Init HTTP group controller
  162. svr.rc.HTTPGroupCtl = group.NewHTTPGroupController(svr.httpVhostRouter)
  163. // Init TCP mux group controller
  164. svr.rc.TCPMuxGroupCtl = group.NewTCPMuxGroupCtl(svr.rc.TCPMuxHTTPConnectMuxer)
  165. // Init 404 not found page
  166. vhost.NotFoundPagePath = cfg.Custom404Page
  167. var (
  168. httpMuxOn bool
  169. httpsMuxOn bool
  170. )
  171. if cfg.BindAddr == cfg.ProxyBindAddr {
  172. if cfg.BindPort == cfg.VhostHTTPPort {
  173. httpMuxOn = true
  174. }
  175. if cfg.BindPort == cfg.VhostHTTPSPort {
  176. httpsMuxOn = true
  177. }
  178. }
  179. // Listen for accepting connections from client.
  180. address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.BindPort))
  181. ln, err := net.Listen("tcp", address)
  182. if err != nil {
  183. return nil, fmt.Errorf("create server listener error, %v", err)
  184. }
  185. svr.muxer = mux.NewMux(ln)
  186. svr.muxer.SetKeepAlive(time.Duration(cfg.Transport.TCPKeepAlive) * time.Second)
  187. go func() {
  188. _ = svr.muxer.Serve()
  189. }()
  190. ln = svr.muxer.DefaultListener()
  191. svr.listener = ln
  192. log.Infof("frps tcp listen on %s", address)
  193. // Listen for accepting connections from client using kcp protocol.
  194. if cfg.KCPBindPort > 0 {
  195. address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.KCPBindPort))
  196. svr.kcpListener, err = netpkg.ListenKcp(address)
  197. if err != nil {
  198. return nil, fmt.Errorf("listen on kcp udp address %s error: %v", address, err)
  199. }
  200. log.Infof("frps kcp listen on udp %s", address)
  201. }
  202. if cfg.QUICBindPort > 0 {
  203. address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.QUICBindPort))
  204. quicTLSCfg := tlsConfig.Clone()
  205. quicTLSCfg.NextProtos = []string{"frp"}
  206. svr.quicListener, err = quic.ListenAddr(address, quicTLSCfg, &quic.Config{
  207. MaxIdleTimeout: time.Duration(cfg.Transport.QUIC.MaxIdleTimeout) * time.Second,
  208. MaxIncomingStreams: int64(cfg.Transport.QUIC.MaxIncomingStreams),
  209. KeepAlivePeriod: time.Duration(cfg.Transport.QUIC.KeepalivePeriod) * time.Second,
  210. })
  211. if err != nil {
  212. return nil, fmt.Errorf("listen on quic udp address %s error: %v", address, err)
  213. }
  214. log.Infof("frps quic listen on %s", address)
  215. }
  216. if cfg.SSHTunnelGateway.BindPort > 0 {
  217. sshGateway, err := ssh.NewGateway(cfg.SSHTunnelGateway, cfg.ProxyBindAddr, svr.sshTunnelListener)
  218. if err != nil {
  219. return nil, fmt.Errorf("create ssh gateway error: %v", err)
  220. }
  221. svr.sshTunnelGateway = sshGateway
  222. log.Infof("frps sshTunnelGateway listen on port %d", cfg.SSHTunnelGateway.BindPort)
  223. }
  224. // Listen for accepting connections from client using websocket protocol.
  225. websocketPrefix := []byte("GET " + netpkg.FrpWebsocketPath)
  226. websocketLn := svr.muxer.Listen(0, uint32(len(websocketPrefix)), func(data []byte) bool {
  227. return bytes.Equal(data, websocketPrefix)
  228. })
  229. svr.websocketListener = netpkg.NewWebsocketListener(websocketLn)
  230. // Create http vhost muxer.
  231. if cfg.VhostHTTPPort > 0 {
  232. rp := vhost.NewHTTPReverseProxy(vhost.HTTPReverseProxyOptions{
  233. ResponseHeaderTimeoutS: cfg.VhostHTTPTimeout,
  234. }, svr.httpVhostRouter)
  235. svr.rc.HTTPReverseProxy = rp
  236. address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.VhostHTTPPort))
  237. server := &http.Server{
  238. Addr: address,
  239. Handler: rp,
  240. }
  241. var l net.Listener
  242. if httpMuxOn {
  243. l = svr.muxer.ListenHttp(1)
  244. } else {
  245. l, err = net.Listen("tcp", address)
  246. if err != nil {
  247. return nil, fmt.Errorf("create vhost http listener error, %v", err)
  248. }
  249. }
  250. go func() {
  251. _ = server.Serve(l)
  252. }()
  253. log.Infof("http service listen on %s", address)
  254. }
  255. // Create https vhost muxer.
  256. if cfg.VhostHTTPSPort > 0 {
  257. var l net.Listener
  258. if httpsMuxOn {
  259. l = svr.muxer.ListenHttps(1)
  260. } else {
  261. address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.VhostHTTPSPort))
  262. l, err = net.Listen("tcp", address)
  263. if err != nil {
  264. return nil, fmt.Errorf("create server listener error, %v", err)
  265. }
  266. log.Infof("https service listen on %s", address)
  267. }
  268. svr.rc.VhostHTTPSMuxer, err = vhost.NewHTTPSMuxer(l, vhostReadWriteTimeout)
  269. if err != nil {
  270. return nil, fmt.Errorf("create vhost httpsMuxer error, %v", err)
  271. }
  272. }
  273. // frp tls listener
  274. svr.tlsListener = svr.muxer.Listen(2, 1, func(data []byte) bool {
  275. // tls first byte can be 0x16 only when vhost https port is not same with bind port
  276. return int(data[0]) == netpkg.FRPTLSHeadByte || int(data[0]) == 0x16
  277. })
  278. // Create nat hole controller.
  279. nc, err := nathole.NewController(time.Duration(cfg.NatHoleAnalysisDataReserveHours) * time.Hour)
  280. if err != nil {
  281. return nil, fmt.Errorf("create nat hole controller error, %v", err)
  282. }
  283. svr.rc.NatHoleController = nc
  284. return svr, nil
  285. }
  286. func (svr *Service) Run(ctx context.Context) {
  287. ctx, cancel := context.WithCancel(ctx)
  288. svr.ctx = ctx
  289. svr.cancel = cancel
  290. // run dashboard web server.
  291. if svr.webServer != nil {
  292. go func() {
  293. log.Infof("dashboard listen on %s", svr.webServer.Address())
  294. if err := svr.webServer.Run(); err != nil {
  295. log.Warnf("dashboard server exit with error: %v", err)
  296. }
  297. }()
  298. }
  299. go svr.HandleListener(svr.sshTunnelListener, true)
  300. if svr.kcpListener != nil {
  301. go svr.HandleListener(svr.kcpListener, false)
  302. }
  303. if svr.quicListener != nil {
  304. go svr.HandleQUICListener(svr.quicListener)
  305. }
  306. go svr.HandleListener(svr.websocketListener, false)
  307. go svr.HandleListener(svr.tlsListener, false)
  308. if svr.rc.NatHoleController != nil {
  309. go svr.rc.NatHoleController.CleanWorker(svr.ctx)
  310. }
  311. if svr.sshTunnelGateway != nil {
  312. go svr.sshTunnelGateway.Run()
  313. }
  314. svr.HandleListener(svr.listener, false)
  315. <-svr.ctx.Done()
  316. // service context may not be canceled by svr.Close(), we should call it here to release resources
  317. if svr.listener != nil {
  318. svr.Close()
  319. }
  320. }
  321. func (svr *Service) Close() error {
  322. if svr.kcpListener != nil {
  323. svr.kcpListener.Close()
  324. svr.kcpListener = nil
  325. }
  326. if svr.quicListener != nil {
  327. svr.quicListener.Close()
  328. svr.quicListener = nil
  329. }
  330. if svr.websocketListener != nil {
  331. svr.websocketListener.Close()
  332. svr.websocketListener = nil
  333. }
  334. if svr.tlsListener != nil {
  335. svr.tlsListener.Close()
  336. svr.tlsConfig = nil
  337. }
  338. if svr.listener != nil {
  339. svr.listener.Close()
  340. svr.listener = nil
  341. }
  342. svr.ctlManager.Close()
  343. if svr.cancel != nil {
  344. svr.cancel()
  345. }
  346. return nil
  347. }
  348. func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool) {
  349. xl := xlog.FromContextSafe(ctx)
  350. var (
  351. rawMsg msg.Message
  352. err error
  353. )
  354. _ = conn.SetReadDeadline(time.Now().Add(connReadTimeout))
  355. if rawMsg, err = msg.ReadMsg(conn); err != nil {
  356. log.Tracef("Failed to read message: %v", err)
  357. conn.Close()
  358. return
  359. }
  360. _ = conn.SetReadDeadline(time.Time{})
  361. switch m := rawMsg.(type) {
  362. case *msg.Login:
  363. // server plugin hook
  364. content := &plugin.LoginContent{
  365. Login: *m,
  366. ClientAddress: conn.RemoteAddr().String(),
  367. }
  368. retContent, err := svr.pluginManager.Login(content)
  369. if err == nil {
  370. m = &retContent.Login
  371. err = svr.RegisterControl(conn, m, internal)
  372. }
  373. // If login failed, send error message there.
  374. // Otherwise send success message in control's work goroutine.
  375. if err != nil {
  376. xl.Warnf("register control error: %v", err)
  377. _ = msg.WriteMsg(conn, &msg.LoginResp{
  378. Version: version.Full(),
  379. Error: util.GenerateResponseErrorString("register control error", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
  380. })
  381. conn.Close()
  382. }
  383. case *msg.NewWorkConn:
  384. if err := svr.RegisterWorkConn(conn, m); err != nil {
  385. conn.Close()
  386. }
  387. case *msg.NewVisitorConn:
  388. if err = svr.RegisterVisitorConn(conn, m); err != nil {
  389. xl.Warnf("register visitor conn error: %v", err)
  390. _ = msg.WriteMsg(conn, &msg.NewVisitorConnResp{
  391. ProxyName: m.ProxyName,
  392. Error: util.GenerateResponseErrorString("register visitor conn error", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
  393. })
  394. conn.Close()
  395. } else {
  396. _ = msg.WriteMsg(conn, &msg.NewVisitorConnResp{
  397. ProxyName: m.ProxyName,
  398. Error: "",
  399. })
  400. }
  401. default:
  402. log.Warnf("Error message type for the new connection [%s]", conn.RemoteAddr().String())
  403. conn.Close()
  404. }
  405. }
  406. // HandleListener accepts connections from client and call handleConnection to handle them.
  407. // If internal is true, it means that this listener is used for internal communication like ssh tunnel gateway.
  408. // TODO(fatedier): Pass some parameters of listener/connection through context to avoid passing too many parameters.
  409. func (svr *Service) HandleListener(l net.Listener, internal bool) {
  410. // Listen for incoming connections from client.
  411. for {
  412. c, err := l.Accept()
  413. if err != nil {
  414. log.Warnf("Listener for incoming connections from client closed")
  415. return
  416. }
  417. // inject xlog object into net.Conn context
  418. xl := xlog.New()
  419. ctx := context.Background()
  420. c = netpkg.NewContextConn(xlog.NewContext(ctx, xl), c)
  421. if !internal {
  422. log.Tracef("start check TLS connection...")
  423. originConn := c
  424. forceTLS := svr.cfg.Transport.TLS.Force
  425. var isTLS, custom bool
  426. c, isTLS, custom, err = netpkg.CheckAndEnableTLSServerConnWithTimeout(c, svr.tlsConfig, forceTLS, connReadTimeout)
  427. if err != nil {
  428. log.Warnf("CheckAndEnableTLSServerConnWithTimeout error: %v", err)
  429. originConn.Close()
  430. continue
  431. }
  432. log.Tracef("check TLS connection success, isTLS: %v custom: %v internal: %v", isTLS, custom, internal)
  433. }
  434. // Start a new goroutine to handle connection.
  435. go func(ctx context.Context, frpConn net.Conn) {
  436. if lo.FromPtr(svr.cfg.Transport.TCPMux) && !internal {
  437. fmuxCfg := fmux.DefaultConfig()
  438. fmuxCfg.KeepAliveInterval = time.Duration(svr.cfg.Transport.TCPMuxKeepaliveInterval) * time.Second
  439. fmuxCfg.LogOutput = io.Discard
  440. fmuxCfg.MaxStreamWindowSize = 6 * 1024 * 1024
  441. session, err := fmux.Server(frpConn, fmuxCfg)
  442. if err != nil {
  443. log.Warnf("Failed to create mux connection: %v", err)
  444. frpConn.Close()
  445. return
  446. }
  447. for {
  448. stream, err := session.AcceptStream()
  449. if err != nil {
  450. log.Debugf("Accept new mux stream error: %v", err)
  451. session.Close()
  452. return
  453. }
  454. go svr.handleConnection(ctx, stream, internal)
  455. }
  456. } else {
  457. svr.handleConnection(ctx, frpConn, internal)
  458. }
  459. }(ctx, c)
  460. }
  461. }
  462. func (svr *Service) HandleQUICListener(l *quic.Listener) {
  463. // Listen for incoming connections from client.
  464. for {
  465. c, err := l.Accept(context.Background())
  466. if err != nil {
  467. log.Warnf("QUICListener for incoming connections from client closed")
  468. return
  469. }
  470. // Start a new goroutine to handle connection.
  471. go func(ctx context.Context, frpConn quic.Connection) {
  472. for {
  473. stream, err := frpConn.AcceptStream(context.Background())
  474. if err != nil {
  475. log.Debugf("Accept new quic mux stream error: %v", err)
  476. _ = frpConn.CloseWithError(0, "")
  477. return
  478. }
  479. go svr.handleConnection(ctx, netpkg.QuicStreamToNetConn(stream, frpConn), false)
  480. }
  481. }(context.Background(), c)
  482. }
  483. }
  484. func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login, internal bool) error {
  485. // If client's RunID is empty, it's a new client, we just create a new controller.
  486. // Otherwise, we check if there is one controller has the same run id. If so, we release previous controller and start new one.
  487. var err error
  488. if loginMsg.RunID == "" {
  489. loginMsg.RunID, err = util.RandID()
  490. if err != nil {
  491. return err
  492. }
  493. }
  494. ctx := netpkg.NewContextFromConn(ctlConn)
  495. xl := xlog.FromContextSafe(ctx)
  496. xl.AppendPrefix(loginMsg.RunID)
  497. ctx = xlog.NewContext(ctx, xl)
  498. xl.Infof("client login info: ip [%s] version [%s] hostname [%s] os [%s] arch [%s]",
  499. ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)
  500. // Check auth.
  501. authVerifier := svr.authVerifier
  502. if internal && loginMsg.ClientSpec.AlwaysAuthPass {
  503. authVerifier = auth.AlwaysPassVerifier
  504. }
  505. if err := authVerifier.VerifyLogin(loginMsg); err != nil {
  506. return err
  507. }
  508. // TODO(fatedier): use SessionContext
  509. ctl, err := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, authVerifier, ctlConn, !internal, loginMsg, svr.cfg)
  510. if err != nil {
  511. xl.Warnf("create new controller error: %v", err)
  512. // don't return detailed errors to client
  513. return fmt.Errorf("unexpected error when creating new controller")
  514. }
  515. if oldCtl := svr.ctlManager.Add(loginMsg.RunID, ctl); oldCtl != nil {
  516. oldCtl.WaitClosed()
  517. }
  518. ctl.Start()
  519. // for statistics
  520. metrics.Server.NewClient()
  521. go func() {
  522. // block until control closed
  523. ctl.WaitClosed()
  524. svr.ctlManager.Del(loginMsg.RunID, ctl)
  525. }()
  526. return nil
  527. }
  528. // RegisterWorkConn register a new work connection to control and proxies need it.
  529. func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn) error {
  530. xl := netpkg.NewLogFromConn(workConn)
  531. ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
  532. if !exist {
  533. xl.Warnf("No client control found for run id [%s]", newMsg.RunID)
  534. return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
  535. }
  536. // server plugin hook
  537. content := &plugin.NewWorkConnContent{
  538. User: plugin.UserInfo{
  539. User: ctl.loginMsg.User,
  540. Metas: ctl.loginMsg.Metas,
  541. RunID: ctl.loginMsg.RunID,
  542. },
  543. NewWorkConn: *newMsg,
  544. }
  545. retContent, err := svr.pluginManager.NewWorkConn(content)
  546. if err == nil {
  547. newMsg = &retContent.NewWorkConn
  548. // Check auth.
  549. err = ctl.authVerifier.VerifyNewWorkConn(newMsg)
  550. }
  551. if err != nil {
  552. xl.Warnf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
  553. _ = msg.WriteMsg(workConn, &msg.StartWorkConn{
  554. Error: util.GenerateResponseErrorString("invalid NewWorkConn", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
  555. })
  556. return fmt.Errorf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
  557. }
  558. return ctl.RegisterWorkConn(workConn)
  559. }
  560. func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn) error {
  561. visitorUser := ""
  562. // TODO(deprecation): Compatible with old versions, can be without runID, user is empty. In later versions, it will be mandatory to include runID.
  563. // If runID is required, it is not compatible with versions prior to v0.50.0.
  564. if newMsg.RunID != "" {
  565. ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
  566. if !exist {
  567. return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
  568. }
  569. visitorUser = ctl.loginMsg.User
  570. }
  571. return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
  572. newMsg.UseEncryption, newMsg.UseCompression, visitorUser)
  573. }