DodoSSH/src/DodoSSH.Client.Ssh/SshNetConnectionFactory.cs

using System.Text;
using Renci.SshNet;
using Renci.SshNet.Common;

namespace DodoSSH.Client.Ssh;

/// <summary>
/// Opens SSH connections with SSH.NET, checking host key trust during the handshake.
/// </summary>
/// <remarks>
/// The pinned fingerprint is looked up <em>before</em> connecting, so the comparison inside SSH.NET's
/// synchronous <c>HostKeyReceived</c> event is a pure equality check with no I/O and no chance of
/// blocking the handshake on a UI round trip. Anything the comparison cannot settle becomes an
/// exception the caller resolves asynchronously.
/// </remarks>
public sealed class SshNetConnectionFactory(IKnownHostStore knownHosts) : ISshConnectionFactory
{
    private static readonly TimeSpan DefaultConnectTimeout = TimeSpan.FromSeconds(15);

    /// <inheritdoc />
    public async Task<ISshConnection> ConnectAsync(
        SshConnectionRequest request,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);

        var client = new SshClient(BuildConnectionInfo(request));
        var gate = new HostKeyGate(knownHosts, request, cancellationToken);

        client.HostKeyReceived += gate.OnHostKeyReceived;

        try
        {
            await client.ConnectAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (Exception exception) when (exception is SshConnectionException or SshAuthenticationException)
        {
            client.Dispose();

            // Translate a refusal we caused ourselves into something the caller can act on. Without
            // this the user sees "connection lost" for what is really "do you trust this key?".
            throw gate.TranslateFailure() ?? exception;
        }
        catch
        {
            client.Dispose();
            throw;
        }

        return new SshNetConnection(client, gate.Presented!);
    }

    /// <summary>
    /// Decides host key trust during the handshake, and remembers enough to explain a refusal.
    /// </summary>
    /// <remarks>
    /// Separate from the connect method because the decision is the security-relevant part and reads
    /// better on its own: pinned and equal accepts, pinned and different refuses as a mismatch,
    /// unpinned refuses as unknown. There is no fourth branch, and there is no prompt.
    /// </remarks>
    private sealed class HostKeyGate(
        IKnownHostStore knownHosts,
        SshConnectionRequest request,
        CancellationToken cancellationToken)
    {
        /// <summary>What the server offered, once the handshake has reached that point.</summary>
        public HostKeyPresentation? Presented { get; private set; }

        private string? pinned;
        private bool mismatch;

        public void OnHostKeyReceived(object? sender, HostKeyEventArgs e)
        {
            var presentation = new HostKeyPresentation(
                request.Host,
                request.Port,
                e.HostKeyName,
                SshHostKeyFingerprint.Format(e.HostKey));

            Presented = presentation;

            // Looked up here rather than before connecting, because the negotiated algorithm is only
            // known now and a server may choose a different one than it did last time.
            //
            // This is the one place the design cannot stay asynchronous: SSH.NET raises host key
            // verification synchronously. It is a local store read rather than a UI round trip, and
            // making the store synchronous instead would rule out a vault-backed implementation.
            pinned = knownHosts
                .FindAsync(request.Host, request.Port, e.HostKeyName, cancellationToken)
                .AsTask()
                .GetAwaiter()
                .GetResult();

            if (pinned is null)
            {
                // Refused, not prompted. The caller decides, off the handshake thread.
                e.CanTrust = false;
                return;
            }

            var matches = SshHostKeyFingerprint.Equal(pinned, presentation.Fingerprint);
            mismatch = !matches;
            e.CanTrust = matches;
        }

        /// <summary>The specific exception for a refusal this gate caused, or null if it did not.</summary>
        public Exception? TranslateFailure()
        {
            if (Presented is not { } presentation)
            {
                return null;
            }

            if (mismatch && pinned is { } pin)
            {
                return new SshHostKeyMismatchException(presentation, pin);
            }

            return pinned is null ? new SshHostKeyUnknownException(presentation) : null;
        }
    }

    /// <remarks>
    /// The known-host lookup inside the synchronous event is the one place this design cannot avoid
    /// blocking. It is a local store read rather than a UI round trip, and the alternative — making
    /// the store synchronous — would rule out the encrypted vault-backed implementation entirely.
    /// </remarks>
    private static ConnectionInfo BuildConnectionInfo(SshConnectionRequest request)
    {
        AuthenticationMethod method = request.Credential switch
        {
            SshPasswordCredential password =>
                new PasswordAuthenticationMethod(request.Username, password.Password),

            SshPrivateKeyCredential key => new PrivateKeyAuthenticationMethod(
                request.Username,
                CreatePrivateKeyFile(key)),

            _ => throw new NotSupportedException(
                $"Credential type {request.Credential.GetType().Name} is not supported."),
        };

        return new ConnectionInfo(request.Host, request.Port, request.Username, method)
        {
            Timeout = request.ConnectTimeout ?? DefaultConnectTimeout,
        };
    }

    private static PrivateKeyFile CreatePrivateKeyFile(SshPrivateKeyCredential credential)
    {
        using var stream = new MemoryStream(credential.PrivateKeyPem, writable: false);

        return credential.Passphrase is null
            ? new PrivateKeyFile(stream)
            : new PrivateKeyFile(stream, credential.Passphrase);
    }
}

/// <summary>An SSH.NET-backed connection.</summary>
internal sealed class SshNetConnection(SshClient client, HostKeyPresentation hostKey) : ISshConnection
{
    /// <inheritdoc />
    public bool IsConnected => client.IsConnected;

    /// <inheritdoc />
    public HostKeyPresentation HostKey { get; } = hostKey;

    /// <inheritdoc />
    public Task<ISshShellSession> OpenShellAsync(TerminalSize size, CancellationToken cancellationToken)
    {
        cancellationToken.ThrowIfCancellationRequested();

        var effective = size.IsUsable ? size : TerminalSize.Default;

        // 4 KiB read buffer inside SSH.NET. Output is coalesced a layer up, so a larger buffer here
        // only delays the first byte reaching the screen.
        var shell = client.CreateShellStream(
            "xterm-256color",
            effective.Columns,
            effective.Rows,
            effective.PixelWidth,
            effective.PixelHeight,
            4096);

        return Task.FromResult<ISshShellSession>(new SshNetShellSession(shell));
    }

    /// <inheritdoc />
    public ValueTask DisposeAsync()
    {
        client.Dispose();
        return ValueTask.CompletedTask;
    }
}

/// <summary>An SSH.NET-backed shell session.</summary>
internal sealed class SshNetShellSession(ShellStream shell) : ISshShellSession
{
    /// <inheritdoc />
    public bool IsOpen => shell.CanRead;

    /// <inheritdoc />
    /// <remarks>
    /// <c>ShellStream</c> does not override <c>ReadAsync</c>, so the base <see cref="Stream"/>
    /// implementation runs the blocking read on a thread-pool thread. Every idle session therefore
    /// parks one thread; see docs/platform-flags.md.
    /// </remarks>
    public ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken) =>
        shell.ReadAsync(buffer, cancellationToken);

    /// <inheritdoc />
    /// <remarks>
    /// The flush is mandatory, not an optimisation. <c>ShellStream.Write</c> accumulates into an
    /// internal buffer and sends nothing until flushed, so without this a keystroke is accepted,
    /// reported as written, and never reaches the remote — the terminal simply stops responding to
    /// input while still displaying output perfectly. SSH.NET's own <c>WriteLine</c> flushes for this
    /// reason, which is why the spike tests never hit it.
    /// <para>
    /// Flushed per write rather than batched: a terminal has to put a keystroke on the wire
    /// immediately, and there is nothing to coalesce — a human types far below any rate at which
    /// batching would matter.
    /// </para>
    /// </remarks>
    public async ValueTask WriteAsync(ReadOnlyMemory<byte> data, CancellationToken cancellationToken)
    {
        await shell.WriteAsync(data, cancellationToken).ConfigureAwait(false);
        await shell.FlushAsync(cancellationToken).ConfigureAwait(false);
    }

    /// <inheritdoc />
    public void Resize(TerminalSize size)
    {
        if (!size.IsUsable)
        {
            // A collapsed pane or a minimised window produces these. Forwarding one leaves the
            // remote's idea of the terminal nonsensical until the next resize arrives.
            return;
        }

        shell.ChangeWindowSize(size.Columns, size.Rows, size.PixelWidth, size.PixelHeight);
    }

    /// <inheritdoc />
    public async ValueTask DisposeAsync()
    {
        await shell.DisposeAsync().ConfigureAwait(false);
    }
}

/// <summary>Convenience helpers over a shell session.</summary>
public static class SshShellSessionExtensions
{
    /// <summary>Writes UTF-8 text to the remote.</summary>
    public static ValueTask WriteTextAsync(
        this ISshShellSession session,
        string text,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(session);

        return session.WriteAsync(Encoding.UTF8.GetBytes(text), cancellationToken);
    }
}