JRedisDesktop/electron/redis/connection.ts

// electron/main/redis/connection.ts
// Redis 连接管理（支持 SSH 隧道）
import Redis, { RedisOptions } from 'ioredis'
import { Client as SSHClient } from 'ssh2'
import { createServer, type AddressInfo } from 'net'
import { readFileSync } from 'fs'
import { resolve } from 'path'
import os from 'os'

export interface ConnectionOptions {
  host: string
  port: number
  username?: string
  password?: string
  tls?: boolean
  tlsCaPath?: string
  tlsCertPath?: string
  tlsKeyPath?: string
  tlsRejectUnauthorized?: boolean
  sentinelEnabled?: boolean
  sentinelHost?: string
  sentinelPort?: number
  sentinelMasterName?: string
  sentinelNodePassword?: string
  cluster?: boolean
  sshEnabled?: boolean
  sshHost?: string
  sshPort?: number
  sshUsername?: string
  sshPassword?: string
  sshPrivateKeyPath?: string
  sshPassphrase?: string
}

export type RedisClient = Redis

interface ClientEntry {
  client: RedisClient
  sshClient?: SSHClient
  tunnelServer?: ReturnType<typeof createServer>
}

const clients = new Map<string, ClientEntry>()

function expandHomePath(p: string): string {
  if (p.startsWith('~')) {
    return resolve(os.homedir(), p.slice(p.startsWith('~/') ? 2 : 1))
  }
  return resolve(p)
}

async function createSshTunnel(opts: ConnectionOptions): Promise<{ localPort: number; sshClient: SSHClient; server: ReturnType<typeof createServer> }> {
  const sshConfig: any = {
    host: opts.sshHost || opts.host,
    port: opts.sshPort || 22,
    username: opts.sshUsername || 'root',
    readyTimeout: 10000,
    keepaliveInterval: 10000,
  }

  if (opts.sshPrivateKeyPath) {
    try {
      sshConfig.privateKey = readFileSync(expandHomePath(opts.sshPrivateKeyPath))
      if (opts.sshPassphrase) {
        sshConfig.passphrase = opts.sshPassphrase
      }
    } catch {
      // private key file not found, fall through to password auth
    }
  }

  if (opts.sshPassword) {
    sshConfig.password = opts.sshPassword
  }

  return new Promise((resolve, reject) => {
    const sshClient = new SSHClient()

    sshClient.on('ready', () => {
      // Create a local TCP server to forward connections
      const server = createServer((socket) => {
        sshClient.forwardOut(
          '127.0.0.1',
          0,
          opts.host,
          opts.port,
          (err, stream) => {
            if (err) {
              socket.destroy()
              return
            }
            socket.pipe(stream).pipe(socket)
          }
        )
      })

      server.listen(0, '127.0.0.1', () => {
        const localPort = (server.address() as AddressInfo).port
        resolve({ localPort, sshClient, server })
      })

      server.on('error', (err) => {
        sshClient.end()
        reject(new Error(`SSH tunnel server error: ${err.message}`))
      })
    })

    sshClient.on('error', (err: Error) => {
      reject(new Error(`SSH connection failed: ${err.message}`))
    })

    sshClient.connect(sshConfig)
  })
}

function buildRedisOptions(opts: ConnectionOptions): RedisOptions {
  const options: RedisOptions = {
    host: opts.host,
    port: opts.port,
    username: opts.username || undefined,
    password: opts.password || undefined,
    stringNumbers: true,
    enableReadyCheck: true,
    family: 0,
    connectTimeout: 10000,
    retryStrategy(times: number) {
      if (times > 3) return null
      return Math.min(times * 200, 1000)
    },
    maxRetriesPerRequest: 3,
  }

  if (opts.tls) {
    const tlsOpts: any = {
      rejectUnauthorized: opts.tlsRejectUnauthorized || false,
      checkServerIdentity: () => undefined,
    }
    if (opts.tlsCaPath) {
      try { tlsOpts.ca = readFileSync(expandHomePath(opts.tlsCaPath)) } catch { /* ignore */ }
    }
    if (opts.tlsCertPath) {
      try { tlsOpts.cert = readFileSync(expandHomePath(opts.tlsCertPath)) } catch { /* ignore */ }
    }
    if (opts.tlsKeyPath) {
      try { tlsOpts.key = readFileSync(expandHomePath(opts.tlsKeyPath)) } catch { /* ignore */ }
    }
    options.tls = tlsOpts
  }

  return options
}

export function getClient(id: string): RedisClient | undefined {
  return clients.get(id)?.client
}

export function isClientConnected(id: string): boolean {
  const entry = clients.get(id)
  return entry?.client?.status === 'ready'
}

export async function connect(id: string, opts: ConnectionOptions): Promise<void> {
  if (clients.has(id)) {
    await disconnect(id)
  }

  return new Promise(async (resolve, reject) => {
    let sshClient: SSHClient | undefined
    let tunnelServer: ReturnType<typeof createServer> | undefined

    try {
      if (opts.sentinelEnabled && opts.sentinelMasterName) {
        // Sentinel mode: connect via sentinel
        const sentinelOptions: RedisOptions = {
          sentinels: [{ host: opts.sentinelHost || opts.host, port: opts.sentinelPort || 26379 }],
          name: opts.sentinelMasterName,
          sentinelPassword: opts.sentinelNodePassword || undefined,
          username: opts.username || undefined,
          password: opts.password || undefined,
          stringNumbers: true,
          enableReadyCheck: true,
          connectTimeout: 10000,
          retryStrategy(times: number) {
            if (times > 3) return null
            return Math.min(times * 200, 1000)
          },
          maxRetriesPerRequest: 3,
        }
        if (opts.tls) {
          sentinelOptions.tls = {
            rejectUnauthorized: opts.tlsRejectUnauthorized || false,
            checkServerIdentity: () => undefined,
          }
        }
        const redis = new Redis(sentinelOptions)

        redis.once('ready', () => {
          clients.set(id, { client: redis })
          resolve()
        })
        redis.once('error', (err: Error) => {
          redis.disconnect()
          reject(err)
        })
        setTimeout(() => {
          if (!clients.has(id)) {
            redis.disconnect()
            reject(new Error('连接超时，请检查主机和端口'))
          }
        }, 12000)
        return
      }

      if (opts.cluster) {
        // Cluster mode
        const clusterOptions: any = {
          redisOptions: {
            username: opts.username || undefined,
            password: opts.password || undefined,
            tls: opts.tls ? {
              rejectUnauthorized: opts.tlsRejectUnauthorized || false,
              checkServerIdentity: () => undefined,
            } : undefined,
          },
          clusterRetryStrategy(times: number) {
            if (times > 3) return null
            return Math.min(times * 200, 1000)
          },
          enableReadyCheck: true,
        }
        const redis = new Redis.Cluster([{ host: opts.host, port: opts.port }], clusterOptions)

        redis.once('ready', () => {
          clients.set(id, { client: redis as any })
          resolve()
        })
        redis.once('error', (err: Error) => {
          redis.disconnect()
          reject(err)
        })
        setTimeout(() => {
          if (!clients.has(id)) {
            redis.disconnect()
            reject(new Error('连接超时，请检查主机和端口'))
          }
        }, 12000)
        return
      }

      if (opts.sshEnabled) {
        const tunnel = await createSshTunnel(opts)
        sshClient = tunnel.sshClient
        tunnelServer = tunnel.server
        // Redirect Redis connection through SSH tunnel
        opts = { ...opts, host: '127.0.0.1', port: tunnel.localPort }
      }

      const redisOptions = buildRedisOptions(opts)
      const redis = new Redis(redisOptions)

      redis.once('ready', () => {
        clients.set(id, { client: redis, sshClient, tunnelServer })
        resolve()
      })

      redis.once('error', (err: Error) => {
        redis.disconnect()
        if (sshClient) sshClient.end()
        if (tunnelServer) tunnelServer.close()
        reject(err)
      })

      setTimeout(() => {
        if (!clients.has(id)) {
          redis.disconnect()
          if (sshClient) sshClient.end()
          if (tunnelServer) tunnelServer.close()
          reject(new Error('连接超时，请检查主机和端口'))
        }
      }, 12000)
    } catch (err: any) {
      reject(err)
    }
  })
}

export async function disconnect(id: string): Promise<void> {
  const entry = clients.get(id)
  if (entry) {
    entry.client.disconnect()
    if (entry.sshClient) entry.sshClient.end()
    if (entry.tunnelServer) entry.tunnelServer.close()
    clients.delete(id)
  }
}

export function disconnectAll(): void {
  for (const [id] of clients) {
    disconnect(id)
  }
}