lolia-frp/server/service.go

// Copyright 2017 fatedier, fatedier@gmail.com
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package server

import (
	"bytes"
	"context"
	"crypto/tls"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/fatedier/golib/crypto"
	"github.com/fatedier/golib/net/mux"
	fmux "github.com/fatedier/yamux"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	quic "github.com/quic-go/quic-go"
	"github.com/samber/lo"

	"github.com/fatedier/frp/pkg/auth"
	v1 "github.com/fatedier/frp/pkg/config/v1"
	"github.com/fatedier/frp/pkg/config/v1/validation"
	modelmetrics "github.com/fatedier/frp/pkg/metrics"
	"github.com/fatedier/frp/pkg/msg"
	"github.com/fatedier/frp/pkg/nathole"
	plugin "github.com/fatedier/frp/pkg/plugin/server"
	"github.com/fatedier/frp/pkg/proto/wire"
	"github.com/fatedier/frp/pkg/ssh"
	"github.com/fatedier/frp/pkg/transport"
	httppkg "github.com/fatedier/frp/pkg/util/http"
	"github.com/fatedier/frp/pkg/util/log"
	netpkg "github.com/fatedier/frp/pkg/util/net"
	"github.com/fatedier/frp/pkg/util/tcpmux"
	"github.com/fatedier/frp/pkg/util/util"
	"github.com/fatedier/frp/pkg/util/version"
	"github.com/fatedier/frp/pkg/util/vhost"
	"github.com/fatedier/frp/pkg/util/xlog"
	"github.com/fatedier/frp/server/controller"
	"github.com/fatedier/frp/server/group"
	"github.com/fatedier/frp/server/ports"
	"github.com/fatedier/frp/server/proxy"
	"github.com/fatedier/frp/server/registry"
	"github.com/fatedier/frp/server/visitor"
)

const (
	connReadTimeout       time.Duration = 10 * time.Second
	connWriteTimeout      time.Duration = 5 * time.Second
	vhostReadWriteTimeout time.Duration = 30 * time.Second
)

var errControlReplaced = errors.New("control was replaced during login")

func init() {
	crypto.DefaultSalt = "frp"
	// Disable quic-go's receive buffer warning.
	os.Setenv("QUIC_GO_DISABLE_RECEIVE_BUFFER_WARNING", "true")
	// Disable quic-go's ECN support by default. It may cause issues on certain operating systems.
	if os.Getenv("QUIC_GO_DISABLE_ECN") == "" {
		os.Setenv("QUIC_GO_DISABLE_ECN", "true")
	}
}

// Server service
type Service struct {
	// Dispatch connections to different handlers listen on same port
	muxer *mux.Mux

	// Accept connections from client
	listener net.Listener

	// Accept connections using kcp
	kcpListener net.Listener

	// Accept connections using quic
	quicListener *quic.Listener

	// Accept connections using websocket
	websocketListener net.Listener

	// Accept frp tls connections
	tlsListener net.Listener

	// Accept pipe connections from ssh tunnel gateway
	sshTunnelListener *netpkg.InternalListener

	// Manage all controllers
	ctlManager *ControlManager

	// Track logical clients keyed by user.clientID (runID fallback when raw clientID is empty).
	clientRegistry *registry.ClientRegistry

	// Manage all proxies
	pxyManager *proxy.Manager

	// Manage all plugins
	pluginManager *plugin.Manager

	// HTTP vhost router
	httpVhostRouter *vhost.Routers

	// All resource managers and controllers
	rc *controller.ResourceController

	// web server for dashboard UI and apis
	webServer     *httppkg.Server
	metricsServer *httppkg.Server
	// HTTP vhost resources are started and owned by NewService.
	httpVhostServer   *http.Server
	httpVhostListener net.Listener

	sshTunnelGateway *ssh.Gateway

	// Auth runtime and encryption materials
	auth *auth.ServerAuth

	tlsConfig *tls.Config

	cfg *v1.ServerConfig

	// service context
	ctx context.Context
	// call cancel to stop service
	cancel      context.CancelFunc
	lifecycleMu sync.Mutex
	closed      bool
	closeOnce   sync.Once
}

func NewService(cfg *v1.ServerConfig) (*Service, error) {
	tlsConfig, err := transport.NewServerTLSConfig(
		cfg.Transport.TLS.CertFile,
		cfg.Transport.TLS.KeyFile,
		cfg.Transport.TLS.TrustedCaFile)
	if err != nil {
		return nil, err
	}

	authRuntime, err := auth.BuildServerAuth(&cfg.Auth)
	if err != nil {
		return nil, err
	}

	clientRegistry := registry.NewClientRegistry()
	svr := &Service{
		ctlManager:     NewControlManager(clientRegistry),
		clientRegistry: clientRegistry,
		pxyManager:     proxy.NewManager(),
		pluginManager:  plugin.NewManager(),
		rc: &controller.ResourceController{
			VisitorManager: visitor.NewManager(),
		},
		sshTunnelListener: netpkg.NewInternalListener(),
		httpVhostRouter:   vhost.NewRouters(),
		auth:              authRuntime,
		tlsConfig:         tlsConfig,
		cfg:               cfg,
		ctx:               context.Background(),
	}

	// Create tcpmux httpconnect multiplexer.
	if cfg.TCPMuxHTTPConnectPort > 0 {
		var l net.Listener
		address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.TCPMuxHTTPConnectPort))
		l, err = net.Listen("tcp", address)
		if err != nil {
			return nil, fmt.Errorf("create server listener error, %v", err)
		}

		svr.rc.TCPMuxHTTPConnectMuxer, err = tcpmux.NewHTTPConnectTCPMuxer(l, cfg.TCPMuxPassthrough, vhostReadWriteTimeout)
		if err != nil {
			return nil, fmt.Errorf("create vhost tcpMuxer error, %v", err)
		}
		log.Infof("tcpmux httpconnect multiplexer listen on %s, passthrough: %v", address, cfg.TCPMuxPassthrough)
	}

	// Init all plugins
	for _, p := range cfg.HTTPPlugins {
		svr.pluginManager.Register(plugin.NewHTTPPluginOptions(p))
		log.Infof("plugin [%s] has been registered", p.Name)
	}
	svr.rc.PluginManager = svr.pluginManager

	// Init HTTP group controller
	svr.rc.HTTPGroupCtl = group.NewHTTPGroupController(svr.httpVhostRouter)

	// Init TCP mux group controller
	svr.rc.TCPMuxGroupCtl = group.NewTCPMuxGroupCtl(svr.rc.TCPMuxHTTPConnectMuxer)

	// Init 404 not found page
	vhost.NotFoundPagePath = cfg.Custom404Page

	var (
		httpMuxOn  bool
		httpsMuxOn bool
	)
	if cfg.BindAddr == cfg.ProxyBindAddr {
		if cfg.BindPort == cfg.VhostHTTPPort {
			httpMuxOn = true
		}
		if cfg.BindPort == cfg.VhostHTTPSPort {
			httpsMuxOn = true
		}
	}

	// Listen for accepting connections from client.
	address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.BindPort))
	ln, err := net.Listen("tcp", address)
	if err != nil {
		return nil, fmt.Errorf("create server listener error, %v", err)
	}

	svr.muxer = mux.NewMux(ln)
	svr.muxer.SetKeepAlive(time.Duration(cfg.Transport.TCPKeepAlive) * time.Second)
	go func() {
		_ = svr.muxer.Serve()
	}()
	ln = svr.muxer.DefaultListener()

	svr.listener = ln
	log.Infof("frps tcp listen on %s", address)

	// Listen for accepting connections from client using kcp protocol.
	if cfg.KCPBindPort > 0 {
		address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.KCPBindPort))
		svr.kcpListener, err = netpkg.ListenKcp(address)
		if err != nil {
			return nil, fmt.Errorf("listen on kcp udp address %s error: %v", address, err)
		}
		log.Infof("frps kcp listen on udp %s", address)
	}

	if cfg.QUICBindPort > 0 {
		address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.QUICBindPort))
		quicTLSCfg := tlsConfig.Clone()
		quicTLSCfg.NextProtos = []string{"frp"}
		svr.quicListener, err = quic.ListenAddr(address, quicTLSCfg, &quic.Config{
			MaxIdleTimeout:     time.Duration(cfg.Transport.QUIC.MaxIdleTimeout) * time.Second,
			MaxIncomingStreams: int64(cfg.Transport.QUIC.MaxIncomingStreams),
			KeepAlivePeriod:    time.Duration(cfg.Transport.QUIC.KeepalivePeriod) * time.Second,
		})
		if err != nil {
			return nil, fmt.Errorf("listen on quic udp address %s error: %v", address, err)
		}
		log.Infof("frps quic listen on %s", address)
	}

	if cfg.SSHTunnelGateway.BindPort > 0 {
		sshGateway, err := ssh.NewGateway(cfg.SSHTunnelGateway, cfg.BindAddr, svr.sshTunnelListener)
		if err != nil {
			return nil, fmt.Errorf("create ssh gateway error: %v", err)
		}
		svr.sshTunnelGateway = sshGateway
		log.Infof("frps sshTunnelGateway listen on port %d", cfg.SSHTunnelGateway.BindPort)
	}

	// Listen for accepting connections from client using websocket protocol.
	websocketPrefix := []byte("GET " + netpkg.FrpWebsocketPath)
	websocketLn := svr.muxer.Listen(0, uint32(len(websocketPrefix)), func(data []byte) bool {
		return bytes.Equal(data, websocketPrefix)
	})
	svr.websocketListener = netpkg.NewWebsocketListener(websocketLn)

	// Create http vhost muxer.
	if cfg.VhostHTTPPort > 0 {
		rp := vhost.NewHTTPReverseProxy(vhost.HTTPReverseProxyOptions{
			ResponseHeaderTimeoutS: cfg.VhostHTTPTimeout,
		}, svr.httpVhostRouter)
		svr.rc.HTTPReverseProxy = rp

		address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.VhostHTTPPort))
		protocols := new(http.Protocols)
		protocols.SetHTTP1(true)
		protocols.SetUnencryptedHTTP2(true)
		server := &http.Server{
			Addr:              address,
			Handler:           rp,
			ReadHeaderTimeout: 60 * time.Second,
			Protocols:         protocols,
		}
		var l net.Listener
		if httpMuxOn {
			l = svr.muxer.ListenHTTP(1)
		} else {
			l, err = net.Listen("tcp", address)
			if err != nil {
				return nil, fmt.Errorf("create vhost http listener error, %v", err)
			}
		}
		svr.httpVhostServer = server
		svr.httpVhostListener = l
		go func() {
			_ = server.Serve(l)
		}()
		log.Infof("http service listen on %s", address)
	}

	// Create https vhost muxer.
	if cfg.VhostHTTPSPort > 0 {
		var l net.Listener
		if httpsMuxOn {
			l = svr.muxer.ListenHTTPS(1)
		} else {
			address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.VhostHTTPSPort))
			l, err = net.Listen("tcp", address)
			if err != nil {
				return nil, fmt.Errorf("create server listener error, %v", err)
			}
			log.Infof("https service listen on %s", address)
		}

		svr.rc.VhostHTTPSMuxer, err = vhost.NewHTTPSMuxer(l, vhostReadWriteTimeout)
		if err != nil {
			return nil, fmt.Errorf("create vhost httpsMuxer error, %v", err)
		}

		// Init HTTPS group controller after HTTPSMuxer is created
		svr.rc.HTTPSGroupCtl = group.NewHTTPSGroupController(svr.rc.VhostHTTPSMuxer)
	}

	// frp tls listener
	svr.tlsListener = svr.muxer.Listen(2, 1, func(data []byte) bool {
		// tls first byte can be 0x16 only when vhost https port is not same with bind port
		return int(data[0]) == netpkg.FRPTLSHeadByte || int(data[0]) == 0x16
	})

	// Create nat hole controller.
	nc, err := nathole.NewController(time.Duration(cfg.NatHoleAnalysisDataReserveHours) * time.Hour)
	if err != nil {
		return nil, fmt.Errorf("create nat hole controller error, %v", err)
	}
	svr.rc.NatHoleController = nc

	// Create HTTP servers after all other fallible steps, rolling back listeners
	// if either server cannot be created.
	if cfg.WebServer.Port > 0 {
		ws, err := httppkg.NewServer(cfg.WebServer)
		if err != nil {
			// All other listeners are already created and must be released on failure.
			_ = svr.Close()
			return nil, err
		}
		svr.webServer = ws
	}
	if cfg.MetricsServer.Port > 0 {
		ms, err := httppkg.NewMetricsServer(cfg.MetricsServer, promhttp.Handler())
		if err != nil {
			_ = svr.Close()
			return nil, err
		}
		svr.metricsServer = ms
	}

	// Collectors are enabled at most once. The in-memory collector is only read
	// by the dashboard API, and Prometheus metrics need an endpoint to serve them.
	enableMem, enablePrometheus := metricsCollectors(cfg)
	if enableMem {
		modelmetrics.EnableMem()
	}
	if enablePrometheus {
		modelmetrics.EnablePrometheus()
	}

	// Start port-manager cleanup workers only after all fallible steps have succeeded.
	svr.rc.TCPPortManager = ports.NewManager("tcp", cfg.ProxyBindAddr, cfg.AllowPorts)
	svr.rc.UDPPortManager = ports.NewManager("udp", cfg.ProxyBindAddr, cfg.AllowPorts)
	svr.rc.TCPGroupCtl = group.NewTCPGroupCtl(svr.rc.TCPPortManager)

	if svr.webServer != nil {
		svr.webServer.RouteRegister(svr.registerRouteHandlers)
	}
	return svr, nil
}

// metricsCollectors reports which collectors should be started for cfg. The
// in-memory collector only backs the dashboard API (webServer), while the
// Prometheus collector needs a listener that serves /metrics: either the new
// metricsServer, or the legacy webServer endpoint gated by enablePrometheus.
func metricsCollectors(cfg *v1.ServerConfig) (mem bool, prometheus bool) {
	mem = cfg.WebServer.Port > 0
	prometheus = cfg.MetricsServer.Port > 0 || (cfg.EnablePrometheus && mem)
	return mem, prometheus
}

func (svr *Service) Run(ctx context.Context) {
	ctx, cancel := context.WithCancel(ctx)
	svr.lifecycleMu.Lock()
	if svr.closed {
		svr.lifecycleMu.Unlock()
		cancel()
		return
	}
	svr.ctx = ctx
	svr.cancel = cancel
	svr.lifecycleMu.Unlock()
	defer svr.Close()
	if ctx.Err() != nil {
		return
	}
	// Cancellation must unblock the synchronous main listener's Accept.
	stopClose := context.AfterFunc(ctx, func() { _ = svr.Close() })
	defer stopClose()

	// run dashboard web server.
	if svr.webServer != nil {
		if svr.cfg.WebServer.PprofEnable {
			log.Warnf("webServer pprof endpoints are deprecated; use metricsServer.pprofEnable instead")
		}
		if svr.cfg.EnablePrometheus {
			log.Warnf("webServer /metrics is deprecated; configure metricsServer.port instead")
		}
		if svr.cfg.MetricsServer.Port == 0 {
			log.Infof("webServer /healthz is deprecated; health probes can migrate to metricsServer.port")
		}
		go func() {
			log.Infof("dashboard listen on %s", svr.webServer.Address())
			if err := svr.webServer.Run(); err != nil {
				log.Warnf("dashboard server exit with error: %v", err)
			}
		}()
	}

	if svr.metricsServer != nil {
		go func() {
			log.Infof("metrics server listen on %s", svr.metricsServer.Address())
			if err := svr.metricsServer.Run(); err != nil && !errors.Is(err, http.ErrServerClosed) {
				log.Warnf("metrics server exit with error: %v", err)
			}
		}()
	}

	go svr.HandleListener(svr.sshTunnelListener, true)

	if svr.kcpListener != nil {
		go svr.HandleListener(svr.kcpListener, false)
	}
	if svr.quicListener != nil {
		go svr.HandleQUICListener(svr.quicListener)
	}
	go svr.HandleListener(svr.websocketListener, false)
	go svr.HandleListener(svr.tlsListener, false)

	if svr.rc.NatHoleController != nil {
		go svr.rc.NatHoleController.CleanWorker(svr.ctx)
	}

	if svr.sshTunnelGateway != nil {
		go svr.sshTunnelGateway.Run()
	}

	svr.HandleListener(svr.listener, false)

	<-svr.ctx.Done()
}

func (svr *Service) Close() error {
	svr.lifecycleMu.Lock()
	svr.closed = true
	cancel := svr.cancel
	svr.lifecycleMu.Unlock()
	if cancel != nil {
		cancel()
	}
	svr.closeOnce.Do(svr.closeResources)
	return nil
}

func (svr *Service) closeResources() {
	if svr.kcpListener != nil {
		svr.kcpListener.Close()
	}
	if svr.quicListener != nil {
		svr.quicListener.Close()
	}
	if svr.websocketListener != nil {
		svr.websocketListener.Close()
	}
	if svr.tlsListener != nil {
		svr.tlsListener.Close()
	}
	if svr.sshTunnelListener != nil {
		svr.sshTunnelListener.Close()
	}
	if svr.listener != nil {
		svr.listener.Close()
	}
	if svr.webServer != nil {
		svr.webServer.Close()
	}
	if svr.metricsServer != nil {
		svr.metricsServer.Close()
	}
	if svr.httpVhostServer != nil {
		svr.httpVhostServer.Close()
		// Close even if Serve has not registered its listener yet.
		svr.httpVhostListener.Close()
	}
	if svr.sshTunnelGateway != nil {
		svr.sshTunnelGateway.Close()
	}
	svr.rc.Close()
	svr.muxer.Close()
	svr.ctlManager.Close()
}

func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool, session io.Closer) {
	xl := xlog.FromContextSafe(ctx)

	acceptedConn, err := svr.acceptConnection(ctx, conn)
	if err != nil {
		log.Tracef("failed to accept frp connection: %v", err)
		conn.Close()
		return
	}
	conn = acceptedConn.conn

	switch m := acceptedConn.firstMsg.(type) {
	case *msg.Login:
		// server plugin hook
		content := &plugin.LoginContent{
			Login:         *m,
			ClientAddress: conn.RemoteAddr().String(),
		}
		retContent, err := svr.pluginManager.Login(content)
		var ctl *Control
		if err == nil {
			m = &retContent.Login
			controlConn := acceptedConn.conn
			if !internal {
				var controlRW io.ReadWriter
				controlRW, err = acceptedConn.newControlReadWriter(conn, svr.auth.EncryptionKey())
				if err == nil {
					controlConn = acceptedConn.messageConnFor(controlRW)
				}
			}
			if err == nil {
				ctl, err = svr.RegisterControl(controlConn, m, internal, acceptedConn.wireProtocol, acceptedConn.udpPacketCodec, session)
			}
		}

		if err != nil {
			xl.Warnf("register control error: %v", err)
			if ctl != nil {
				svr.ctlManager.Remove(ctl)
			}
			if writeErr := writeWithDeadline(conn, connWriteTimeout, func() error {
				return acceptedConn.conn.WriteMsg(&msg.LoginResp{
					Version: version.Full(),
					Error:   util.GenerateResponseErrorString("register control error", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
				})
			}); writeErr != nil {
				xl.Warnf("write login error response error: %v", writeErr)
			}
			if ctl != nil {
				_ = ctl.Close()
			} else {
				conn.Close()
			}
			return
		}
		if err = svr.completeControlLogin(ctl, func() error {
			return writeWithDeadline(conn, connWriteTimeout, func() error {
				return acceptedConn.conn.WriteMsg(&msg.LoginResp{
					Version: version.Full(),
					RunID:   ctl.runID,
					Error:   "",
				})
			})
		}); err != nil {
			xl.Warnf("complete control login error: %v", err)
			svr.ctlManager.Remove(ctl)
			_ = ctl.Close()
			return
		}
	case *msg.NewWorkConn:
		if err := svr.RegisterWorkConn(
			acceptedConn.conn,
			m,
			acceptedConn.wireProtocol,
			acceptedConn.clientHelloPresent,
		); err != nil {
			_ = acceptedConn.conn.WriteMsg(&msg.StartWorkConn{
				Error: util.GenerateResponseErrorString("invalid NewWorkConn", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
			})
			conn.Close()
		}
	case *msg.NewVisitorConn:
		if err = svr.RegisterVisitorConn(conn, m, acceptedConn.wireProtocol); err != nil {
			xl.Warnf("register visitor conn error: %v", err)
			_ = acceptedConn.conn.WriteMsg(&msg.NewVisitorConnResp{
				ProxyName: m.ProxyName,
				Error:     util.GenerateResponseErrorString("register visitor conn error", err, lo.FromPtr(svr.cfg.DetailedErrorsToClient)),
			})
			conn.Close()
		} else {
			_ = acceptedConn.conn.WriteMsg(&msg.NewVisitorConnResp{
				ProxyName: m.ProxyName,
				Error:     "",
			})
		}
	default:
		log.Warnf("error message type for the new connection [%s]", conn.RemoteAddr().String())
		conn.Close()
	}
}

func (svr *Service) completeControlLogin(ctl *Control, writeSuccess func() error) error {
	committed, err := svr.ctlManager.completeLogin(ctl, writeSuccess)
	if err != nil {
		return err
	}
	if !committed {
		return errControlReplaced
	}
	return nil
}

type acceptedConnection struct {
	conn               *msg.Conn
	wireProtocol       string
	clientHelloPresent bool
	udpPacketCodec     string
	cryptoContext      *wire.CryptoContext
	firstMsg           msg.Message
}

func (svr *Service) acceptConnection(ctx context.Context, conn net.Conn) (*acceptedConnection, error) {
	_ = conn.SetReadDeadline(time.Now().Add(connReadTimeout))
	checkedConn, isV2, err := wire.CheckMagic(conn)
	if err != nil {
		return nil, fmt.Errorf("read wire protocol magic: %w", err)
	}

	wireProtocol := wire.ProtocolV1
	if isV2 {
		wireProtocol = wire.ProtocolV2
	}

	conn = netpkg.NewContextConn(ctx, checkedConn)
	acceptedConn := &acceptedConnection{wireProtocol: wireProtocol}
	if isV2 {
		wireConn := wire.NewConn(conn)
		rw := msg.NewV2ReadWriterWithConn(wireConn)
		acceptedConn.conn = msg.NewConn(conn, rw)
		acceptedConn.firstMsg, err = acceptedConn.readFirstV2Msg(conn, wireConn)
	} else {
		rw := msg.NewV1ReadWriter(conn)
		acceptedConn.conn = msg.NewConn(conn, rw)
		acceptedConn.firstMsg, err = acceptedConn.conn.ReadMsg()
	}
	if err != nil {
		return nil, err
	}
	_ = conn.SetReadDeadline(time.Time{})
	return acceptedConn, nil
}

func writeWithDeadline(conn net.Conn, timeout time.Duration, writeFn func() error) error {
	_ = conn.SetWriteDeadline(time.Now().Add(timeout))
	defer func() {
		_ = conn.SetWriteDeadline(time.Time{})
	}()
	return writeFn()
}

func (ac *acceptedConnection) messageConnFor(rw io.ReadWriter) *msg.Conn {
	return msg.NewConn(ac.conn, msg.NewReadWriter(rw, ac.wireProtocol))
}

func (ac *acceptedConnection) newControlReadWriter(rw io.ReadWriter, key []byte) (io.ReadWriter, error) {
	if ac.wireProtocol == wire.ProtocolV2 {
		if ac.cryptoContext == nil {
			return nil, fmt.Errorf("missing v2 crypto negotiation")
		}
		return netpkg.NewAEADCryptoReadWriter(
			rw,
			key,
			netpkg.AEADCryptoRoleServer,
			ac.cryptoContext.Algorithm,
			ac.cryptoContext.TranscriptHash,
		)
	}
	return netpkg.NewCryptoReadWriter(rw, key)
}

func (ac *acceptedConnection) readFirstV2Msg(conn net.Conn, wireConn *wire.Conn) (msg.Message, error) {
	frame, err := wireConn.ReadFrame()
	if err != nil {
		return nil, fmt.Errorf("read v2 frame: %w", err)
	}
	if frame.Type == wire.FrameTypeClientHello {
		ac.clientHelloPresent = true
		if err := ac.handleClientHello(conn, wireConn, frame); err != nil {
			return nil, err
		}
		frame, err = wireConn.ReadFrame()
		if err != nil {
			return nil, fmt.Errorf("read first v2 message frame: %w", err)
		}
	}

	m, err := msg.DecodeV2MessageFrame(frame)
	if err != nil {
		return nil, fmt.Errorf("decode v2 message: %w", err)
	}
	return m, nil
}

func (ac *acceptedConnection) handleClientHello(conn net.Conn, wireConn *wire.Conn, frame *wire.Frame) error {
	var hello wire.ClientHello
	if err := wireConn.UnmarshalFrame(frame, &hello); err != nil {
		return fmt.Errorf("decode ClientHello: %w", err)
	}

	serverHello, err := wire.NewServerHello(hello)
	if err != nil {
		serverHello = wire.DefaultServerHello()
		serverHello.Error = err.Error()
		if writeErr := writeWithDeadline(conn, connWriteTimeout, func() error {
			return wireConn.WriteJSONFrame(wire.FrameTypeServerHello, serverHello)
		}); writeErr != nil {
			return fmt.Errorf("%w; write ServerHello error: %v", err, writeErr)
		}
		return err
	}
	serverHelloFrame, err := wire.NewJSONFrame(wire.FrameTypeServerHello, serverHello)
	if err != nil {
		return fmt.Errorf("encode ServerHello: %w", err)
	}
	cryptoContext := wire.NewCryptoContext(
		serverHello.Selected.Crypto.Algorithm,
		frame.Payload,
		serverHelloFrame.Payload,
	)
	if err := writeWithDeadline(conn, connWriteTimeout, func() error {
		return wireConn.WriteFrame(serverHelloFrame)
	}); err != nil {
		return fmt.Errorf("write ServerHello: %w", err)
	}
	ac.cryptoContext = cryptoContext
	ac.udpPacketCodec = serverHello.Selected.Message.UDPPacketCodec
	return nil
}

// HandleListener accepts connections from client and call handleConnection to handle them.
// If internal is true, it means that this listener is used for internal communication like ssh tunnel gateway.
// TODO(fatedier): Pass some parameters of listener/connection through context to avoid passing too many parameters.
func (svr *Service) HandleListener(l net.Listener, internal bool) {
	// Listen for incoming connections from client.
	for {
		c, err := l.Accept()
		if err != nil {
			log.Warnf("listener for incoming connections from client closed")
			return
		}
		// inject xlog object into net.Conn context
		xl := xlog.New()
		ctx := context.Background()

		c = netpkg.NewContextConn(xlog.NewContext(ctx, xl), c)

		if !internal {
			log.Tracef("start check TLS connection...")
			originConn := c
			forceTLS := svr.cfg.Transport.TLS.Force
			var isTLS, custom bool
			c, isTLS, custom, err = netpkg.CheckAndEnableTLSServerConnWithTimeout(c, svr.tlsConfig, forceTLS, connReadTimeout)
			if err != nil {
				log.Warnf("checkAndEnableTLSServerConnWithTimeout error: %v", err)
				originConn.Close()
				continue
			}
			log.Tracef("check TLS connection success, isTLS: %v custom: %v internal: %v", isTLS, custom, internal)
		}

		// Start a new goroutine to handle connection.
		go func(ctx context.Context, frpConn net.Conn) {
			if lo.FromPtr(svr.cfg.Transport.TCPMux) && !internal {
				fmuxCfg := fmux.DefaultConfig()
				fmuxCfg.KeepAliveInterval = time.Duration(svr.cfg.Transport.TCPMuxKeepaliveInterval) * time.Second
				// Use trace level for yamux logs
				fmuxCfg.LogOutput = xlog.NewTraceWriter(xlog.FromContextSafe(ctx))
				fmuxCfg.MaxStreamWindowSize = 6 * 1024 * 1024
				session, err := fmux.Server(frpConn, fmuxCfg)
				if err != nil {
					log.Warnf("failed to create mux connection: %v", err)
					frpConn.Close()
					return
				}

				for {
					stream, err := session.AcceptStream()
					if err != nil {
						log.Debugf("accept new mux stream error: %v", err)
						session.Close()
						return
					}
					go svr.handleConnection(ctx, stream, internal, session)
				}
			} else {
				svr.handleConnection(ctx, frpConn, internal, nil)
			}
		}(ctx, c)
	}
}

func (svr *Service) HandleQUICListener(l *quic.Listener) {
	// Listen for incoming connections from client.
	for {
		c, err := l.Accept(context.Background())
		if err != nil {
			log.Warnf("quic listener for incoming connections from client closed")
			return
		}
		// Start a new goroutine to handle connection.
		go func(ctx context.Context, frpConn *quic.Conn) {
			for {
				stream, err := frpConn.AcceptStream(context.Background())
				if err != nil {
					log.Debugf("accept new quic mux stream error: %v", err)
					_ = frpConn.CloseWithError(0, "")
					return
				}
				go svr.handleConnection(ctx, netpkg.QuicStreamToNetConn(stream, frpConn), false, nil)
			}
		}(context.Background(), c)
	}
}

func (svr *Service) RegisterControl(
	ctlConn *msg.Conn,
	loginMsg *msg.Login,
	internal bool,
	wireProtocol string,
	udpPacketCodec string,
	session io.Closer,
) (*Control, error) {
	switch wireProtocol {
	case wire.ProtocolV1:
		if udpPacketCodec != "" {
			return nil, fmt.Errorf("UDP packet codec %q requires wire protocol v2", udpPacketCodec)
		}
	case wire.ProtocolV2:
		if udpPacketCodec != "" && udpPacketCodec != wire.UDPPacketCodecBinary {
			return nil, fmt.Errorf("unsupported UDP packet codec selection: %s", udpPacketCodec)
		}
	default:
		return nil, fmt.Errorf("unsupported wire protocol: %s", wireProtocol)
	}
	// If client's RunID is empty, it's a new client, we just create a new controller.
	// Otherwise, we check if there is one controller has the same run id. If so, we release previous controller and start new one.
	var err error
	if loginMsg.RunID == "" {
		loginMsg.RunID, err = util.RandID()
		if err != nil {
			return nil, err
		}
	}
	if err := validation.ValidateRunID(loginMsg.RunID); err != nil {
		return nil, fmt.Errorf("invalid run id: %w", err)
	}

	ctx := netpkg.NewContextFromConn(ctlConn)
	xl := xlog.FromContextSafe(ctx)
	xl.AppendPrefix(loginMsg.RunID)
	ctx = xlog.NewContext(ctx, xl)
	xl.Infof("client login info: ip [%s] version [%s] hostname [%s] os [%s] arch [%s]",
		ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)

	// Check auth.
	authVerifier := svr.auth.Verifier
	if internal && loginMsg.ClientSpec.AlwaysAuthPass {
		authVerifier = auth.AlwaysPassVerifier
	}
	if err := authVerifier.VerifyLogin(loginMsg); err != nil {
		return nil, err
	}

	ctl, err := NewControl(ctx, &SessionContext{
		RC:             svr.rc,
		PxyManager:     svr.pxyManager,
		PluginManager:  svr.pluginManager,
		AuthVerifier:   authVerifier,
		EncryptionKey:  svr.auth.EncryptionKey(),
		Conn:           ctlConn,
		Session:        session,
		LoginMsg:       loginMsg,
		ServerCfg:      svr.cfg,
		WireProtocol:   wireProtocol,
		UDPPacketCodec: udpPacketCodec,
	})
	if err != nil {
		xl.Warnf("create new controller error: %v", err)
		// don't return detailed errors to client
		return nil, fmt.Errorf("unexpected error when creating new controller")
	}

	if err := svr.ctlManager.Add(ctl); err != nil {
		return ctl, err
	}
	ctl.WaitForHandoff()

	active, err := svr.ctlManager.Activate(ctl)
	if err != nil {
		return ctl, err
	}
	if !active {
		return ctl, errControlReplaced
	}

	return ctl, nil
}

// RegisterWorkConn register a new work connection to control and proxies need it.
func (svr *Service) RegisterWorkConn(
	workConn *msg.Conn,
	newMsg *msg.NewWorkConn,
	workWireProtocol string,
	workClientHelloPresent bool,
) error {
	if workClientHelloPresent {
		return fmt.Errorf("ClientHello is not allowed on work connections")
	}
	xl := netpkg.NewLogFromConn(workConn)
	ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
	if !exist {
		xl.Warnf("no client control found for run id [%s]", newMsg.RunID)
		return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
	}
	if workWireProtocol != ctl.sessionCtx.WireProtocol {
		return fmt.Errorf("work connection wire protocol mismatch: got %s want %s", workWireProtocol, ctl.sessionCtx.WireProtocol)
	}

	// server plugin hook
	content := &plugin.NewWorkConnContent{
		User: plugin.UserInfo{
			User:  ctl.sessionCtx.LoginMsg.User,
			Metas: ctl.sessionCtx.LoginMsg.Metas,
			RunID: ctl.sessionCtx.LoginMsg.RunID,
		},
		NewWorkConn: *newMsg,
	}
	retContent, err := svr.pluginManager.NewWorkConn(content)
	if err == nil {
		newMsg = &retContent.NewWorkConn
		// Check auth.
		err = ctl.sessionCtx.AuthVerifier.VerifyNewWorkConn(newMsg)
	}
	if err != nil {
		xl.Warnf("invalid NewWorkConn with run id [%s]", newMsg.RunID)
		return err
	}
	return svr.ctlManager.RegisterWorkConn(ctl, proxy.NewWorkConn(workConn))
}

func (svr *Service) RegisterVisitorConn(visitorConn net.Conn, newMsg *msg.NewVisitorConn, wireProtocol string) error {
	admit := func(visitorUser, visitorWireProtocol, visitorUDPPacketCodec string) error {
		if visitorWireProtocol == "" {
			visitorWireProtocol = wireProtocol
		}
		return svr.rc.VisitorManager.NewConn(newMsg.ProxyName, visitorConn, newMsg.Timestamp, newMsg.SignKey,
			newMsg.UseEncryption, newMsg.UseCompression, visitorUser, visitorWireProtocol, visitorUDPPacketCodec)
	}
	// TODO(deprecation): Compatible with old versions, can be without runID, user is empty. In later versions, it will be mandatory to include runID.
	// If runID is required, it is not compatible with versions prior to v0.50.0.
	if newMsg.RunID != "" {
		admitted, err := svr.ctlManager.admitVisitorByRunID(newMsg.RunID, func(visitorUser, controlWireProtocol, controlUDPPacketCodec string) error {
			if wireProtocol != controlWireProtocol {
				return fmt.Errorf("visitor connection wire protocol mismatch: got %s want %s", wireProtocol, controlWireProtocol)
			}
			return admit(visitorUser, controlWireProtocol, controlUDPPacketCodec)
		})
		if err != nil {
			return err
		}
		if !admitted {
			return fmt.Errorf("no client control found for run id [%s]", newMsg.RunID)
		}
		return nil
	}
	return admit("", wireProtocol, "")
}