DodoSSH/src/DodoSSH.Client.Session/VaultKnownHostStore.cs

using DodoSSH.Client.Domain;
using DodoSSH.Client.Ssh;
using DodoSSH.Client.Sync;

namespace DodoSSH.Client.Session;

/// <summary>
/// Host key trust kept in the vault, answered from memory.
/// </summary>
/// <remarks>
/// <para>
/// <b>Why a snapshot and not a lookup.</b> <see cref="FindAsync"/> is called from inside SSH.NET's
/// synchronous host key event, which the connection factory has no choice but to block on — see
/// <c>SshNetConnectionFactory</c>, where the comment explains why that cannot be avoided. A store that read
/// SQLite and ran an AEAD open per lookup would put a disk round trip and a decryption on the thread
/// completing the key exchange, once per host key offered, on every connection. So the vault is read when it
/// unlocks and after each synchronisation pass, and the handshake gets a dictionary lookup.
/// </para>
/// <para>
/// <b>What that costs, stated rather than hidden.</b> Trust recorded on another machine is not visible until
/// the next pass brings it down, which is within the minute the shell already syncs on. The failure that
/// causes is a first-contact prompt for a host somebody else approved seconds ago — a prompt the user can
/// answer correctly, since the fingerprint is on screen — and it resolves itself. The reverse mistake would
/// be the serious one, and it cannot happen here: a pin recorded on this machine goes into the snapshot as
/// part of recording it, and invalidates any read that was already in flight.
/// </para>
/// <para>
/// <b>Process-lifetime object, session-scoped contents.</b> The connection factory is composed once, at
/// startup, and outlives every unlock; the vault behind this store does not. So the lifecycle is explicit:
/// <see cref="OpenAsync"/> when a vault unlocks, <see cref="RefreshAsync"/> after a pass,
/// <see cref="Close"/> when it locks. While closed every lookup answers "not pinned", which refuses
/// connections rather than allowing them — the safe direction for the one case that can reach it, a vault
/// locked while a handshake was in flight.
/// </para>
/// </remarks>
public sealed class VaultKnownHostStore : IKnownHostStore
{
    private readonly Lock gate = new();

    private Dictionary<string, PinnedHostKey> pins = new(KnownHostIdentity.Comparer);
    private Binding? binding;

    /// <summary>
    /// Bumped whenever the snapshot changes underneath a read, so a read in flight cannot install a stale
    /// answer over a newer one.
    /// </summary>
    /// <remarks>
    /// The races are ordinary rather than theoretical. A background pass ends with a refresh, and between
    /// that refresh's listing and its assignment the user may lock the vault, approve a new host key, or
    /// withdraw trust from one — and every one of those would otherwise be undone a moment later by the
    /// arriving snapshot.
    /// </remarks>
    private int generation;

    /// <summary>Whether a vault is open behind this store.</summary>
    public bool IsOpen
    {
        get
        {
            lock (gate)
            {
                return binding is not null;
            }
        }
    }

    /// <summary>Reads an unlocked vault's pins, and starts writing new ones to it.</summary>
    /// <param name="session">The unlocked session. Its active vault is the one used.</param>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <remarks>
    /// The previous session's pins are dropped before the new vault is read, not after. Between the two
    /// every host looks unvisited, which is one listing long and errs towards asking; keeping them would
    /// mean one account's trust decisions briefly answering for another's.
    /// </remarks>
    public async Task OpenAsync(VaultSession session, CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(session);

        var opened = new Binding(session.KnownHosts, session.ActiveVaultId);
        int stamp;

        lock (gate)
        {
            binding = opened;
            pins = new Dictionary<string, PinnedHostKey>(KnownHostIdentity.Comparer);
            stamp = ++generation;
        }

        var snapshot = await ReadAsync(opened, cancellationToken).ConfigureAwait(false);

        Install(stamp, snapshot);
    }

    /// <summary>
    /// Re-reads the vault, picking up trust recorded on another machine.
    /// </summary>
    /// <remarks>
    /// Does nothing while closed, so a synchronisation pass that finishes after the vault was locked cannot
    /// bring its contents back.
    /// </remarks>
    public async Task RefreshAsync(CancellationToken cancellationToken)
    {
        Binding? current;
        int stamp;

        lock (gate)
        {
            current = binding;
            stamp = generation;
        }

        if (current is null)
        {
            return;
        }

        var snapshot = await ReadAsync(current, cancellationToken).ConfigureAwait(false);

        Install(stamp, snapshot);
    }

    /// <summary>Forgets the vault and everything read from it. What locking means here.</summary>
    public void Close()
    {
        lock (gate)
        {
            binding = null;
            pins = new Dictionary<string, PinnedHostKey>(KnownHostIdentity.Comparer);
            generation++;
        }
    }

    /// <inheritdoc />
    public ValueTask<string?> FindAsync(
        string host,
        int port,
        string algorithm,
        CancellationToken cancellationToken)
    {
        lock (gate)
        {
            var identity = KnownHostIdentity.For(host, port, algorithm);

            return ValueTask.FromResult(pins.GetValueOrDefault(identity)?.Secret.Fingerprint);
        }
    }

    /// <inheritdoc />
    /// <remarks>From the snapshot, like <see cref="FindAsync"/>, and for the same reason.</remarks>
    public ValueTask<IReadOnlyList<HostKeyPresentation>> FindForEndpointAsync(
        string host,
        int port,
        CancellationToken cancellationToken)
    {
        lock (gate)
        {
            IReadOnlyList<HostKeyPresentation> found = pins.Values
                .Where(pin =>
                    KnownHostIdentity.Comparer.Equals(pin.Secret.Host, host) && pin.Secret.Port == port)
                .Select(pin => new HostKeyPresentation(
                    pin.Secret.Host, pin.Secret.Port, pin.Secret.Algorithm, pin.Secret.Fingerprint))
                .ToArray();

            return ValueTask.FromResult(found);
        }
    }

    /// <inheritdoc />
    /// <exception cref="InvalidOperationException">
    /// The vault is not open, so there is nowhere to record trust.
    /// </exception>
    public async ValueTask TrustAsync(
        HostKeyPresentation presentation,
        CancellationToken cancellationToken)
    {
        ArgumentNullException.ThrowIfNull(presentation);

        var bound = Require();

        var identity = KnownHostIdentity.For(
            presentation.Host, presentation.Port, presentation.Algorithm);

        var existing = Pinned(identity);

        if (existing is not null
            && SshHostKeyFingerprint.Equal(existing.Secret.Fingerprint, presentation.Fingerprint))
        {
            // Already trusted, to the byte. Writing it again would queue an outbox operation that changes
            // nothing and push it to every other machine as a modification.
            return;
        }

        var pin = new KnownHostSecret
        {
            Host = presentation.Host,
            Port = presentation.Port,
            Algorithm = presentation.Algorithm,
            Fingerprint = presentation.Fingerprint,
        };

        var entityId = await StoreAsync(bound, existing, pin, cancellationToken).ConfigureAwait(false);

        lock (gate)
        {
            if (!ReferenceEquals(binding, bound))
            {
                // The vault was locked, or another one was opened, while this was being written. The item is
                // in that vault's outbox and will be there when it is next opened; it must not answer for
                // whatever is open now.
                return;
            }

            pins[identity] = new PinnedHostKey(entityId, pin);
            generation++;
        }
    }

    /// <inheritdoc />
    /// <exception cref="InvalidOperationException">
    /// The vault is not open, so there is nothing to forget.
    /// </exception>
    public async ValueTask<int> ForgetAsync(
        string host,
        int port,
        CancellationToken cancellationToken)
    {
        var bound = Require();

        // Read from the vault rather than from the snapshot, and this is the one operation that must:
        // withdrawing trust has to reach every pin for the endpoint, including a duplicate the snapshot
        // shadowed. A pin left behind is a host that keeps refusing to connect for a reason the user
        // believes they have already dealt with.
        var listing = await bound.KnownHosts
            .ListAsync(bound.VaultId, cancellationToken)
            .ConfigureAwait(false);

        var doomed = listing.Items
            .Where(item =>
                KnownHostIdentity.Comparer.Equals(item.Secret.Host, host) && item.Secret.Port == port)
            .ToArray();

        foreach (var item in doomed)
        {
            await bound.KnownHosts
                .DeleteAsync(bound.VaultId, item.EntityId, cancellationToken)
                .ConfigureAwait(false);
        }

        Invalidate();

        // Re-read rather than patched. The next listing is what these deletions mean, and reproducing that
        // arithmetic against the snapshot is how the two would come to disagree.
        await RefreshAsync(cancellationToken).ConfigureAwait(false);

        return doomed.Length;
    }

    /// <summary>The vault this store writes to, and which vault inside it.</summary>
    private sealed record Binding(KnownHostRepository KnownHosts, Guid VaultId);

    /// <summary>A pin, and the vault item it came from, so re-trusting updates rather than duplicates.</summary>
    /// <param name="EntityId">The item holding this pin.</param>
    /// <param name="Secret">The pin.</param>
    private sealed record PinnedHostKey(Guid EntityId, KnownHostSecret Secret)
    {
        /// <summary>Whether this build may re-encode the item, or a newer one wrote it.</summary>
        internal bool IsWritable { get; init; } = true;
    }

    /// <summary>
    /// Records a pin, updating the item that already held one for this endpoint where there is one.
    /// </summary>
    /// <remarks>
    /// An item a newer client wrote is left alone and a fresh one is created beside it. Re-encoding it would
    /// drop fields this build has no concept of, which is the rule the whole client follows for read-only
    /// items — and refusing outright, as the editors do, would leave the user unable to connect to a rebuilt
    /// server at all. The new item wins the lookup by the tie-break in <see cref="ReadAsync"/>, and a client
    /// that understands both can reconcile them.
    /// </remarks>
    private static async Task<Guid> StoreAsync(
        Binding bound,
        PinnedHostKey? existing,
        KnownHostSecret pin,
        CancellationToken cancellationToken)
    {
        if (existing is { IsWritable: true } writable)
        {
            await bound.KnownHosts
                .UpdateAsync(bound.VaultId, writable.EntityId, pin, cancellationToken)
                .ConfigureAwait(false);

            return writable.EntityId;
        }

        return await bound.KnownHosts
            .CreateAsync(bound.VaultId, pin, cancellationToken)
            .ConfigureAwait(false);
    }

    /// <summary>
    /// Reads every readable pin in the vault into a lookup.
    /// </summary>
    /// <remarks>
    /// <para>
    /// Two items can name the same endpoint and algorithm: two machines that first met a host while unable
    /// to reach each other each minted one. Where they agree — the ordinary case, since they saw the same
    /// server — the duplicate is invisible. Where they do not, the later item wins, ordered by an id that is
    /// a UUIDv7 and therefore by when the trust was recorded. Any total order would do for correctness; what
    /// matters is that every machine picks the same one, and that the wrong choice is recoverable rather than
    /// permanent, which <see cref="ForgetAsync"/> makes it.
    /// </para>
    /// <para>
    /// A pin that will not decrypt is skipped, and its endpoint then looks unvisited. That is the safe
    /// reading: an unreadable pin cannot be compared against anything, so the only honest answers are "ask
    /// the user" and "refuse", and asking is the one that leaves them a way forward. The count is not lost —
    /// the vault view reports undecryptable items of every kind.
    /// </para>
    /// </remarks>
    private static async Task<Dictionary<string, PinnedHostKey>> ReadAsync(
        Binding bound,
        CancellationToken cancellationToken)
    {
        var listing = await bound.KnownHosts
            .ListAsync(bound.VaultId, cancellationToken)
            .ConfigureAwait(false);

        var snapshot = new Dictionary<string, PinnedHostKey>(KnownHostIdentity.Comparer);

        foreach (var item in listing.Items.OrderBy(item => item.EntityId))
        {
            var identity = KnownHostIdentity.For(
                item.Secret.Host, item.Secret.Port, item.Secret.Algorithm);

            snapshot[identity] = new PinnedHostKey(item.EntityId, item.Secret)
            {
                IsWritable = !item.IsReadOnly,
            };
        }

        return snapshot;
    }

    /// <remarks>
    /// Replaces the snapshot wholesale rather than merging into it, which is what makes a pin withdrawn on
    /// another machine actually disappear here. Anything recorded on this machine while the read was in
    /// flight has already bumped the generation, so it is this snapshot that gets dropped and not that pin.
    /// </remarks>
    private void Install(int stamp, Dictionary<string, PinnedHostKey> snapshot)
    {
        lock (gate)
        {
            if (stamp == generation)
            {
                pins = snapshot;
            }
        }
    }

    private void Invalidate()
    {
        lock (gate)
        {
            generation++;
        }
    }

    private PinnedHostKey? Pinned(string identity)
    {
        lock (gate)
        {
            return pins.GetValueOrDefault(identity);
        }
    }

    private Binding Require()
    {
        lock (gate)
        {
            return binding
                ?? throw new InvalidOperationException(
                    "The vault is locked, so host key trust cannot be changed.");
        }
    }
}