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, ISftpSessionFactory
{
    private static readonly TimeSpan DefaultConnectTimeout = TimeSpan.FromSeconds(15);

    /// <summary>
    /// How much of a file SSH.NET reads or writes per SFTP request.
    /// </summary>
    /// <remarks>
    /// SSH.NET's default is 32 KiB, which is a request per 32 KiB and a round trip's latency between each on
    /// a link the window would happily keep full. 64 KiB is the largest an OpenSSH server accepts without
    /// negotiation, so it is the ceiling rather than a guess — anything above it is answered with a shorter
    /// read, which SSH.NET handles but which buys nothing.
    /// </remarks>
    private const uint SftpBufferSize = 64 * 1024;

    /// <summary>Where a <see cref="SshLoopbackProxy"/> is, and the only address one is ever dialled at.</summary>
    /// <remarks>
    /// The literal rather than <c>IPAddress.Loopback.ToString()</c>, and rather than "localhost": SSH.NET
    /// takes the proxy host as a string and resolves it, so a name would put a DNS lookup — and whatever
    /// the machine's hosts file says <c>localhost</c> means — inside the connect path of every proxied
    /// connection. It is also the half of the loopback promise this assembly can keep on its own; the other
    /// half is that the proxy bound there, which is the caller's to get right.
    /// </remarks>
    private const string LoopbackAddress = "127.0.0.1";

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

        var client = new SshClient(BuildConnectionInfo(request));

        var gate = await ConnectThroughHostKeyGateAsync(client, request, cancellationToken)
            .ConfigureAwait(false);

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

    /// <inheritdoc />
    /// <remarks>
    /// A second connection to the host rather than a second channel on one that may already be open — see
    /// <see cref="ISftpSession"/> for why SSH.NET leaves no choice. Everything that guards a shell guards this
    /// too, because it is the same handshake: the same host key gate, the same pin, the same two refusals.
    /// </remarks>
    public async Task<ISftpSession> OpenSftpAsync(
        SshConnectionRequest request,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(request);

        var client = new SftpClient(BuildConnectionInfo(request)) { BufferSize = SftpBufferSize };

        var gate = await ConnectThroughHostKeyGateAsync(client, request, cancellationToken)
            .ConfigureAwait(false);

        // Read once, here, rather than per call. SftpClient.WorkingDirectory canonicalises against the server
        // on first read, so leaving it to the property would put a round trip behind something that reads
        // like a field — and the session's own remark promises this is the one path known without asking.
        string home;

        try
        {
            home = client.WorkingDirectory;
        }
        catch
        {
            client.Dispose();
            throw;
        }

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

    /// <summary>
    /// Runs the handshake with host key trust attached, and translates a refusal this factory caused.
    /// </summary>
    /// <remarks>
    /// Shared by the shell and the file-transfer paths over <c>BaseClient</c>, which is where SSH.NET puts
    /// both <c>ConnectAsync</c> and <c>HostKeyReceived</c>. The alternative was the same twelve lines twice,
    /// and the half worth getting wrong is the translation: without it a user who has never seen a host is
    /// told the connection was lost.
    /// </remarks>
    private async Task<HostKeyGate> ConnectThroughHostKeyGateAsync(
        BaseClient client,
        SshConnectionRequest request,
        CancellationToken cancellationToken)
    {
        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 gate;
    }

    /// <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>
    /// <para>
    /// 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.
    /// </para>
    /// <para>
    /// ◆ <b>A proxied connection differs here and nowhere else.</b> The host, the port, the account and the
    /// credential are the target's either way, and so is everything the gate above reads — which is what
    /// makes a host reached through a bastion or a relay get pinned under its own name rather than under
    /// <c>127.0.0.1</c>. The proxy is a route, not a destination, and this is the one method that needs to
    /// know the difference. See <see cref="SshLoopbackProxy"/>.
    /// </para>
    /// <para>
    /// <c>Timeout</c> is assigned after the branch rather than in two initialisers, because it covers the
    /// whole of getting there — the proxy handshake included — and having it stated once is what stops the
    /// two paths quietly drifting to different waits.
    /// </para>
    /// </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."),
        };

        var info = request.Proxy is { } proxy
            ? new ConnectionInfo(
                request.Host,
                request.Port,
                request.Username,
                ProxyTypes.Socks5,
                LoopbackAddress,
                proxy.Port,

                // No proxy credentials, and empty rather than null: SSH.NET offers username/password
                // authentication to a SOCKS5 server only when it has been given one, and both proxies this
                // client will ever use are on its own loopback interface, where a password would be a
                // secret shared between two halves of the same process.
                string.Empty,
                string.Empty,
                method)
            : new ConnectionInfo(request.Host, request.Port, request.Username, method);

        info.Timeout = request.ConnectTimeout ?? DefaultConnectTimeout;

        return info;
    }

    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);
    }
}