mirror of
https://github.com/fatedier/frp.git
synced 2026-01-11 22:23:12 +00:00
optimize some code (#3801)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user