DodoSSH/src/DodoSSH.Client.Sync/VaultKeyring.cs

using System.Diagnostics.CodeAnalysis;
using System.Security.Cryptography;
using DodoSSH.Client.Storage;
using DodoSSH.Crypto;

namespace DodoSSH.Client.Sync;

/// <summary>
/// The vault keys held for the duration of an unlocked session.
/// </summary>
/// <remarks>
/// <para>
/// One place that holds plaintext vault keys, so there is one place that clears them. Every store and
/// every cipher call borrows a key from here rather than keeping a copy, which is what makes
/// "the keys exist only while unlocked" a property of the code and not of everyone's discipline.
/// </para>
/// <para>
/// <b>A vault has a key per generation, and this holds every one it was granted.</b> A rotation does not
/// re-encrypt what is already stored — each item keeps the generation it was sealed under — so reading a
/// rotated vault means opening items under two or three different keys, chosen per item rather than per
/// vault. Writing uses the newest, which is what <see cref="TryGet"/> answers; reading an item asks for
/// the generation that item names, which is <see cref="TryGetAt"/>.
/// </para>
/// <para>
/// A grant that will not open is not an error: it means the vault has been rekeyed and this client's
/// grant has not been re-wrapped yet, or the grant was fabricated. Both leave the vault temporarily
/// unreadable and both are reported rather than thrown, so one bad grant does not take the other vaults
/// down with it.
/// </para>
/// </remarks>
public sealed class VaultKeyring : IDisposable
{
    private readonly Dictionary<Guid, Dictionary<uint, byte[]>> keys = [];
    private readonly Dictionary<Guid, uint> generations = [];
    private bool disposed;

    private VaultKeyring()
    {
    }

    /// <summary>Vaults whose grant could not be opened, and which are therefore unreadable.</summary>
    public IReadOnlyList<Guid> Unopened { get; private set; } = [];

    /// <summary>
    /// Opens every grant the bundle can.
    /// </summary>
    /// <param name="bundle">The unlocked identity keys.</param>
    /// <param name="vaults">The cached vault list, each with its wrapped key.</param>
    public static VaultKeyring Open(UserSecretBundle bundle, IReadOnlyList<StoredVault> vaults)
    {
        ArgumentNullException.ThrowIfNull(bundle);
        ArgumentNullException.ThrowIfNull(vaults);

        var keyring = new VaultKeyring();
        var unopened = new List<Guid>();

        try
        {
            foreach (var vault in vaults)
            {
                // The history first, and never conditional on the current generation opening. A member
                // who has been rotated past but not yet re-wrapped can still read everything written
                // before the rotation, and dropping those keys because the newest grant is missing
                // would turn "you cannot see the last hour's changes" into "the vault is empty".
                keyring.OpenPriorWraps(bundle, vault);

                if (vault.WrappedVaultKey is null)
                {
                    // The server said so itself: a grant awaiting re-wrap after a rekey.
                    unopened.Add(vault.VaultId);
                    continue;
                }

                var key = VaultKeys.TryUnwrap(
                    bundle.EncryptionKey,
                    vault.WrappedVaultKey,
                    vault.VaultId,
                    vault.KeyGeneration);

                if (key is null)
                {
                    unopened.Add(vault.VaultId);
                    continue;
                }

                keyring.Adopt(vault.VaultId, key, vault.KeyGeneration);
            }

            keyring.Unopened = unopened;
            return keyring;
        }
        catch
        {
            keyring.Dispose();
            throw;
        }
    }

    /// <summary>
    /// Takes a vault key this session has just generated.
    /// </summary>
    /// <param name="vaultId">The vault.</param>
    /// <param name="vaultKey">
    /// The plaintext key. <b>The keyring takes ownership</b> and zeroes it on disposal; the caller must
    /// not keep a reference or zero it itself.
    /// </param>
    /// <param name="keyGeneration">The generation this key is for.</param>
    /// <remarks>
    /// <para>
    /// Two cases, and they are the same operation: creating a team vault, and rotating one. Both
    /// generate the key here, wrap it to this user and send the wrap, so the plaintext is already in
    /// this process and unwrapping the server's copy back would be a round trip to learn something it
    /// just chose. Adopting it also means the vault is usable immediately rather than at the next
    /// unlock, which is what somebody who has just pressed a button expects.
    /// </para>
    /// <para>
    /// This generation becomes the one writes are sealed under. A key for a generation the vault has
    /// moved <em>past</em> goes in through <see cref="AdoptPrior"/> instead, which is not the same
    /// operation: it makes old items readable and must not walk the write target backwards.
    /// </para>
    /// </remarks>
    public void Adopt(Guid vaultId, byte[] vaultKey, uint keyGeneration)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        ArgumentNullException.ThrowIfNull(vaultKey);

        Store(vaultId, vaultKey, keyGeneration);

        Promote(vaultId, keyGeneration);
    }

    /// <summary>
    /// Takes a vault key for a generation the vault has already moved past.
    /// </summary>
    /// <param name="vaultId">The vault.</param>
    /// <param name="vaultKey">
    /// The plaintext key. <b>The keyring takes ownership</b>, exactly as <see cref="Adopt"/> does.
    /// </param>
    /// <param name="keyGeneration">The superseded generation this key opens.</param>
    /// <remarks>
    /// Holding one of these is what lets a rotated vault be read at all: items are not re-encrypted by a
    /// rotation, so everything written before it is still sealed under the key it was written with.
    /// Nothing is ever <em>written</em> under one, which is why this does not touch the current
    /// generation and does not make an otherwise unreadable vault readable.
    /// </remarks>
    public void AdoptPrior(Guid vaultId, byte[] vaultKey, uint keyGeneration)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        ArgumentNullException.ThrowIfNull(vaultKey);

        Store(vaultId, vaultKey, keyGeneration);
    }

    /// <summary>
    /// Tries to open a vault that has become readable since the session was unlocked.
    /// </summary>
    /// <returns>Whether the grant opened.</returns>
    /// <remarks>
    /// What a share looks like from the receiving end: the vault was in the list all along, listed and
    /// unreadable, and a member holding Share has now wrapped its key. Re-opening it here rather than
    /// waiting for a relock is the difference between "someone shared a vault with you" arriving and
    /// arriving tomorrow.
    /// </remarks>
    public bool TryAdmit(UserSecretBundle bundle, StoredVault vault)
    {
        ObjectDisposedException.ThrowIf(disposed, this);
        ArgumentNullException.ThrowIfNull(bundle);
        ArgumentNullException.ThrowIfNull(vault);

        // Attempted whatever happens to the current generation, and before it. A share of a vault that
        // has been rotated since it was created arrives as a current wrap plus its history, and the
        // history is not a consolation prize — without it the recipient sees a vault full of items that
        // will not decrypt.
        OpenPriorWraps(bundle, vault);

        if (vault.WrappedVaultKey is null)
        {
            return false;
        }

        if (Held(vault.VaultId, vault.KeyGeneration) is not null)
        {
            // Already open at this generation. Promoted rather than returned early, because a vault
            // that was rotated and re-granted arrives here with a generation this keyring has been
            // treating as historic, and it is now the one writes belong under.
            Promote(vault.VaultId, vault.KeyGeneration);
            return true;
        }

        var key = VaultKeys.TryUnwrap(
            bundle.EncryptionKey, vault.WrappedVaultKey, vault.VaultId, vault.KeyGeneration);

        if (key is null)
        {
            return false;
        }

        Adopt(vault.VaultId, key, vault.KeyGeneration);

        return true;
    }

    /// <summary>Records that a vault cannot be read, so the interface can say so.</summary>
    /// <remarks>
    /// <para>
    /// The counterpart of <see cref="TryAdmit"/> for the case where the grant did not open. Kept
    /// explicit rather than inferred from the absence of a key, because "no key" is also what a vault
    /// this session has never heard of looks like.
    /// </para>
    /// <para>
    /// <b>It also gives up the write target, and that is the load-bearing half.</b> The usual way to
    /// reach here is another client having rotated the vault: this session still holds the previous
    /// generation's key and it is no longer the current one. Going on treating it as current would seal
    /// new items under a superseded key — readable here, unreadable to everybody else, and with nothing
    /// to show the author that anything was wrong. The keys themselves are kept, because the items
    /// already written under them are still readable through <see cref="TryGetAt"/>.
    /// </para>
    /// </remarks>
    public void MarkUnreadable(Guid vaultId)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        generations.Remove(vaultId);

        if (!Unopened.Contains(vaultId))
        {
            Unopened = [.. Unopened, vaultId];
        }
    }

    /// <summary>
    /// Borrows a vault's current key: the one new items are sealed under.
    /// </summary>
    /// <remarks>
    /// The returned memory is the keyring's own buffer, not a copy, and is zeroed when the keyring is
    /// disposed. Callers must not retain it past the operation they borrowed it for.
    /// <para>
    /// False for a vault this session holds only the history of — one rotated past a grant that has not
    /// been re-wrapped yet. That is deliberate: writing under a superseded key would produce an item
    /// nobody else could read, and the honest answer is that the vault is not writable until the new
    /// key arrives.
    /// </para>
    /// </remarks>
    public bool TryGet(Guid vaultId, out ReadOnlyMemory<byte> vaultKey, out uint keyGeneration)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        if (generations.TryGetValue(vaultId, out var current)
            && Held(vaultId, current) is { } key)
        {
            vaultKey = key;
            keyGeneration = current;
            return true;
        }

        vaultKey = default;
        keyGeneration = 0;
        return false;
    }

    /// <summary>
    /// Borrows the key one particular generation of a vault was sealed under.
    /// </summary>
    /// <param name="vaultId">The vault.</param>
    /// <param name="keyGeneration">The generation the item names.</param>
    /// <param name="vaultKey">The key, borrowed on the same terms as <see cref="TryGet"/>.</param>
    /// <returns>Whether this session holds that generation.</returns>
    /// <remarks>
    /// What every read goes through, because an item names the generation it was sealed under and a
    /// rotated vault holds items from more than one. False means that item is unreadable here and says
    /// nothing about the rest of the vault — which is why a caller counts it rather than failing.
    /// </remarks>
    public bool TryGetAt(Guid vaultId, uint keyGeneration, out ReadOnlyMemory<byte> vaultKey)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        if (Held(vaultId, keyGeneration) is { } key)
        {
            vaultKey = key;
            return true;
        }

        vaultKey = default;
        return false;
    }

    /// <summary>
    /// Every generation of one vault's key that this session holds, oldest first.
    /// </summary>
    /// <remarks>
    /// Read when sharing: a recipient given only the newest key would find the vault's history
    /// undecryptable, so the sharing client wraps each of these in turn. It is the only party that can
    /// — the server holds ciphertext, and the recipient holds nothing yet.
    /// </remarks>
    public IReadOnlyList<uint> GenerationsHeld(Guid vaultId)
    {
        ObjectDisposedException.ThrowIf(disposed, this);

        return keys.TryGetValue(vaultId, out var held) ? [.. held.Keys.Order()] : [];
    }

    /// <summary>Whether this vault can be read and written at its current generation.</summary>
    public bool CanRead(Guid vaultId) =>
        !disposed
        && generations.TryGetValue(vaultId, out var current)
        && Held(vaultId, current) is not null;

    /// <inheritdoc />
    public void Dispose()
    {
        if (disposed)
        {
            return;
        }

        disposed = true;

        foreach (var held in keys.Values)
        {
            foreach (var key in held.Values)
            {
                CryptographicOperations.ZeroMemory(key);
            }
        }

        keys.Clear();
        generations.Clear();
    }

    /// <summary>Opens whatever superseded generations this vault came with.</summary>
    /// <remarks>
    /// A wrap that will not open is skipped rather than reported. It means one historic grant is
    /// unusable — the items under that generation stay unreadable and are counted as such where they
    /// are listed — and it is not a reason to refuse the generations that did open.
    /// </remarks>
    private void OpenPriorWraps(UserSecretBundle bundle, StoredVault vault)
    {
        foreach (var wrap in vault.PriorKeyWraps ?? [])
        {
            if (Held(vault.VaultId, wrap.KeyGeneration) is not null)
            {
                continue;
            }

            var key = VaultKeys.TryUnwrap(
                bundle.EncryptionKey, wrap.WrappedKey, vault.VaultId, wrap.KeyGeneration);

            if (key is not null)
            {
                AdoptPrior(vault.VaultId, key, wrap.KeyGeneration);
            }
        }
    }

    private byte[]? Held(Guid vaultId, uint keyGeneration) =>
        keys.TryGetValue(vaultId, out var held) && held.TryGetValue(keyGeneration, out var key)
            ? key
            : null;

    private void Store(Guid vaultId, byte[] vaultKey, uint keyGeneration)
    {
        if (!keys.TryGetValue(vaultId, out var held))
        {
            held = [];
            keys[vaultId] = held;
        }

        if (held.TryGetValue(keyGeneration, out var previous))
        {
            CryptographicOperations.ZeroMemory(previous);
        }

        held[keyGeneration] = vaultKey;
    }

    private void Promote(Guid vaultId, uint keyGeneration)
    {
        generations[vaultId] = keyGeneration;

        Unopened = [.. Unopened.Where(id => id != vaultId)];
    }
}

/// <summary>Thrown when an operation needs a vault key the keyring does not hold.</summary>
/// <remarks>
/// An exception rather than a silent no-op, because every caller that reaches this point has already
/// been given the chance to check <see cref="VaultKeyring.CanRead"/>. Continuing without the key would
/// mean writing an item nobody can open.
/// </remarks>
[SuppressMessage(
    "Design",
    "CA1032:Implement standard exception constructors",
    Justification = "The vault id is required context; a message-only constructor would lose it.")]
public sealed class VaultUnreadableException(Guid vaultId)
    : InvalidOperationException(
        $"Vault {vaultId} has no usable key. Its grant is missing or awaiting re-wrap after a rekey.")
{
    /// <summary>The vault that cannot be read.</summary>
    public Guid VaultId { get; } = vaultId;
}