optimize some code (#3801)

This commit is contained in:
fatedier
2023-11-27 15:47:49 +08:00
committed by GitHub
parent d5b41f1e14
commit 69ae2b0b69
52 changed files with 880 additions and 600 deletions

View File

@@ -30,7 +30,6 @@ import (
quic "github.com/quic-go/quic-go"
"github.com/samber/lo"
"github.com/fatedier/frp/assets"
"github.com/fatedier/frp/pkg/auth"
v1 "github.com/fatedier/frp/pkg/config/v1"
modelmetrics "github.com/fatedier/frp/pkg/metrics"
@@ -39,8 +38,9 @@ import (
plugin "github.com/fatedier/frp/pkg/plugin/server"
"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"
utilnet "github.com/fatedier/frp/pkg/util/net"
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"
@@ -79,7 +79,8 @@ type Service struct {
// Accept frp tls connections
tlsListener net.Listener
virtualListener *utilnet.InternalListener
// Accept pipe connections from ssh tunnel gateway
sshTunnelListener *netpkg.InternalListener
// Manage all controllers
ctlManager *ControlManager
@@ -96,6 +97,9 @@ type Service struct {
// All resource managers and controllers
rc *controller.ResourceController
// web server for dashboard UI and apis
webServer *httppkg.Server
sshTunnelGateway *ssh.Gateway
// Verifies authentication based on selected method
@@ -111,16 +115,30 @@ type Service struct {
cancel context.CancelFunc
}
func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
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
return nil, err
}
svr = &Service{
var webServer *httppkg.Server
if cfg.WebServer.Port > 0 {
ws, err := httppkg.NewServer(cfg.WebServer)
if err != nil {
return nil, err
}
webServer = ws
modelmetrics.EnableMem()
if cfg.EnablePrometheus {
modelmetrics.EnablePrometheus()
}
}
svr := &Service{
ctlManager: NewControlManager(),
pxyManager: proxy.NewManager(),
pluginManager: plugin.NewManager(),
@@ -129,12 +147,16 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
TCPPortManager: ports.NewManager("tcp", cfg.ProxyBindAddr, cfg.AllowPorts),
UDPPortManager: ports.NewManager("udp", cfg.ProxyBindAddr, cfg.AllowPorts),
},
virtualListener: utilnet.NewInternalListener(),
httpVhostRouter: vhost.NewRouters(),
authVerifier: auth.NewAuthVerifier(cfg.Auth),
tlsConfig: tlsConfig,
cfg: cfg,
ctx: context.Background(),
sshTunnelListener: netpkg.NewInternalListener(),
httpVhostRouter: vhost.NewRouters(),
authVerifier: auth.NewAuthVerifier(cfg.Auth),
webServer: webServer,
tlsConfig: tlsConfig,
cfg: cfg,
ctx: context.Background(),
}
if webServer != nil {
webServer.RouteRegister(svr.registerRouteHandlers)
}
// Create tcpmux httpconnect multiplexer.
@@ -143,14 +165,12 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.TCPMuxHTTPConnectPort))
l, err = net.Listen("tcp", address)
if err != nil {
err = fmt.Errorf("create server listener error, %v", err)
return
return nil, fmt.Errorf("create server listener error, %v", err)
}
svr.rc.TCPMuxHTTPConnectMuxer, err = tcpmux.NewHTTPConnectTCPMuxer(l, cfg.TCPMuxPassthrough, vhostReadWriteTimeout)
if err != nil {
err = fmt.Errorf("create vhost tcpMuxer error, %v", err)
return
return nil, fmt.Errorf("create vhost tcpMuxer error, %v", err)
}
log.Info("tcpmux httpconnect multiplexer listen on %s, passthough: %v", address, cfg.TCPMuxPassthrough)
}
@@ -191,8 +211,7 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
address := net.JoinHostPort(cfg.BindAddr, strconv.Itoa(cfg.BindPort))
ln, err := net.Listen("tcp", address)
if err != nil {
err = fmt.Errorf("create server listener error, %v", err)
return
return nil, fmt.Errorf("create server listener error, %v", err)
}
svr.muxer = mux.NewMux(ln)
@@ -208,10 +227,9 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
// 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 = utilnet.ListenKcp(address)
svr.kcpListener, err = netpkg.ListenKcp(address)
if err != nil {
err = fmt.Errorf("listen on kcp udp address %s error: %v", address, err)
return
return nil, fmt.Errorf("listen on kcp udp address %s error: %v", address, err)
}
log.Info("frps kcp listen on udp %s", address)
}
@@ -226,28 +244,26 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
KeepAlivePeriod: time.Duration(cfg.Transport.QUIC.KeepalivePeriod) * time.Second,
})
if err != nil {
err = fmt.Errorf("listen on quic udp address %s error: %v", address, err)
return
return nil, fmt.Errorf("listen on quic udp address %s error: %v", address, err)
}
log.Info("frps quic listen on %s", address)
}
if cfg.SSHTunnelGateway.BindPort > 0 {
sshGateway, err := ssh.NewGateway(cfg.SSHTunnelGateway, cfg.ProxyBindAddr, svr.virtualListener)
sshGateway, err := ssh.NewGateway(cfg.SSHTunnelGateway, cfg.ProxyBindAddr, svr.sshTunnelListener)
if err != nil {
err = fmt.Errorf("create ssh gateway error: %v", err)
return nil, err
return nil, fmt.Errorf("create ssh gateway error: %v", err)
}
svr.sshTunnelGateway = sshGateway
log.Info("frps sshTunnelGateway listen on port %d", cfg.SSHTunnelGateway.BindPort)
}
// Listen for accepting connections from client using websocket protocol.
websocketPrefix := []byte("GET " + utilnet.FrpWebsocketPath)
websocketPrefix := []byte("GET " + netpkg.FrpWebsocketPath)
websocketLn := svr.muxer.Listen(0, uint32(len(websocketPrefix)), func(data []byte) bool {
return bytes.Equal(data, websocketPrefix)
})
svr.websocketListener = utilnet.NewWebsocketListener(websocketLn)
svr.websocketListener = netpkg.NewWebsocketListener(websocketLn)
// Create http vhost muxer.
if cfg.VhostHTTPPort > 0 {
@@ -267,8 +283,7 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
} else {
l, err = net.Listen("tcp", address)
if err != nil {
err = fmt.Errorf("create vhost http listener error, %v", err)
return
return nil, fmt.Errorf("create vhost http listener error, %v", err)
}
}
go func() {
@@ -286,55 +301,30 @@ func NewService(cfg *v1.ServerConfig) (svr *Service, err error) {
address := net.JoinHostPort(cfg.ProxyBindAddr, strconv.Itoa(cfg.VhostHTTPSPort))
l, err = net.Listen("tcp", address)
if err != nil {
err = fmt.Errorf("create server listener error, %v", err)
return
return nil, fmt.Errorf("create server listener error, %v", err)
}
log.Info("https service listen on %s", address)
}
svr.rc.VhostHTTPSMuxer, err = vhost.NewHTTPSMuxer(l, vhostReadWriteTimeout)
if err != nil {
err = fmt.Errorf("create vhost httpsMuxer error, %v", err)
return
return nil, fmt.Errorf("create vhost httpsMuxer error, %v", err)
}
}
// 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]) == utilnet.FRPTLSHeadByte || int(data[0]) == 0x16
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 {
err = fmt.Errorf("create nat hole controller error, %v", err)
return
return nil, fmt.Errorf("create nat hole controller error, %v", err)
}
svr.rc.NatHoleController = nc
var statsEnable bool
// Create dashboard web server.
if cfg.WebServer.Port > 0 {
// Init dashboard assets
assets.Load(cfg.WebServer.AssetsDir)
address := net.JoinHostPort(cfg.WebServer.Addr, strconv.Itoa(cfg.WebServer.Port))
err = svr.RunDashboardServer(address)
if err != nil {
err = fmt.Errorf("create dashboard web server error, %v", err)
return
}
log.Info("Dashboard listen on %s", address)
statsEnable = true
}
if statsEnable {
modelmetrics.EnableMem()
if cfg.EnablePrometheus {
modelmetrics.EnablePrometheus()
}
}
return
return svr, nil
}
func (svr *Service) Run(ctx context.Context) {
@@ -342,7 +332,17 @@ func (svr *Service) Run(ctx context.Context) {
svr.ctx = ctx
svr.cancel = cancel
go svr.HandleListener(svr.virtualListener, true)
// run dashboard web server.
if svr.webServer != nil {
go func() {
log.Info("dashboard listen on %s", svr.webServer.Address())
if err := svr.webServer.Run(); err != nil {
log.Warn("dashboard server exit with error: %v", err)
}
}()
}
go svr.HandleListener(svr.sshTunnelListener, true)
if svr.kcpListener != nil {
go svr.HandleListener(svr.kcpListener, false)
@@ -398,7 +398,7 @@ func (svr *Service) Close() error {
return nil
}
func (svr *Service) handleConnection(ctx context.Context, conn net.Conn) {
func (svr *Service) handleConnection(ctx context.Context, conn net.Conn, internal bool) {
xl := xlog.FromContextSafe(ctx)
var (
@@ -424,7 +424,7 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn) {
retContent, err := svr.pluginManager.Login(content)
if err == nil {
m = &retContent.Login
err = svr.RegisterControl(conn, m)
err = svr.RegisterControl(conn, m, internal)
}
// If login failed, send error message there.
@@ -461,6 +461,9 @@ func (svr *Service) handleConnection(ctx context.Context, conn net.Conn) {
}
}
// 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 {
@@ -473,19 +476,21 @@ func (svr *Service) HandleListener(l net.Listener, internal bool) {
xl := xlog.New()
ctx := context.Background()
c = utilnet.NewContextConn(xlog.NewContext(ctx, xl), c)
c = netpkg.NewContextConn(xlog.NewContext(ctx, xl), c)
log.Trace("start check TLS connection...")
originConn := c
forceTLS := svr.cfg.Transport.TLS.Force && !internal
var isTLS, custom bool
c, isTLS, custom, err = utilnet.CheckAndEnableTLSServerConnWithTimeout(c, svr.tlsConfig, forceTLS, connReadTimeout)
if err != nil {
log.Warn("CheckAndEnableTLSServerConnWithTimeout error: %v", err)
originConn.Close()
continue
if !internal {
log.Trace("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.Warn("CheckAndEnableTLSServerConnWithTimeout error: %v", err)
originConn.Close()
continue
}
log.Trace("check TLS connection success, isTLS: %v custom: %v internal: %v", isTLS, custom, internal)
}
log.Trace("check TLS connection success, isTLS: %v custom: %v", isTLS, custom)
// Start a new goroutine to handle connection.
go func(ctx context.Context, frpConn net.Conn) {
@@ -508,10 +513,10 @@ func (svr *Service) HandleListener(l net.Listener, internal bool) {
session.Close()
return
}
go svr.handleConnection(ctx, stream)
go svr.handleConnection(ctx, stream, internal)
}
} else {
svr.handleConnection(ctx, frpConn)
svr.handleConnection(ctx, frpConn, internal)
}
}(ctx, c)
}
@@ -534,13 +539,13 @@ func (svr *Service) HandleQUICListener(l *quic.Listener) {
_ = frpConn.CloseWithError(0, "")
return
}
go svr.handleConnection(ctx, utilnet.QuicStreamToNetConn(stream, frpConn))
go svr.handleConnection(ctx, netpkg.QuicStreamToNetConn(stream, frpConn), false)
}
}(context.Background(), c)
}
}
func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login) error {
func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login, internal bool) error {
// 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
@@ -551,7 +556,7 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login) error
}
}
ctx := utilnet.NewContextFromConn(ctlConn)
ctx := netpkg.NewContextFromConn(ctlConn)
xl := xlog.FromContextSafe(ctx)
xl.AppendPrefix(loginMsg.RunID)
ctx = xlog.NewContext(ctx, xl)
@@ -559,11 +564,16 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login) error
ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)
// Check auth.
if err := svr.authVerifier.VerifyLogin(loginMsg); err != nil {
authVerifier := svr.authVerifier
if internal && loginMsg.ClientSpec.AlwaysAuthPass {
authVerifier = auth.AlwaysPassVerifier
}
if err := authVerifier.VerifyLogin(loginMsg); err != nil {
return err
}
ctl, err := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, svr.authVerifier, ctlConn, loginMsg, svr.cfg)
// TODO(fatedier): use SessionContext
ctl, err := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, authVerifier, ctlConn, !internal, loginMsg, svr.cfg)
if err != nil {
xl.Warn("create new controller error: %v", err)
// don't return detailed errors to client
@@ -588,7 +598,7 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login) error
// RegisterWorkConn register a new work connection to control and proxies need it.
func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn) error {
xl := utilnet.NewLogFromConn(workConn)
xl := netpkg.NewLogFromConn(workConn)
ctl, exist := svr.ctlManager.GetByID(newMsg.RunID)
if !exist {
xl.Warn("No client control found for run id [%s]", newMsg.RunID)
@@ -607,7 +617,7 @@ func (svr *Service) RegisterWorkConn(workConn net.Conn, newMsg *msg.NewWorkConn)
if err == nil {
newMsg = &retContent.NewWorkConn
// Check auth.
err = svr.authVerifier.VerifyNewWorkConn(newMsg)
err = ctl.authVerifier.VerifyNewWorkConn(newMsg)
}
if err != nil {
xl.Warn("invalid NewWorkConn with run id [%s]", newMsg.RunID)