mirror of
https://github.com/Syngnat/GoNavi.git
synced 2026-06-03 04:59:46 +08:00
♻️ refactor(database/ssh): SSH隧道架构重构与多数据源适配
- 架构升级:从driver专属拨号器改为通用本地端口转发模式
- 并发安全:sync.Once保护Close操作,RWMutex保护状态访问,双向errc等待
- 连接池化:GetOrCreateLocalForwarder/GetOrCreateSSHClient实现缓存复用
- SQL安全:kingbase_impl.go引入esc函数,防止双引号注入(""ldf_server""问题)
- Schema动态化:三级fallback(schema.table解析→dbName参数→current_schema())
- 代码复用:scanRows统一行扫描逻辑,normalizeQueryValueWithDBType增强类型处理
Close #40
This commit is contained in:
@@ -3,8 +3,10 @@ package ssh
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"os"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"GoNavi-Wails/internal/connection"
|
||||
@@ -110,3 +112,264 @@ func RegisterSSHNetwork(sshConfig connection.SSHConfig) (string, error) {
|
||||
|
||||
return netName, nil
|
||||
}
|
||||
|
||||
// sshClientCache stores SSH clients to avoid creating multiple connections
|
||||
var (
|
||||
sshClientCache = make(map[string]*ssh.Client)
|
||||
sshClientCacheMu sync.RWMutex
|
||||
localForwarders = make(map[string]*LocalForwarder)
|
||||
forwarderMu sync.RWMutex
|
||||
)
|
||||
|
||||
// LocalForwarder represents a local port forwarder through SSH
|
||||
type LocalForwarder struct {
|
||||
LocalAddr string
|
||||
RemoteAddr string
|
||||
SSHClient *ssh.Client
|
||||
listener net.Listener
|
||||
closeChan chan struct{}
|
||||
closeOnce sync.Once // 防止重复关闭
|
||||
closed bool // 关闭状态标记
|
||||
closedMu sync.RWMutex
|
||||
}
|
||||
|
||||
// NewLocalForwarder creates a new local port forwarder
|
||||
// It listens on a random local port and forwards all connections through SSH tunnel
|
||||
func NewLocalForwarder(sshConfig connection.SSHConfig, remoteHost string, remotePort int) (*LocalForwarder, error) {
|
||||
client, err := GetOrCreateSSHClient(sshConfig)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("建立 SSH 连接失败:%w", err)
|
||||
}
|
||||
|
||||
// Listen on localhost with a random port
|
||||
listener, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建本地监听器失败:%w", err)
|
||||
}
|
||||
|
||||
localAddr := listener.Addr().String()
|
||||
remoteAddr := fmt.Sprintf("%s:%d", remoteHost, remotePort)
|
||||
|
||||
forwarder := &LocalForwarder{
|
||||
LocalAddr: localAddr,
|
||||
RemoteAddr: remoteAddr,
|
||||
SSHClient: client,
|
||||
listener: listener,
|
||||
closeChan: make(chan struct{}),
|
||||
}
|
||||
|
||||
// Start forwarding in background
|
||||
go forwarder.forward()
|
||||
|
||||
logger.Infof("已创建 SSH 端口转发:本地 %s -> 远程 %s", localAddr, remoteAddr)
|
||||
return forwarder, nil
|
||||
}
|
||||
|
||||
// forward handles the port forwarding
|
||||
func (f *LocalForwarder) forward() {
|
||||
for {
|
||||
localConn, err := f.listener.Accept()
|
||||
if err != nil {
|
||||
// Check if we're shutting down
|
||||
select {
|
||||
case <-f.closeChan:
|
||||
return
|
||||
default:
|
||||
logger.Warnf("接受本地连接失败:%v", err)
|
||||
// listener可能已关闭,退出循环
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
go f.handleConnection(localConn)
|
||||
}
|
||||
}
|
||||
|
||||
// handleConnection handles a single connection
|
||||
func (f *LocalForwarder) handleConnection(localConn net.Conn) {
|
||||
defer localConn.Close()
|
||||
|
||||
// Connect to remote through SSH with timeout
|
||||
remoteConn, err := f.SSHClient.Dial("tcp", f.RemoteAddr)
|
||||
if err != nil {
|
||||
logger.Warnf("通过 SSH 连接到远程 %s 失败:%v", f.RemoteAddr, err)
|
||||
return
|
||||
}
|
||||
defer remoteConn.Close()
|
||||
|
||||
// Bidirectional copy with error channel
|
||||
errc := make(chan error, 2)
|
||||
|
||||
// Copy from local to remote
|
||||
go func() {
|
||||
_, err := io.Copy(remoteConn, localConn)
|
||||
if err != nil {
|
||||
logger.Warnf("本地->远程数据复制错误:%v", err)
|
||||
}
|
||||
errc <- err
|
||||
}()
|
||||
|
||||
// Copy from remote to local
|
||||
go func() {
|
||||
_, err := io.Copy(localConn, remoteConn)
|
||||
if err != nil {
|
||||
logger.Warnf("远程->本地数据复制错误:%v", err)
|
||||
}
|
||||
errc <- err
|
||||
}()
|
||||
|
||||
// Wait for BOTH goroutines to complete
|
||||
<-errc
|
||||
<-errc
|
||||
}
|
||||
|
||||
// Close closes the forwarder (thread-safe, can be called multiple times)
|
||||
func (f *LocalForwarder) Close() error {
|
||||
var err error
|
||||
f.closeOnce.Do(func() {
|
||||
f.closedMu.Lock()
|
||||
f.closed = true
|
||||
f.closedMu.Unlock()
|
||||
|
||||
close(f.closeChan)
|
||||
err = f.listener.Close()
|
||||
if err != nil {
|
||||
logger.Warnf("关闭端口转发监听器失败:%v", err)
|
||||
}
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// IsClosed returns whether the forwarder is closed
|
||||
func (f *LocalForwarder) IsClosed() bool {
|
||||
f.closedMu.RLock()
|
||||
defer f.closedMu.RUnlock()
|
||||
return f.closed
|
||||
}
|
||||
|
||||
// GetOrCreateLocalForwarder returns a cached forwarder or creates a new one
|
||||
func GetOrCreateLocalForwarder(sshConfig connection.SSHConfig, remoteHost string, remotePort int) (*LocalForwarder, error) {
|
||||
key := fmt.Sprintf("%s:%d:%s->%s:%d",
|
||||
sshConfig.Host, sshConfig.Port, sshConfig.User,
|
||||
remoteHost, remotePort)
|
||||
|
||||
forwarderMu.RLock()
|
||||
forwarder, exists := localForwarders[key]
|
||||
forwarderMu.RUnlock()
|
||||
|
||||
// Check if exists and is still valid
|
||||
if exists && forwarder != nil && !forwarder.IsClosed() {
|
||||
logger.Infof("复用已有端口转发:%s", key)
|
||||
return forwarder, nil
|
||||
}
|
||||
|
||||
// Remove stale forwarder from cache
|
||||
if exists {
|
||||
forwarderMu.Lock()
|
||||
delete(localForwarders, key)
|
||||
forwarderMu.Unlock()
|
||||
}
|
||||
|
||||
forwarder, err := NewLocalForwarder(sshConfig, remoteHost, remotePort)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
forwarderMu.Lock()
|
||||
localForwarders[key] = forwarder
|
||||
forwarderMu.Unlock()
|
||||
|
||||
return forwarder, nil
|
||||
}
|
||||
|
||||
// CloseAllForwarders closes all local forwarders
|
||||
func CloseAllForwarders() {
|
||||
forwarderMu.Lock()
|
||||
defer forwarderMu.Unlock()
|
||||
|
||||
for key, forwarder := range localForwarders {
|
||||
if forwarder != nil {
|
||||
_ = forwarder.Close()
|
||||
logger.Infof("已关闭端口转发:%s", key)
|
||||
}
|
||||
}
|
||||
localForwarders = make(map[string]*LocalForwarder)
|
||||
}
|
||||
|
||||
|
||||
// getSSHClientCacheKey generates a unique cache key for SSH config
|
||||
func getSSHClientCacheKey(config connection.SSHConfig) string {
|
||||
return fmt.Sprintf("%s:%d:%s", config.Host, config.Port, config.User)
|
||||
}
|
||||
|
||||
// GetOrCreateSSHClient returns a cached SSH client or creates a new one
|
||||
func GetOrCreateSSHClient(config connection.SSHConfig) (*ssh.Client, error) {
|
||||
key := getSSHClientCacheKey(config)
|
||||
|
||||
sshClientCacheMu.RLock()
|
||||
client, exists := sshClientCache[key]
|
||||
sshClientCacheMu.RUnlock()
|
||||
|
||||
if exists && client != nil {
|
||||
// Test if connection is still alive by creating a test session
|
||||
session, err := client.NewSession()
|
||||
if err == nil {
|
||||
session.Close()
|
||||
logger.Infof("复用已有 SSH 连接:%s", key)
|
||||
return client, nil
|
||||
}
|
||||
// Connection is dead, remove from cache
|
||||
logger.Warnf("SSH 连接已断开,重新建立:%s (错误: %v)", key, err)
|
||||
sshClientCacheMu.Lock()
|
||||
delete(sshClientCache, key)
|
||||
sshClientCacheMu.Unlock()
|
||||
// Try to close the dead client
|
||||
_ = client.Close()
|
||||
}
|
||||
|
||||
// Create new SSH client
|
||||
client, err := connectSSH(config)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Cache the client
|
||||
sshClientCacheMu.Lock()
|
||||
sshClientCache[key] = client
|
||||
sshClientCacheMu.Unlock()
|
||||
|
||||
logger.Infof("已缓存 SSH 连接:%s", key)
|
||||
return client, nil
|
||||
}
|
||||
|
||||
// DialThroughSSH creates a connection through SSH tunnel
|
||||
// This is a generic dialer that can be used by any database driver
|
||||
func DialThroughSSH(config connection.SSHConfig, network, address string) (net.Conn, error) {
|
||||
client, err := GetOrCreateSSHClient(config)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("建立 SSH 连接失败:%w", err)
|
||||
}
|
||||
|
||||
conn, err := client.Dial(network, address)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("通过 SSH 隧道连接到 %s 失败:%w", address, err)
|
||||
}
|
||||
|
||||
logger.Infof("已通过 SSH 隧道连接到:%s", address)
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
// CloseAllSSHClients closes all cached SSH clients
|
||||
func CloseAllSSHClients() {
|
||||
sshClientCacheMu.Lock()
|
||||
defer sshClientCacheMu.Unlock()
|
||||
|
||||
for key, client := range sshClientCache {
|
||||
if client != nil {
|
||||
_ = client.Close()
|
||||
logger.Infof("已关闭 SSH 连接:%s", key)
|
||||
}
|
||||
}
|
||||
sshClientCache = make(map[string]*ssh.Client)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user