opsxdb/apps/server/internal/db/ssh_tunnel.go

package db

import (
	"fmt"
	"io"
	"net"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	"golang.org/x/crypto/ssh"
)

// SSHTunnel holds a local TCP listener that forwards through an SSH session
// to the remote database host:port (local port forward semantics).
type SSHTunnel struct {
	LocalPort int
	listener  net.Listener
	client    *ssh.Client
	cancel    chan struct{}
	wg        sync.WaitGroup
	once      sync.Once
}

func (t *SSHTunnel) Close() {
	if t == nil {
		return
	}
	t.once.Do(func() {
		close(t.cancel)
		if t.listener != nil {
			_ = t.listener.Close()
		}
		if t.client != nil {
			_ = t.client.Close()
		}
		t.wg.Wait()
	})
}

// StartLocalForward dials SSH and listens on 127.0.0.1:0, forwarding to c.Host:c.Port.
func StartLocalForward(c ConnReq) (*SSHTunnel, error) {
	sshHost := strings.TrimSpace(c.SSHHost)
	if sshHost == "" {
		return nil, fmt.Errorf("请填写 SSH 主机")
	}
	sshPort := c.SSHPort
	if sshPort <= 0 {
		sshPort = 22
	}
	sshUser := strings.TrimSpace(c.SSHUser)
	if sshUser == "" {
		return nil, fmt.Errorf("请填写 SSH 用户")
	}

	auth, err := sshAuthMethods(c)
	if err != nil {
		return nil, err
	}

	cfg := &ssh.ClientConfig{
		User:            sshUser,
		Auth:            auth,
		HostKeyCallback: ssh.InsecureIgnoreHostKey(), //nolint:gosec // product tunnel; host-key pinning later
		Timeout:         12 * time.Second,
	}

	addr := net.JoinHostPort(sshHost, fmt.Sprintf("%d", sshPort))
	client, err := ssh.Dial("tcp", addr, cfg)
	if err != nil {
		return nil, fmt.Errorf("无法连接跳板机 %s: %w", addr, err)
	}

	ln, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		_ = client.Close()
		return nil, fmt.Errorf("本地监听失败: %w", err)
	}
	localPort := ln.Addr().(*net.TCPAddr).Port

	remoteHost := c.Host
	if remoteHost == "" {
		remoteHost = "127.0.0.1"
	}
	remotePort := c.Port
	if remotePort <= 0 {
		remotePort = 3306
	}
	remoteAddr := net.JoinHostPort(remoteHost, fmt.Sprintf("%d", remotePort))

	tun := &SSHTunnel{
		LocalPort: localPort,
		listener:  ln,
		client:    client,
		cancel:    make(chan struct{}),
	}

	tun.wg.Add(1)
	go func() {
		defer tun.wg.Done()
		for {
			conn, err := ln.Accept()
			if err != nil {
				select {
				case <-tun.cancel:
					return
				default:
					return
				}
			}
			tun.wg.Add(1)
			go func(local net.Conn) {
				defer tun.wg.Done()
				defer local.Close()
				remote, err := client.Dial("tcp", remoteAddr)
				if err != nil {
					return
				}
				defer remote.Close()
				pipe(local, remote)
			}(conn)
		}
	}()

	return tun, nil
}

func sshAuthMethods(c ConnReq) ([]ssh.AuthMethod, error) {
	authMode := strings.ToLower(strings.TrimSpace(c.SSHAuth))
	if authMode == "" {
		authMode = "password"
	}
	switch authMode {
	case "password":
		if c.SSHPassword == "" {
			return nil, fmt.Errorf("请填写 SSH 密码")
		}
		return []ssh.AuthMethod{ssh.Password(c.SSHPassword)}, nil
	case "privatekey", "key", "private_key":
		keyMaterial, err := loadSSHPrivateKey(c.SSHPrivateKey)
		if err != nil {
			return nil, err
		}
		var signer ssh.Signer
		if strings.TrimSpace(c.SSHPassphrase) != "" {
			signer, err = ssh.ParsePrivateKeyWithPassphrase(keyMaterial, []byte(c.SSHPassphrase))
		} else {
			signer, err = ssh.ParsePrivateKey(keyMaterial)
		}
		if err != nil {
			return nil, fmt.Errorf("解析私钥失败: %w", err)
		}
		return []ssh.AuthMethod{ssh.PublicKeys(signer)}, nil
	default:
		return nil, fmt.Errorf("不支持的 SSH 认证方式: %s", c.SSHAuth)
	}
}

func loadSSHPrivateKey(raw string) ([]byte, error) {
	s := strings.TrimSpace(raw)
	if s == "" {
		return nil, fmt.Errorf("请填写私钥 PEM 或路径")
	}
	if strings.Contains(s, "BEGIN") && strings.Contains(s, "PRIVATE KEY") {
		return []byte(s), nil
	}
	path := s
	if strings.HasPrefix(path, "~/") {
		home, err := os.UserHomeDir()
		if err != nil {
			return nil, fmt.Errorf("展开 ~ 失败: %w", err)
		}
		path = filepath.Join(home, path[2:])
	}
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("读取私钥文件失败: %w", err)
	}
	return b, nil
}

func pipe(a, b net.Conn) {
	var wg sync.WaitGroup
	wg.Add(2)
	go func() {
		defer wg.Done()
		_, _ = io.Copy(a, b)
		_ = a.Close()
	}()
	go func() {
		defer wg.Done()
		_, _ = io.Copy(b, a)
		_ = b.Close()
	}()
	wg.Wait()
}

// ApplySSHTunnel starts a local forward when sshEnabled and rewrites Host/Port to 127.0.0.1:<local>.
// Caller must Close the returned tunnel (may be nil).
func ApplySSHTunnel(c ConnReq) (ConnReq, *SSHTunnel, error) {
	if !c.SSHEnabled {
		return c, nil, nil
	}
	tun, err := StartLocalForward(c)
	if err != nil {
		return c, nil, fmt.Errorf("SSH: %w", err)
	}
	out := c
	out.Host = "127.0.0.1"
	out.Port = tun.LocalPort
	// Avoid nested tunnel if helpers re-enter
	out.SSHEnabled = false
	return out, tun, nil
}