Merge remote-tracking branch 'upstream/dev' into fix/882-sqlserver-ddl-unsupported

2026-08-15 11:14:31 +08:00 · 2026-08-09 13:33:05 +08:00
parent 52488cf3f7 f74ebb05a5
commit 664c7a05d3
4 changed files with 159 additions and 102 deletions
--- a/internal/db/mongodb_impl.go
+++ b/internal/db/mongodb_impl.go
@@ -28,7 +28,6 @@ type MongoDB struct {
 	client      *mongo.Client
 	database    string
 	pingTimeout time.Duration
-	forwarder   *ssh.LocalForwarder
 }

 type mongoProxyDialer struct {
@@ -39,6 +38,25 @@ func (d *mongoProxyDialer) DialContext(ctx context.Context, network, address str
 	return proxytunnel.DialContext(ctx, d.proxyConfig, network, address)
 }

+type mongoSSHDialer struct {
+	sshConfig   connection.SSHConfig
+	dialContext func(context.Context, connection.SSHConfig, string, string) (net.Conn, error)
+}
+
+func (d *mongoSSHDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	return d.dialContext(ctx, d.sshConfig, network, address)
+}
+
+func mongoConnectionDialer(config connection.ConnectionConfig) options.ContextDialer {
+	if config.UseSSH {
+		return &mongoSSHDialer{sshConfig: config.SSH, dialContext: ssh.DialContextThroughSSH}
+	}
+	if config.UseProxy {
+		return &mongoProxyDialer{proxyConfig: config.Proxy}
+	}
+	return nil
+}
+
 const defaultMongoPort = 27017
 const mongoObjectIDLocatorColumn = "__gonavi_mongodb_id_locator__"

@@ -335,48 +353,9 @@ func (m *MongoDB) Connect(config connection.ConnectionConfig) (err error) {
 	runConfig := applyMongoURI(config)
 	connectConfig := runConfig
 	sshRouteHint := ""
-
-	if runConfig.UseSSH && runConfig.MongoSRV {
-		return fmt.Errorf("MongoDB SRV 记录模式暂不支持 SSH 隧道")
-	}
-
 	if runConfig.UseSSH {
-		seeds := collectMongoSeeds(runConfig)
-		if len(seeds) == 0 {
-			seeds = append(seeds, normalizeMongoAddress(runConfig.Host, runConfig.Port))
-		}
-		targetHost, targetPort, ok := parseHostPortWithDefault(seeds[0], defaultMongoPort)
-		if !ok {
-			return fmt.Errorf("MongoDB 连接失败：无效地址 %s", seeds[0])
-		}
-
-		logger.Infof("MongoDB 使用 SSH 连接：地址=%s:%d", targetHost, targetPort)
-
-		forwarder, err := ssh.AcquireLocalForwarder(runConfig.SSH, targetHost, targetPort)
-		if err != nil {
-			return fmt.Errorf("创建 SSH 隧道失败：%w", err)
-		}
-		m.forwarder = forwarder
-
-		host, portStr, err := net.SplitHostPort(forwarder.LocalAddr)
-		if err != nil {
-			return fmt.Errorf("解析本地转发地址失败：%w", err)
-		}
-
-		port, err := strconv.Atoi(portStr)
-		if err != nil {
-			return fmt.Errorf("解析本地端口失败：%w", err)
-		}
-
-		localConfig := runConfig
-		localConfig.Host = host
-		localConfig.Port = port
-		localConfig.UseSSH = false
-		localConfig.URI = ""
-		localConfig.Hosts = []string{normalizeMongoAddress(host, port)}
-		connectConfig = localConfig
-		sshRouteHint = fmt.Sprintf("SSH隧道 %s -> %s:%d", forwarder.LocalAddr, targetHost, targetPort)
-		logger.Infof("MongoDB 通过本地端口转发连接：%s -> %s:%d", forwarder.LocalAddr, targetHost, targetPort)
+		sshRouteHint = fmt.Sprintf("SSH隧道 %s:%d", runConfig.SSH.Host, runConfig.SSH.Port)
+		logger.Infof("MongoDB 使用 SSH 隧道连接：跳板=%s:%d", runConfig.SSH.Host, runConfig.SSH.Port)
 	}

 	m.pingTimeout = getConnectTimeout(connectConfig)
@@ -431,8 +410,8 @@ func (m *MongoDB) Connect(config connection.ConnectionConfig) (err error) {
 			if tlsConfig != nil {
 				clientOpts.SetTLSConfig(tlsConfig)
 			}
-			if attemptConfig.UseProxy {
-				clientOpts.SetDialer(&mongoProxyDialer{proxyConfig: attemptConfig.Proxy})
+			if dialer := mongoConnectionDialer(attemptConfig); dialer != nil {
+				clientOpts.SetDialer(dialer)
 			}
 			client, err := mongo.Connect(clientOpts)
 			if err != nil {
@@ -478,13 +457,6 @@ func (m *MongoDB) Connect(config connection.ConnectionConfig) (err error) {
 }

 func (m *MongoDB) Close() error {
-	if m.forwarder != nil {
-		if err := m.forwarder.Release(); err != nil {
-			logger.Warnf("关闭 MongoDB SSH 端口转发失败：%v", err)
-		}
-		m.forwarder = nil
-	}
-
 	if m.client != nil {
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
--- a/internal/db/mongodb_impl_uri_test.go
+++ b/internal/db/mongodb_impl_uri_test.go
@@ -3,6 +3,9 @@
 package db

 import (
+	"context"
+	"errors"
+	"net"
 	"strings"
 	"testing"

@@ -11,6 +14,59 @@ import (
 	"go.mongodb.org/mongo-driver/v2/bson"
 )

+func TestMongoSSHDialerRoutesAllMembersThroughSSH(t *testing.T) {
+	tests := []struct {
+		name       string
+		uri        string
+		wantScheme string
+	}{
+		{name: "standard", uri: "mongodb://mongo.internal:27017/app", wantScheme: "mongodb://"},
+		{name: "srv", uri: "mongodb+srv://cluster.example.test/app", wantScheme: "mongodb+srv://"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := applyMongoURI(connection.ConnectionConfig{
+				URI:    tt.uri,
+				UseSSH: true,
+				SSH: connection.SSHConfig{
+					Host: "bastion.example.test",
+					Port: 22,
+					User: "operator",
+				},
+			})
+			if uri := (&MongoDB{}).getURI(config); !strings.HasPrefix(uri, tt.wantScheme) {
+				t.Fatalf("expected URI scheme %q, got %q", tt.wantScheme, uri)
+			}
+
+			dialer := mongoConnectionDialer(config)
+			sshDialer, ok := dialer.(*mongoSSHDialer)
+			if !ok {
+				t.Fatalf("expected SSH dialer, got %T", dialer)
+			}
+
+			wantErr := errors.New("dial stopped")
+			var addresses []string
+			sshDialer.dialContext = func(_ context.Context, _ connection.SSHConfig, network, address string) (net.Conn, error) {
+				if network != "tcp" {
+					t.Fatalf("expected tcp network, got %q", network)
+				}
+				addresses = append(addresses, address)
+				return nil, wantErr
+			}
+
+			for _, address := range []string{"mongo-1.internal:27017", "mongo-2.internal:27018"} {
+				if _, err := dialer.DialContext(context.Background(), "tcp", address); !errors.Is(err, wantErr) {
+					t.Fatalf("expected SSH dial error for %s, got %v", address, err)
+				}
+			}
+			if strings.Join(addresses, ",") != "mongo-1.internal:27017,mongo-2.internal:27018" {
+				t.Fatalf("unexpected SSH targets: %v", addresses)
+			}
+		})
+	}
+}
+
 func TestApplyMongoURI_ExplicitHostDoesNotAdoptURIHosts(t *testing.T) {
 	config := connection.ConnectionConfig{
 		Host: "10.10.10.10",
--- a/internal/db/mongodb_impl_v1.go
+++ b/internal/db/mongodb_impl_v1.go
@@ -29,7 +29,6 @@ type MongoDBV1 struct {
 	client      *mongo.Client
 	database    string
 	pingTimeout time.Duration
-	forwarder   *ssh.LocalForwarder
 }

 type mongoProxyDialer struct {
@@ -40,6 +39,25 @@ func (d *mongoProxyDialer) DialContext(ctx context.Context, network, address str
 	return proxytunnel.DialContext(ctx, d.proxyConfig, network, address)
 }

+type mongoSSHDialer struct {
+	sshConfig   connection.SSHConfig
+	dialContext func(context.Context, connection.SSHConfig, string, string) (net.Conn, error)
+}
+
+func (d *mongoSSHDialer) DialContext(ctx context.Context, network, address string) (net.Conn, error) {
+	return d.dialContext(ctx, d.sshConfig, network, address)
+}
+
+func mongoConnectionDialer(config connection.ConnectionConfig) options.ContextDialer {
+	if config.UseSSH {
+		return &mongoSSHDialer{sshConfig: config.SSH, dialContext: ssh.DialContextThroughSSH}
+	}
+	if config.UseProxy {
+		return &mongoProxyDialer{proxyConfig: config.Proxy}
+	}
+	return nil
+}
+
 const defaultMongoPort = 27017
 const mongoObjectIDLocatorColumn = "__gonavi_mongodb_id_locator__"

@@ -336,48 +354,9 @@ func (m *MongoDBV1) Connect(config connection.ConnectionConfig) (err error) {
 	runConfig := applyMongoURI(config)
 	connectConfig := runConfig
 	sshRouteHint := ""
-
-	if runConfig.UseSSH && runConfig.MongoSRV {
-		return fmt.Errorf("MongoDB SRV 记录模式暂不支持 SSH 隧道")
-	}
-
 	if runConfig.UseSSH {
-		seeds := collectMongoSeeds(runConfig)
-		if len(seeds) == 0 {
-			seeds = append(seeds, normalizeMongoAddress(runConfig.Host, runConfig.Port))
-		}
-		targetHost, targetPort, ok := parseHostPortWithDefault(seeds[0], defaultMongoPort)
-		if !ok {
-			return fmt.Errorf("MongoDB 连接失败：无效地址 %s", seeds[0])
-		}
-
-		logger.Infof("MongoDB 使用 SSH 连接：地址=%s:%d", targetHost, targetPort)
-
-		forwarder, err := ssh.AcquireLocalForwarder(runConfig.SSH, targetHost, targetPort)
-		if err != nil {
-			return fmt.Errorf("创建 SSH 隧道失败：%w", err)
-		}
-		m.forwarder = forwarder
-
-		host, portStr, err := net.SplitHostPort(forwarder.LocalAddr)
-		if err != nil {
-			return fmt.Errorf("解析本地转发地址失败：%w", err)
-		}
-
-		port, err := strconv.Atoi(portStr)
-		if err != nil {
-			return fmt.Errorf("解析本地端口失败：%w", err)
-		}
-
-		localConfig := runConfig
-		localConfig.Host = host
-		localConfig.Port = port
-		localConfig.UseSSH = false
-		localConfig.URI = ""
-		localConfig.Hosts = []string{normalizeMongoAddress(host, port)}
-		connectConfig = localConfig
-		sshRouteHint = fmt.Sprintf("SSH隧道 %s -> %s:%d", forwarder.LocalAddr, targetHost, targetPort)
-		logger.Infof("MongoDB 通过本地端口转发连接：%s -> %s:%d", forwarder.LocalAddr, targetHost, targetPort)
+		sshRouteHint = fmt.Sprintf("SSH隧道 %s:%d", runConfig.SSH.Host, runConfig.SSH.Port)
+		logger.Infof("MongoDB(v1) 使用 SSH 隧道连接：跳板=%s:%d", runConfig.SSH.Host, runConfig.SSH.Port)
 	}

 	m.pingTimeout = getConnectTimeout(connectConfig)
@@ -432,8 +411,8 @@ func (m *MongoDBV1) Connect(config connection.ConnectionConfig) (err error) {
 			if tlsConfig != nil {
 				clientOpts.SetTLSConfig(tlsConfig)
 			}
-			if attemptConfig.UseProxy {
-				clientOpts.SetDialer(&mongoProxyDialer{proxyConfig: attemptConfig.Proxy})
+			if dialer := mongoConnectionDialer(attemptConfig); dialer != nil {
+				clientOpts.SetDialer(dialer)
 			}
 			connectCtx, connectCancel := context.WithTimeout(context.Background(), m.pingTimeout)
 			client, err := mongo.Connect(connectCtx, clientOpts)
@@ -481,13 +460,6 @@ func (m *MongoDBV1) Connect(config connection.ConnectionConfig) (err error) {
 }

 func (m *MongoDBV1) Close() error {
-	if m.forwarder != nil {
-		if err := m.forwarder.Release(); err != nil {
-			logger.Warnf("关闭 MongoDB SSH 端口转发失败：%v", err)
-		}
-		m.forwarder = nil
-	}
-
 	if m.client != nil {
 		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
 		defer cancel()
--- a/internal/db/mongodb_impl_v1_uri_test.go
+++ b/internal/db/mongodb_impl_v1_uri_test.go
@@ -3,6 +3,10 @@
 package db

 import (
+	"context"
+	"errors"
+	"net"
+	"strings"
 	"testing"

 	"GoNavi-Wails/internal/connection"
@@ -11,6 +15,59 @@ import (
 	"go.mongodb.org/mongo-driver/bson/primitive"
 )

+func TestMongoSSHDialerV1RoutesAllMembersThroughSSH(t *testing.T) {
+	tests := []struct {
+		name       string
+		uri        string
+		wantScheme string
+	}{
+		{name: "standard", uri: "mongodb://mongo.internal:27017/app", wantScheme: "mongodb://"},
+		{name: "srv", uri: "mongodb+srv://cluster.example.test/app", wantScheme: "mongodb+srv://"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := applyMongoURI(connection.ConnectionConfig{
+				URI:    tt.uri,
+				UseSSH: true,
+				SSH: connection.SSHConfig{
+					Host: "bastion.example.test",
+					Port: 22,
+					User: "operator",
+				},
+			})
+			if uri := (&MongoDBV1{}).getURI(config); !strings.HasPrefix(uri, tt.wantScheme) {
+				t.Fatalf("expected URI scheme %q, got %q", tt.wantScheme, uri)
+			}
+
+			dialer := mongoConnectionDialer(config)
+			sshDialer, ok := dialer.(*mongoSSHDialer)
+			if !ok {
+				t.Fatalf("expected SSH dialer, got %T", dialer)
+			}
+
+			wantErr := errors.New("dial stopped")
+			var addresses []string
+			sshDialer.dialContext = func(_ context.Context, _ connection.SSHConfig, network, address string) (net.Conn, error) {
+				if network != "tcp" {
+					t.Fatalf("expected tcp network, got %q", network)
+				}
+				addresses = append(addresses, address)
+				return nil, wantErr
+			}
+
+			for _, address := range []string{"mongo-1.internal:27017", "mongo-2.internal:27018"} {
+				if _, err := dialer.DialContext(context.Background(), "tcp", address); !errors.Is(err, wantErr) {
+					t.Fatalf("expected SSH dial error for %s, got %v", address, err)
+				}
+			}
+			if strings.Join(addresses, ",") != "mongo-1.internal:27017,mongo-2.internal:27018" {
+				t.Fatalf("unexpected SSH targets: %v", addresses)
+			}
+		})
+	}
+}
+
 func TestApplyMongoURIV1_ExplicitHostDoesNotAdoptURIHosts(t *testing.T) {
 	config := connection.ConnectionConfig{
 		Host: "10.10.10.10",