perf: 支持 mongodb 配置 auth_source

jumpserver · Mar 12, 2024 · fe89922 · fe89922
1 parent a95606a
commit fe89922
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 5 deletions.
diff --git a/pkg/jms-sdk-go/model/platform.go b/pkg/jms-sdk-go/model/platform.go
@@ -60,6 +60,9 @@ type ProtocolSetting struct {
 	TelnetUsernamePrompt string `json:"username_prompt"`
 	TelnetPasswordPrompt string `json:"password_prompt"`
 	TelnetSuccessPrompt  string `json:"success_prompt"`
+
+	// for mongodb
+	AuthSource string `json:"auth_source"`
 }
 
 type Protocol struct {

diff --git a/pkg/proxy/server.go b/pkg/proxy/server.go
@@ -589,7 +589,9 @@ func (s *Server) getMongoDBConn(localTunnelAddr *net.TCPAddr) (srvConn *srvconn.
 		host = "127.0.0.1"
 		port = localTunnelAddr.Port
 	}
-
+	platform := s.connOpts.authInfo.Platform
+	protocolSetting := platform.GetProtocol("mongodb")
+	authSource := protocolSetting.Setting.AuthSource
 	srvConn, err = srvconn.NewMongoDBConnection(
 		srvconn.SqlHost(host),
 		srvconn.SqlPort(port),
@@ -600,6 +602,7 @@ func (s *Server) getMongoDBConn(localTunnelAddr *net.TCPAddr) (srvConn *srvconn.
 		srvconn.SqlCaCert(asset.SecretInfo.CaCert),
 		srvconn.SqlCertKey(asset.SecretInfo.ClientKey),
 		srvconn.SqlAllowInvalidCert(asset.SpecInfo.AllowInvalidCert),
+		srvconn.SqlAuthSource(authSource),
 		srvconn.SqlPtyWin(srvconn.Windows{
 			Width:  s.UserConn.Pty().Window.Width,
 			Height: s.UserConn.Pty().Window.Height,

diff --git a/pkg/srvconn/conn_mongodb.go b/pkg/srvconn/conn_mongodb.go
@@ -130,10 +130,19 @@ func addMongoParamsWithSSL(args *sqlOption, params map[string]string) {
 	}
 }
 
+func (opt *sqlOption) GetAuthSource() string {
+	//  authSource 默认是 admin，通过 platform 的 protocol 设置，修改这个认证的值
+	// https://www.mongodb.com/docs/manual/reference/connection-string/#mongodb-urioption-urioption.authSource
+	if opt.AuthSource == "" {
+		return "admin"
+	}
+	return opt.AuthSource
+}
+
 func (opt *sqlOption) MongoDBCommandArgs() []string {
 	host := net.JoinHostPort(opt.Host, strconv.Itoa(opt.Port))
 	params := map[string]string{
-		"authSource": "admin",
+		"authSource": opt.GetAuthSource(),
 	}
 	addMongoParamsWithSSL(opt, params)
 	uri := BuildMongoDBURI(
@@ -149,10 +158,8 @@ func (opt *sqlOption) MongoDBCommandArgs() []string {
 
 func checkMongoDBAccount(args *sqlOption) error {
 	host := net.JoinHostPort(args.Host, strconv.Itoa(args.Port))
-	// todo: authSource 暂且只使用 admin， 待后续可配置后，修改这个认证的值
-	// https://www.mongodb.com/docs/manual/reference/connection-string/#mongodb-urioption-urioption.authSource
 	params := map[string]string{
-		"authSource": "admin",
+		"authSource": args.GetAuthSource(),
 		"connect":    "direct",
 	}
 	addMongoParamsWithSSL(args, params)

diff --git a/pkg/srvconn/conn_sql_opt.go b/pkg/srvconn/conn_sql_opt.go
@@ -23,6 +23,8 @@ type sqlOption struct {
 	win Windows
 
 	disableMySQLAutoRehash bool
+
+	AuthSource string
 }
 
 type SqlOption func(*sqlOption)
@@ -93,6 +95,12 @@ func SqlPtyWin(win Windows) SqlOption {
 	}
 }
 
+func SqlAuthSource(authSource string) SqlOption {
+	return func(args *sqlOption) {
+		args.AuthSource = authSource
+	}
+}
+
 const (
 	maxSQLConnCount = 1
 	maxIdleTime     = time.Second * 15