floci-ui/packages/api/src/services/sshCheck.ts

import {createConnection} from 'node:net'

const DEFAULT_TIMEOUT_MS = 3_000
const MAX_BANNER_BYTES = 8_192

type Ec2DiagnosticInstance = {
    instanceId?: string
    sshHost?: string
    sshPort?: number
    containerStatus?: 'running' | 'stopped' | 'unavailable'
}

type DiagnosticResponse = {
    ec2?: {instances?: Ec2DiagnosticInstance[]}
}

export type SshTarget = {
    host: string
    port: number
}

export type SshCheckResult = {
    ok: true
    banner: string
    checkedFrom: 'UI backend'
    checkedAt: string
    target: SshTarget
}

export type SshCheckErrorCode =
    | 'DIAGNOSTICS_UNAVAILABLE'
    | 'INSTANCE_NOT_RUNNING'
    | 'CONNECTION_REFUSED'
    | 'CONNECTION_TIMEOUT'
    | 'SSH_BANNER_MISSING'

export class SshCheckError extends Error {
    constructor(
        readonly code: SshCheckErrorCode,
        message: string,
        readonly status: 404 | 409 | 502 | 503 | 504,
    ) {
        super(message)
        this.name = 'SshCheckError'
    }
}

export type DiagnosticFetcher = () => Promise<DiagnosticResponse>
export type SshSocketConnector = (target: SshTarget, timeoutMs: number) => Promise<string>

type SshCheckDependencies = {
    fetchDiagnostics?: DiagnosticFetcher
    connect?: SshSocketConnector
    loopbackHost?: string
    now?: () => Date
    timeoutMs?: number
}

async function fetchDiagnostics(): Promise<DiagnosticResponse> {
    const endpoint = process.env.FLOCI_ENDPOINT
    if (!endpoint) {
        throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', 'TheCloud diagnostics are unavailable', 503)
    }

    let response: Response
    try {
        response = await fetch(`${endpoint}/_floci/diagnose`, {signal: AbortSignal.timeout(DEFAULT_TIMEOUT_MS)})
    } catch {
        throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', 'TheCloud diagnostics are unavailable', 503)
    }
    if (!response.ok) {
        throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', 'TheCloud diagnostics are unavailable', 503)
    }

    try {
        return await response.json() as DiagnosticResponse
    } catch {
        throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', 'TheCloud diagnostics returned an invalid response', 502)
    }
}

export function readSshBanner(target: SshTarget, timeoutMs: number): Promise<string> {
    return new Promise((resolve, reject) => {
        const socket = createConnection(target)
        let settled = false
        let buffer = Buffer.alloc(0)

        const finish = (error?: SshCheckError, banner?: string) => {
            if (settled) return
            settled = true
            clearTimeout(timer)
            socket.destroy()
            if (error) reject(error)
            else resolve(banner!)
        }

        const timer = setTimeout(() => {
            finish(new SshCheckError('CONNECTION_TIMEOUT', `SSH connection to ${target.host}:${target.port} timed out`, 504))
        }, timeoutMs)

        socket.on('data', (chunk: Buffer) => {
            buffer = Buffer.concat([buffer, chunk])
            if (buffer.length > MAX_BANNER_BYTES) {
                finish(new SshCheckError('SSH_BANNER_MISSING', 'SSH server did not send a valid banner', 502))
                return
            }

            const lines = buffer.toString('utf8').split(/\r?\n/)
            for (let index = 0; index < lines.length - 1; index++) {
                const line = lines[index]
                if (/^SSH-\d+\.\d+-[\x21-\x7e]+(?: [\x20-\x7e]*)?$/.test(line)
                    && Buffer.byteLength(line, 'ascii') <= 253) {
                    finish(undefined, line)
                    return
                }
            }
        })
        socket.on('end', () => {
            finish(new SshCheckError('SSH_BANNER_MISSING', 'SSH server closed without sending a valid banner', 502))
        })
        socket.on('close', () => {
            finish(new SshCheckError('SSH_BANNER_MISSING', 'SSH server closed without sending a valid banner', 502))
        })
        socket.on('error', (error: NodeJS.ErrnoException) => {
            if (error.code === 'ECONNREFUSED') {
                finish(new SshCheckError('CONNECTION_REFUSED', `SSH connection to ${target.host}:${target.port} was refused`, 503))
                return
            }
            if (error.code === 'ETIMEDOUT') {
                finish(new SshCheckError('CONNECTION_TIMEOUT', `SSH connection to ${target.host}:${target.port} timed out`, 504))
                return
            }
            finish(new SshCheckError('CONNECTION_REFUSED', `SSH connection to ${target.host}:${target.port} failed`, 503))
        })
    })
}

export function createSshCheckService(dependencies: SshCheckDependencies = {}) {
    const diagnosticFetcher = dependencies.fetchDiagnostics ?? fetchDiagnostics
    const connector = dependencies.connect ?? readSshBanner
    const sshHost = dependencies.loopbackHost || process.env.FLOCI_EC2_SSH_HOST || 'localhost'
    const now = dependencies.now ?? (() => new Date())
    const timeoutMs = dependencies.timeoutMs ?? DEFAULT_TIMEOUT_MS

    return {
        async check(instanceId: string): Promise<SshCheckResult> {
            let diagnostics: DiagnosticResponse
            try {
                diagnostics = await diagnosticFetcher()
            } catch (error) {
                if (error instanceof SshCheckError) throw error
                throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', 'TheCloud diagnostics are unavailable', 503)
            }

            if (!diagnostics || typeof diagnostics !== 'object'
                || !Array.isArray(diagnostics.ec2?.instances)) {
                throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', 'TheCloud diagnostics returned an invalid response', 502)
            }
            const instance = diagnostics.ec2.instances.find((candidate) => candidate.instanceId === instanceId)
            if (!instance) {
                throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', `SSH diagnostics for instance ${instanceId} are unavailable`, 404)
            }
            if (!instance.containerStatus) {
                throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', `SSH diagnostics for instance ${instanceId} are unavailable`, 404)
            }
            if (instance.containerStatus !== 'running') {
                throw new SshCheckError('INSTANCE_NOT_RUNNING', `Instance ${instanceId} is not running`, 409)
            }
            if (typeof instance.sshPort !== 'number' || !Number.isInteger(instance.sshPort)
                || instance.sshPort < 1 || instance.sshPort > 65_535) {
                throw new SshCheckError('DIAGNOSTICS_UNAVAILABLE', `SSH diagnostics for instance ${instanceId} are unavailable`, 404)
            }

            const target = {
                host: sshHost,
                port: instance.sshPort,
            }
            const banner = await connector(target, timeoutMs)
            if (!/^SSH-\d+\.\d+-/.test(banner)) {
                throw new SshCheckError('SSH_BANNER_MISSING', 'SSH server did not send a valid banner', 502)
            }

            return {
                ok: true,
                banner,
                checkedFrom: 'UI backend',
                checkedAt: now().toISOString(),
                target,
            }
        },
    }
}

export const sshCheckService = createSshCheckService()