1
0

service.go 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658
  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. ReadHeaderTimeout: 60 * time.Second,
  247. }
  248. var l net.Listener
  249. if httpMuxOn {
  250. l = svr.muxer.ListenHTTP(1)
  251. } else {
  252. l, err = net.Listen("tcp", address)
  253. if err != nil {
  254. return nil, fmt.Errorf("create vhost http listener error, %v", err)
  255. }
  256. }
  257. go func() {
  258. _ = server.Serve(l)
  259. }()
  260. log.Infof("http service listen on %s", address)
  261. }
  262. // Create https vhost muxer.
  263. if cfg.VhostHTTPSPort > 0 {
  264. var l net.Listener
  265. if httpsMuxOn {
  266. l = svr.muxer.ListenHTTPS(1)
  267. } else {
  268. address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.VhostHTTPSPort))
  269. l, err = net.Listen("tcp", address)
  270. if err != nil {
  271. return nil, fmt.Errorf("create server listener error, %v", err)
  272. }
  273. log.Infof("https service listen on %s", address)
  274. }
  275. svr.rc.VhostHTTPSMuxer, err = vhost.NewHTTPSMuxer(l, vhostReadWriteTimeout)
  276. if err != nil {
  277. return nil, fmt.Errorf("create vhost httpsMuxer error, %v", err)
  278. }
  279. }
  280. // frp tls listener
  281. svr.tlsListener = svr.muxer.Listen(2, 1, func(data []byte) bool {
  282. // tls first byte can be 0x16 only when vhost https port is not same with bind port
  283. return int(data[0]) == netpkg.FRPTLSHeadByte || int(data[0]) == 0x16
  284. })
  285. // Create nat hole controller.
  286. nc, err := nathole.NewController(time.Duration(cfg.NatHoleAnalysisDataReserveHours) * time.Hour)
  287. if err != nil {
  288. return nil, fmt.Errorf("create nat hole controller error, %v", err)
  289. }
  290. svr.rc.NatHoleController = nc
  291. return svr, nil
  292. }
  293. func (svr *Service) Run(ctx context.Context) {
  294. ctx, cancel := context.WithCancel(ctx)
  295. svr.ctx = ctx
  296. svr.cancel = cancel
  297. // run dashboard web server.
  298. if svr.webServer != nil {
  299. go func() {
  300. log.Infof("dashboard listen on %s", svr.webServer.Address())
  301. if err := svr.webServer.Run(); err != nil {
  302. log.Warnf("dashboard server exit with error: %v", err)
  303. }
  304. }()
  305. }
  306. go svr.HandleListener(svr.sshTunnelListener, true)
  307. if svr.kcpListener != nil {
  308. go svr.HandleListener(svr.kcpListener, false)
  309. }
  310. if svr.quicListener != nil {
  311. go svr.HandleQUICListener(svr.quicListener)
  312. }
  313. go svr.HandleListener(svr.websocketListener, false)
  314. go svr.HandleListener(svr.tlsListener, false)
  315. if svr.rc.NatHoleController != nil {
  316. go svr.rc.NatHoleController.CleanWorker(svr.ctx)
  317. }
  318. if svr.sshTunnelGateway != nil {
  319. go svr.sshTunnelGateway.Run()
  320. }
  321. svr.HandleListener(svr.listener, false)
  322. <-svr.ctx.Done()
  323. // service context may not be canceled by svr.Close(), we should call it here to release resources
  324. if svr.listener != nil {
  325. svr.Close()
  326. }
  327. }
  328. func (svr *Service) Close() error {
  329. if svr.kcpListener != nil {
  330. svr.kcpListener.Close()
  331. svr.kcpListener = nil
  332. }
  333. if svr.quicListener != nil {
  334. svr.quicListener.Close()
  335. svr.quicListener = nil
  336. }
  337. if svr.websocketListener != nil {
  338. svr.websocketListener.Close()
  339. svr.websocketListener = nil
  340. }
  341. if svr.tlsListener != nil {
  342. svr.tlsListener.Close()
  343. svr.tlsConfig = nil
  344. }
  345. if svr.listener != nil {
  346. svr.listener.Close()
  347. svr.listener = nil
  348. }
  349. svr.ctlManager.Close()
  350. if svr.cancel != nil {
  351. svr.cancel()
  352. }
  353. return nil
  354. }
  355. func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool) {
  356. xl := xlog.FromContextSafe(ctx)
  357. var (
  358. rawMsg msg.Message
  359. err error
  360. )
  361. _ = conn.SetReadDeadline(time.Now().Add(connReadTimeout))
  362. if rawMsg, err = msg.ReadMsg(conn); err != nil {
  363. log.Tracef("Failed to read message: %v", err)
  364. conn.Close()
  365. return
  366. }
  367. _ = conn.SetReadDeadline(time.Time{})
  368. switch m := rawMsg.(type) {
  369. case *msg.Login:
  370. // server plugin hook
  371. content := &plugin.LoginContent{
  372. Login: *m,
  373. ClientAddress: conn.RemoteAddr().String(),
  374. }
  375. retContent, err := svr.pluginManager.Login(content)
  376. if err == nil {
  377. m = &retContent.Login
  378. err = svr.RegisterControl(conn, m, internal)
  379. }
  380. // If login failed, send error message there.
  381. // Otherwise send success message in control's work goroutine.
  382. if err != nil {
  383. xl.Warnf("register control error: %v", err)
  384. _ = msg.WriteMsg(conn, &msg.LoginResp{
  385. Version: version.Full(),
  386. Error: util.GenerateResponseErrorString("register control error", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
  387. })
  388. conn.Close()
  389. }
  390. case *msg.NewWorkConn:
  391. if err := svr.RegisterWorkConn(conn, m); err != nil {
  392. conn.Close()
  393. }
  394. case *msg.NewVisitorConn:
  395. if err = svr.RegisterVisitorConn(conn, m); err != nil {
  396. xl.Warnf("register visitor conn error: %v", err)
  397. _ = msg.WriteMsg(conn, &msg.NewVisitorConnResp{
  398. ProxyName: m.ProxyName,
  399. Error: util.GenerateResponseErrorString("register visitor conn error", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
  400. })
  401. conn.Close()
  402. } else {
  403. _ = msg.WriteMsg(conn, &msg.NewVisitorConnResp{
  404. ProxyName: m.ProxyName,
  405. Error: "",
  406. })
  407. }
  408. default:
  409. log.Warnf("Error message type for the new connection [%s]", conn.RemoteAddr().String())
  410. conn.Close()
  411. }
  412. }
  413. // HandleListener accepts connections from client and call handleConnection to handle them.
  414. // If internal is true, it means that this listener is used for internal communication like ssh tunnel gateway.
  415. // TODO(fatedier): Pass some parameters of listener/connection through context to avoid passing too many parameters.
  416. func (svr *Service) HandleListener(l net.Listener, internal bool) {
  417. // Listen for incoming connections from client.
  418. for {
  419. c, err := l.Accept()
  420. if err != nil {
  421. log.Warnf("Listener for incoming connections from client closed")
  422. return
  423. }
  424. // inject xlog object into net.Conn context
  425. xl := xlog.New()
  426. ctx := context.Background()
  427. c = netpkg.NewContextConn(xlog.NewContext(ctx, xl), c)
  428. if !internal {
  429. log.Tracef("start check TLS connection...")
  430. originConn := c
  431. forceTLS := svr.cfg.Transport.TLS.Force
  432. var isTLS, custom bool
  433. c, isTLS, custom, err = netpkg.CheckAndEnableTLSServerConnWithTimeout(c, svr.tlsConfig, forceTLS, connReadTimeout)
  434. if err != nil {
  435. log.Warnf("CheckAndEnableTLSServerConnWithTimeout error: %v", err)
  436. originConn.Close()
  437. continue
  438. }
  439. log.Tracef("check TLS connection success, isTLS: %v custom: %v internal: %v", isTLS, custom, internal)
  440. }
  441. // Start a new goroutine to handle connection.
  442. go func(ctx context.Context, frpConn net.Conn) {
  443. if lo.FromPtr(svr.cfg.Transport.TCPMux) && !internal {
  444. fmuxCfg := fmux.DefaultConfig()
  445. fmuxCfg.KeepAliveInterval = time.Duration(svr.cfg.Transport.TCPMuxKeepaliveInterval) * time.Second
  446. fmuxCfg.LogOutput = io.Discard
  447. fmuxCfg.MaxStreamWindowSize = 6 * 1024 * 1024
  448. session, err := fmux.Server(frpConn, fmuxCfg)
  449. if err != nil {
  450. log.Warnf("Failed to create mux connection: %v", err)
  451. frpConn.Close()
  452. return
  453. }
  454. for {
  455. stream, err := session.AcceptStream()
  456. if err != nil {
  457. log.Debugf("Accept new mux stream error: %v", err)
  458. session.Close()
  459. return
  460. }
  461. go svr.handleConnection(ctx, stream, internal)
  462. }
  463. } else {
  464. svr.handleConnection(ctx, frpConn, internal)
  465. }
  466. }(ctx, c)
  467. }
  468. }
  469. func (svr *Service) HandleQUICListener(l *quic.Listener) {
  470. // Listen for incoming connections from client.
  471. for {
  472. c, err := l.Accept(context.Background())
  473. if err != nil {
  474. log.Warnf("QUICListener for incoming connections from client closed")
  475. return
  476. }
  477. // Start a new goroutine to handle connection.
  478. go func(ctx context.Context, frpConn quic.Connection) {
  479. for {
  480. stream, err := frpConn.AcceptStream(context.Background())
  481. if err != nil {
  482. log.Debugf("Accept new quic mux stream error: %v", err)
  483. _ = frpConn.CloseWithError(0, "")
  484. return
  485. }
  486. go svr.handleConnection(ctx, netpkg.QuicStreamToNetConn(stream, frpConn), false)
  487. }
  488. }(context.Background(), c)
  489. }
  490. }
  491. func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login, internal bool) error {
  492. // If client's RunID is empty, it's a new client, we just create a new controller.
  493. // Otherwise, we check if there is one controller has the same run id. If so, we release previous controller and start new one.
  494. var err error
  495. if loginMsg.RunID == "" {
  496. loginMsg.RunID, err = util.RandID()
  497. if err != nil {
  498. return err
  499. }
  500. }
  501. ctx := netpkg.NewContextFromConn(ctlConn)
  502. xl := xlog.FromContextSafe(ctx)
  503. xl.AppendPrefix(loginMsg.RunID)
  504. ctx = xlog.NewContext(ctx, xl)
  505. xl.Infof("client login info: ip [%s] version [%s] hostname [%s] os [%s] arch [%s]",
  506. ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)
  507. // Check auth.
  508. authVerifier := svr.authVerifier
  509. if internal && loginMsg.ClientSpec.AlwaysAuthPass {
  510. authVerifier = auth.AlwaysPassVerifier
  511. }
  512. if err := authVerifier.VerifyLogin(loginMsg); err != nil {
  513. return err
  514. }
  515. // TODO(fatedier): use SessionContext
  516. ctl, err := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, authVerifier, ctlConn, !internal, loginMsg, svr.cfg)
  517. if err != nil {
  518. xl.Warnf("create new controller error: %v", err)
  519. // don't return detailed errors to client
  520. return fmt.Errorf("unexpected error when creating new controller")
  521. }
  522. if oldCtl := svr.ctlManager.Add(loginMsg.RunID, ctl); oldCtl != nil {
  523. oldCtl.WaitClosed()
  524. }
  525. ctl.Start()
  526. // for statistics
  527. metrics.Server.NewClient()
  528. go func() {
  529. // block until control closed
  530. ctl.WaitClosed()
  531. svr.ctlManager.Del(loginMsg.RunID, ctl)
  532. }()
  533. return nil
  534. }
  535. // RegisterWorkConn register a new work connection to control and proxies need it.
  536. func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn) error {
  537. xl := netpkg.NewLogFromConn(workConn)
  538. ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
  539. if !exist {
  540. xl.Warnf("No client control found for run id [%s]", newMsg.RunID)
  541. return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
  542. }
  543. // server plugin hook
  544. content := &plugin.NewWorkConnContent{
  545. User: plugin.UserInfo{
  546. User: ctl.loginMsg.User,
  547. Metas: ctl.loginMsg.Metas,
  548. RunID: ctl.loginMsg.RunID,
  549. },
  550. NewWorkConn: *newMsg,
  551. }
  552. retContent, err := svr.pluginManager.NewWorkConn(content)
  553. if err == nil {
  554. newMsg = &retContent.NewWorkConn
  555. // Check auth.
  556. err = ctl.authVerifier.VerifyNewWorkConn(newMsg)
  557. }
  558. if err != nil {
  559. xl.Warnf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
  560. _ = msg.WriteMsg(workConn, &msg.StartWorkConn{
  561. Error: util.GenerateResponseErrorString("invalid NewWorkConn", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
  562. })
  563. return fmt.Errorf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
  564. }
  565. return ctl.RegisterWorkConn(workConn)
  566. }
  567. func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn) error {
  568. visitorUser := ""
  569. // TODO(deprecation): Compatible with old versions, can be without runID, user is empty. In later versions, it will be mandatory to include runID.
  570. // If runID is required, it is not compatible with versions prior to v0.50.0.
  571. if newMsg.RunID != "" {
  572. ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
  573. if !exist {
  574. return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
  575. }
  576. visitorUser = ctl.loginMsg.User
  577. }
  578. return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
  579. newMsg.UseEncryption, newMsg.UseCompression, visitorUser)
  580. }