refactor: separate auth config from runtime and defer token resolution (#5105)

server/control.go

@@ -106,6 +106,8 @@ type Control struct {
 
 	// verifies authentication based on selected method
 	authVerifier auth.Verifier
+	// key used for connection encryption
+	encryptionKey []byte
 
 	// other components can use this to communicate with client
 	msgTransporter transport.MessageTransporter
@@ -157,6 +159,7 @@ func NewControl(
 	pxyManager *proxy.Manager,
 	pluginManager *plugin.Manager,
 	authVerifier auth.Verifier,
+	encryptionKey []byte,
 	ctlConn net.Conn,
 	ctlConnEncrypted bool,
 	loginMsg *msg.Login,
@@ -171,6 +174,7 @@ func NewControl(
 		pxyManager:    pxyManager,
 		pluginManager: pluginManager,
 		authVerifier:  authVerifier,
+		encryptionKey: encryptionKey,
 		conn:          ctlConn,
 		loginMsg:      loginMsg,
 		workConnCh:    make(chan net.Conn, poolCount+10),
@@ -186,7 +190,7 @@ func NewControl(
 	ctl.lastPing.Store(time.Now())
 
 	if ctlConnEncrypted {
-		cryptoRW, err := netpkg.NewCryptoReadWriter(ctl.conn, []byte(ctl.serverCfg.Auth.Token))
+		cryptoRW, err := netpkg.NewCryptoReadWriter(ctl.conn, ctl.encryptionKey)
 		if err != nil {
 			return nil, err
 		}
@@ -478,6 +482,7 @@ func (ctl *Control) RegisterProxy(pxyMsg *msg.NewProxy) (remoteAddr string, err
 		GetWorkConnFn:      ctl.GetWorkConn,
 		Configurer:         pxyConf,
 		ServerCfg:          ctl.serverCfg,
+		EncryptionKey:      ctl.encryptionKey,
 	})
 	if err != nil {
 		return remoteAddr, err

server/proxy/http.go

@@ -165,7 +165,7 @@ func (pxy *HTTPProxy) GetRealConn(remoteAddr string) (workConn net.Conn, err err
 
 	var rwc io.ReadWriteCloser = tmpConn
 	if pxy.cfg.Transport.UseEncryption {
-		rwc, err = libio.WithEncryption(rwc, []byte(pxy.serverCfg.Auth.Token))
+		rwc, err = libio.WithEncryption(rwc, pxy.encryptionKey)
 		if err != nil {
 			xl.Errorf("create encryption stream error: %v", err)
 			return

server/proxy/proxy.go

@@ -68,6 +68,7 @@ type BaseProxy struct {
 	poolCount     int
 	getWorkConnFn GetWorkConnFn
 	serverCfg     *v1.ServerConfig
+	encryptionKey []byte
 	limiter       *rate.Limiter
 	userInfo      plugin.UserInfo
 	loginMsg      *msg.Login
@@ -213,7 +214,6 @@ func (pxy *BaseProxy) handleUserTCPConnection(userConn net.Conn) {
 	xl := xlog.FromContextSafe(pxy.Context())
 	defer userConn.Close()
 
-	serverCfg := pxy.serverCfg
 	cfg := pxy.configurer.GetBaseConfig()
 	// server plugin hook
 	rc := pxy.GetResourceController()
@@ -240,7 +240,7 @@ func (pxy *BaseProxy) handleUserTCPConnection(userConn net.Conn) {
 	xl.Tracef("handler user tcp connection, use_encryption: %t, use_compression: %t",
 		cfg.Transport.UseEncryption, cfg.Transport.UseCompression)
 	if cfg.Transport.UseEncryption {
-		local, err = libio.WithEncryption(local, []byte(serverCfg.Auth.Token))
+		local, err = libio.WithEncryption(local, pxy.encryptionKey)
 		if err != nil {
 			xl.Errorf("create encryption stream error: %v", err)
 			return
@@ -279,6 +279,7 @@ type Options struct {
 	GetWorkConnFn      GetWorkConnFn
 	Configurer         v1.ProxyConfigurer
 	ServerCfg          *v1.ServerConfig
+	EncryptionKey      []byte
 }
 
 func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) {
@@ -298,6 +299,7 @@ func NewProxy(ctx context.Context, options *Options) (pxy Proxy, err error) {
 		poolCount:     options.PoolCount,
 		getWorkConnFn: options.GetWorkConnFn,
 		serverCfg:     options.ServerCfg,
+		encryptionKey: options.EncryptionKey,
 		limiter:       limiter,
 		xl:            xl,
 		ctx:           xlog.NewContext(ctx, xl),

server/proxy/udp.go

@@ -205,7 +205,7 @@ func (pxy *UDPProxy) Run() (remoteAddr string, err error) {
 
 			var rwc io.ReadWriteCloser = workConn
 			if pxy.cfg.Transport.UseEncryption {
-				rwc, err = libio.WithEncryption(rwc, []byte(pxy.serverCfg.Auth.Token))
+				rwc, err = libio.WithEncryption(rwc, pxy.encryptionKey)
 				if err != nil {
 					xl.Errorf("create encryption stream error: %v", err)
 					workConn.Close()

server/service.go

@@ -113,8 +113,8 @@ type Service struct {
 
 	sshTunnelGateway *ssh.Gateway
 
-	// Verifies authentication based on selected method
-	authVerifier auth.Verifier
+	// Auth runtime and encryption materials
+	auth *auth.ServerAuth
 
 	tlsConfig *tls.Config
 
@@ -149,6 +149,11 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) {
 		}
 	}
 
+	authRuntime, err := auth.BuildServerAuth(&cfg.Auth)
+	if err != nil {
+		return nil, err
+	}
+
 	svr := &Service{
 		ctlManager:    NewControlManager(),
 		pxyManager:    proxy.NewManager(),
@@ -160,7 +165,7 @@ func NewService(cfg *v1.ServerConfig) (*Service, error) {
 		},
 		sshTunnelListener: netpkg.NewInternalListener(),
 		httpVhostRouter:   vhost.NewRouters(),
-		authVerifier:      auth.NewAuthVerifier(cfg.Auth),
+		auth:              authRuntime,
 		webServer:         webServer,
 		tlsConfig:         tlsConfig,
 		cfg:               cfg,
@@ -586,7 +591,7 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login, inter
 		ctlConn.RemoteAddr().String(), loginMsg.Version, loginMsg.Hostname, loginMsg.Os, loginMsg.Arch)
 
 	// Check auth.
-	authVerifier := svr.authVerifier
+	authVerifier := svr.auth.Verifier
 	if internal && loginMsg.ClientSpec.AlwaysAuthPass {
 		authVerifier = auth.AlwaysPassVerifier
 	}
@@ -595,7 +600,7 @@ func (svr *Service) RegisterControl(ctlConn net.Conn, loginMsg *msg.Login, inter
 	}
 
 	// TODO(fatedier): use SessionContext
-	ctl, err := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, authVerifier, ctlConn, !internal, loginMsg, svr.cfg)
+	ctl, err := NewControl(ctx, svr.rc, svr.pxyManager, svr.pluginManager, authVerifier, svr.auth.EncryptionKey(), ctlConn, !internal, loginMsg, svr.cfg)
 	if err != nil {
 		xl.Warnf("create new controller error: %v", err)
 		// don't return detailed errors to client