1
0

service.go 19 KB

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