DodoSSH/tests/DodoSSH.Client.App.Tests/FakeVaultServer.Teams.cs

using DodoSSH.Client.Api;
using DodoSSH.Contracts;
using DodoSSH.Crypto;

namespace DodoSSH.Client.App.Tests;

/// <summary>
/// The team, directory and grant half of the fake server.
/// </summary>
/// <remarks>
/// <para>
/// <b>The key log is real.</b> Entries are chained with <see cref="KeyLogChain.ComputeEntryHash"/> exactly
/// as the server chains them, because the client refuses to wrap a vault key to a directory answer that
/// does not appear in a log whose chain verifies — so a fake that returned a plausible-looking log would
/// make every sharing test pass against a check that was never exercised. It also means a test can break
/// the chain deliberately and watch the client refuse.
/// </para>
/// <para>
/// Everything else is deliberately thin. Roles, slugs and idempotency are the server's rules and are
/// tested against the real one in <c>DodoSSH.Api.Tests</c>; what the shell needs from here is that a team
/// can be created, a member added, and a vault key wrapped and recorded.
/// </para>
/// </remarks>
internal sealed partial class FakeVaultServer : ITeamApi, IDirectoryApi, IVaultGrantApi
{
    private readonly List<TeamSummary> teams = [];
    private readonly Dictionary<Guid, List<TeamMemberSummary>> members = [];
    private readonly Dictionary<Guid, VaultSummary> teamVaults = [];
    private readonly Dictionary<(Guid VaultId, Guid UserId), IssueVaultGrantRequest> grants = [];
    private readonly List<KeyLogRecord> keyLog = [];
    private readonly List<DirectoryEntry> directory = [];
    private readonly Dictionary<Guid, List<TeamInvitationSummary>> invitations = [];

    /// <inheritdoc />
    public ITeamApi Teams => this;

    /// <inheritdoc />
    public IDirectoryApi Directory => this;

    /// <inheritdoc />
    public IVaultGrantApi Grants => this;

    /// <summary>Grants this fake has been asked to record, for a test to assert on.</summary>
    internal IReadOnlyDictionary<(Guid VaultId, Guid UserId), IssueVaultGrantRequest> IssuedGrants => grants;

    /// <summary>
    /// When true, the log served omits its last entry's link, so its chain no longer verifies.
    /// </summary>
    /// <remarks>
    /// The switch a test flips to prove the client refuses rather than shares. A fake with no way to be
    /// wrong can only ever confirm the happy path.
    /// </remarks>
    internal bool CorruptKeyLog { get; set; }

    /// <summary>Registers another account, as though they had signed in and enrolled here.</summary>
    /// <returns>Their user id.</returns>
    internal Guid AddAccount(string email, string displayName)
    {
        var userId = Guid.CreateVersion7();

        // Real keys rather than filler: the client recomputes the fingerprint over both halves and refuses
        // an entry whose fingerprint does not match, so random bytes would fail for the wrong reason.
        using var bundle = UserSecretBundle.Create(DateTimeOffset.UnixEpoch);

        var sequence = AppendKeyLog(
            userId, bundle.EncryptionPublicKey, bundle.SigningPublicKey, new byte[64]);

        directory.Add(new DirectoryEntry(
            userId,
            email,
            displayName,
            bundle.EncryptionPublicKey,
            bundle.SigningPublicKey,
            DshCrypto.ComputeFingerprint(bundle.EncryptionPublicKey, bundle.SigningPublicKey),
            KeyGeneration: 1,
            KeyLogSequence: sequence));

        return userId;
    }

    /// <inheritdoc />
    public Task<IReadOnlyList<TeamSummary>> ListTeamsAsync(CancellationToken cancellationToken) =>
        Task.FromResult<IReadOnlyList<TeamSummary>>([.. teams]);

    /// <inheritdoc />
    public Task<TeamSummary> CreateTeamAsync(
        CreateTeamRequest request,
        CancellationToken cancellationToken)
    {
        var team = new TeamSummary(
            request.TeamId,
            request.Name,
            request.Slug,
            request.Description,
            TeamMemberRole.Owner,
            MemberCount: 1,
            VaultCount: 0,
            DateTimeOffset.UnixEpoch);

        teams.Add(team);

        members[team.TeamId] =
        [
            new TeamMemberSummary(
                UserId,
                "alice@example.com",
                "Alice Example",
                TeamMemberRole.Owner,
                TeamMemberStatus.Active,
                IsEnrolled: true,
                DateTimeOffset.UnixEpoch,
                DateTimeOffset.UnixEpoch),
        ];

        return Task.FromResult(team);
    }

    /// <inheritdoc />
    public Task<TeamSummary> UpdateTeamAsync(
        Guid teamId,
        UpdateTeamRequest request,
        CancellationToken cancellationToken)
    {
        var index = teams.FindIndex(team => team.TeamId == teamId);

        if (index < 0)
        {
            throw new DodoSshApiException(
                System.Net.HttpStatusCode.NotFound, ProblemCodes.InvalidTeam, "No such team.");
        }

        // The slug is deliberately not touched, matching the server: a rename changes the display
        // name only. A fake that also moved the slug would let a test assert behaviour nothing has.
        teams[index] = teams[index] with
        {
            Name = request.Name,
            Description = request.Description,
        };

        return Task.FromResult(teams[index]);
    }

    /// <inheritdoc />
    /// <remarks>
    /// The vault refusal is reproduced rather than skipped, unlike the other server rules here. It is
    /// the one whose consequence the shell has to render — a status line explaining why nothing
    /// happened — so a fake that always succeeded would leave that path untested.
    /// </remarks>
    public Task<bool> ArchiveTeamAsync(Guid teamId, CancellationToken cancellationToken)
    {
        var index = teams.FindIndex(team => team.TeamId == teamId);

        if (index < 0)
        {
            return Task.FromResult(false);
        }

        if (teamVaults.Values.Any(vault => vault.TeamId == teamId))
        {
            throw new DodoSshApiException(
                System.Net.HttpStatusCode.Conflict,
                ProblemCodes.TeamNotEmpty,
                "This team still owns vaults, and archiving it would take them away from everybody "
                + "holding a key — including you.");
        }

        teams.RemoveAt(index);
        members.Remove(teamId);
        invitations.Remove(teamId);

        return Task.FromResult(true);
    }

    /// <inheritdoc />
    /// <remarks>
    /// Both rows move, because a fake that only promoted the recipient would let a test pass while
    /// the team was owned twice — which is the exact failure the real service uses a transaction to
    /// make impossible.
    /// </remarks>
    public Task TransferTeamOwnershipAsync(
        Guid teamId,
        TransferTeamOwnershipRequest request,
        CancellationToken cancellationToken)
    {
        var list = members.GetValueOrDefault(teamId, []);
        var incoming = list.FindIndex(member => member.UserId == request.UserId);

        if (incoming < 0)
        {
            throw new DodoSshApiException(
                System.Net.HttpStatusCode.BadRequest,
                ProblemCodes.InvalidTeam,
                "That account is not an active member of this team.");
        }

        var outgoing = list.FindIndex(member => member.Role == TeamMemberRole.Owner);

        list[incoming] = list[incoming] with { Role = TeamMemberRole.Owner };

        if (outgoing >= 0)
        {
            list[outgoing] = list[outgoing] with { Role = TeamMemberRole.Admin };
        }

        var index = teams.FindIndex(team => team.TeamId == teamId);

        if (index >= 0)
        {
            teams[index] = teams[index] with { Role = TeamMemberRole.Admin };
        }

        return Task.CompletedTask;
    }

    /// <summary>
    /// When set, a member read waits on it before answering.
    /// </summary>
    /// <remarks>
    /// Every other method here answers from memory and therefore completes before its caller's await
    /// ever suspends, which hides anything the screen only gets wrong while a read is in flight — the
    /// state a real server leaves it in for the length of a round trip. A test that wants that state
    /// holds the gate.
    /// </remarks>
    internal TaskCompletionSource? MemberReadGate { get; set; }

    /// <summary>How many member reads have been asked for, for a test to assert on.</summary>
    internal int MemberReads { get; private set; }

    /// <inheritdoc />
    public async Task<IReadOnlyList<TeamMemberSummary>> ListTeamMembersAsync(
        Guid teamId,
        CancellationToken cancellationToken)
    {
        MemberReads++;

        if (MemberReadGate is { } gate)
        {
            await gate.Task.WaitAsync(cancellationToken).ConfigureAwait(false);
        }

        return members.TryGetValue(teamId, out var list) ? [.. list] : [];
    }

    /// <inheritdoc />
    public Task<TeamMemberSummary> AddTeamMemberAsync(
        Guid teamId,
        AddTeamMemberRequest request,
        CancellationToken cancellationToken)
    {
        var entry = directory.Find(candidate => candidate.UserId == request.UserId)
            ?? throw new DodoSshApiException(
                System.Net.HttpStatusCode.BadRequest,
                ProblemCodes.InvalidTeam,
                "No such account on this server.");

        // LastActiveAt is left null: this account has been added, not seen. The owner's row carries a
        // real one, so both branches of the interface's "last active / never" split are exercised.
        var member = new TeamMemberSummary(
            entry.UserId,
            entry.Email,
            entry.DisplayName,
            request.Role,
            TeamMemberStatus.Active,
            IsEnrolled: true,
            DateTimeOffset.UnixEpoch,
            LastActiveAt: null);

        members[teamId] = [.. members.GetValueOrDefault(teamId, []), member];

        Recount(teamId);

        return Task.FromResult(member);
    }

    /// <inheritdoc />
    public Task<IReadOnlyList<TeamInvitationSummary>> ListTeamInvitationsAsync(
        Guid teamId,
        CancellationToken cancellationToken) =>
        Task.FromResult<IReadOnlyList<TeamInvitationSummary>>(
            invitations.TryGetValue(teamId, out var list) ? [.. list] : []);

    /// <inheritdoc />
    public Task<TeamInvitationSummary> CreateTeamInvitationAsync(
        Guid teamId,
        CreateTeamInvitationRequest request,
        CancellationToken cancellationToken)
    {
        var list = invitations.GetValueOrDefault(teamId, []);

        if (list.Exists(invitation =>
            invitation.State == TeamInvitationState.Pending
            && string.Equals(invitation.Email, request.Email, StringComparison.OrdinalIgnoreCase)))
        {
            throw new DodoSshApiException(
                System.Net.HttpStatusCode.BadRequest,
                ProblemCodes.InvalidTeamInvitation,
                "There is already an invitation to that address for this team.");
        }

        var invited = new TeamInvitationSummary(
            request.InvitationId,
            request.Email,
            request.Role,
            TeamInvitationState.Pending,
            UserId,
            DateTimeOffset.UnixEpoch,
            DateTimeOffset.UnixEpoch.AddDays(14),
            AcceptedAt: null);

        invitations[teamId] = [.. list, invited];

        return Task.FromResult(invited);
    }

    /// <inheritdoc />
    public Task<bool> RevokeTeamInvitationAsync(
        Guid teamId,
        Guid invitationId,
        CancellationToken cancellationToken)
    {
        var list = invitations.GetValueOrDefault(teamId, []);
        var index = list.FindIndex(invitation =>
            invitation.InvitationId == invitationId
            && invitation.State == TeamInvitationState.Pending);

        if (index < 0)
        {
            return Task.FromResult(false);
        }

        // Kept and marked rather than removed, as the server keeps it: the screen has to be able to
        // say an invitation was withdrawn rather than letting it vanish and read as never sent.
        list[index] = list[index] with { State = TeamInvitationState.Revoked };

        return Task.FromResult(true);
    }

    /// <inheritdoc />
    public Task<TeamMemberSummary> ChangeTeamMemberRoleAsync(
        Guid teamId,
        Guid userId,
        ChangeTeamMemberRoleRequest request,
        CancellationToken cancellationToken)
    {
        var list = members.GetValueOrDefault(teamId, []);
        var index = list.FindIndex(member => member.UserId == userId);

        if (index < 0)
        {
            throw new DodoSshApiException(
                System.Net.HttpStatusCode.BadRequest,
                ProblemCodes.InvalidTeam,
                "That account is not an active member of this team.");
        }

        list[index] = list[index] with { Role = request.Role };

        return Task.FromResult(list[index]);
    }

    /// <inheritdoc />
    public Task<bool> RemoveTeamMemberAsync(
        Guid teamId,
        Guid userId,
        CancellationToken cancellationToken)
    {
        var list = members.GetValueOrDefault(teamId, []);
        var removed = list.RemoveAll(member => member.UserId == userId) > 0;

        // Every grant they held from this team goes with them, as the real service revokes them in the
        // same transaction. A fake that removed the membership and left the grants would let a test
        // "prove" a revocation that had not happened.
        foreach (var vaultId in teamVaults.Values
            .Where(vault => vault.TeamId == teamId)
            .Select(vault => vault.VaultId))
        {
            grants.Remove((vaultId, userId));
        }

        Recount(teamId);

        return Task.FromResult(removed);
    }

    /// <inheritdoc />
    public Task<VaultSummary> CreateTeamVaultAsync(
        Guid teamId,
        CreateTeamVaultRequest request,
        CancellationToken cancellationToken)
    {
        var vault = new VaultSummary(
            request.VaultId,
            request.Name,
            IsPersonal: false,
            TeamId: teamId,
            KeyGeneration: 1,
            Permissions: 31,
            request.WrappedVaultKey,
            RekeyRequired: false);

        teamVaults[vault.VaultId] = vault;

        Recount(teamId);

        return Task.FromResult(vault);
    }

    /// <inheritdoc />
    public Task<IReadOnlyList<DirectoryEntry>> LookupByEmailAsync(
        string email,
        CancellationToken cancellationToken) =>
        Task.FromResult<IReadOnlyList<DirectoryEntry>>(
        [
            .. directory.Where(entry =>
                string.Equals(entry.Email, email, StringComparison.OrdinalIgnoreCase)),
        ]);

    /// <inheritdoc />
    public Task<DirectoryEntry?> LookupByIdAsync(Guid userId, CancellationToken cancellationToken) =>
        Task.FromResult(directory.Find(entry => entry.UserId == userId));

    /// <inheritdoc />
    public Task<KeyLogPage> ReadKeyLogAsync(
        long afterSequence,
        int? limit,
        CancellationToken cancellationToken)
    {
        var page = keyLog.Where(entry => entry.Sequence > afterSequence).ToList();

        if (CorruptKeyLog && page.Count > 0)
        {
            // One byte, in the field the chain is built from. Enough to break the link and nothing else,
            // which is what a tampered log would look like.
            var last = page[^1];
            page[^1] = last with { EncryptionPublicKey = [.. last.EncryptionPublicKey.Reverse()] };
        }

        var head = keyLog.Count == 0
            ? KeyLogChain.CreateGenesisPreviousHash()
            : keyLog[^1].Hash;

        return Task.FromResult(new KeyLogPage(page, keyLog.Count, head, HasMore: false));
    }

    /// <inheritdoc />
    public Task<VaultGrantsResponse> ListVaultGrantsAsync(
        Guid vaultId,
        CancellationToken cancellationToken) =>
        Task.FromResult(new VaultGrantsResponse(
            vaultId,
            KeyGeneration: 1,
            RekeyRequired: false,
            Grants:
            [
                .. grants.Where(entry => entry.Key.VaultId == vaultId).Select(entry =>
                    new VaultGrantSummary(
                        entry.Key.UserId,
                        directory.Find(candidate => candidate.UserId == entry.Key.UserId)?.Email,
                        null,
                        KeyGeneration: 1,
                        VaultGrantState.Active,
                        UserId,
                        DateTimeOffset.UnixEpoch,
                        null)),
            ]));

    /// <inheritdoc />
    public Task IssueVaultGrantAsync(
        Guid vaultId,
        IssueVaultGrantRequest request,
        CancellationToken cancellationToken)
    {
        grants[(vaultId, request.RecipientUserId)] = request;

        return Task.CompletedTask;
    }

    /// <inheritdoc />
    public Task<bool> RevokeVaultGrantAsync(
        Guid vaultId,
        Guid userId,
        CancellationToken cancellationToken) =>
        Task.FromResult(grants.Remove((vaultId, userId)));

    /// <summary>Publishes the enrolling account's own key, in the directory and the key log.</summary>
    private void RegisterSelf(KeyStatement statement, byte[] statementSignature)
    {
        if (directory.Exists(entry => entry.UserId == UserId))
        {
            return;
        }

        var sequence = AppendKeyLog(
            UserId, statement.EncryptionPublicKey, statement.SigningPublicKey, statementSignature);

        directory.Add(new DirectoryEntry(
            UserId,
            "alice@example.com",
            "Alice Example",
            statement.EncryptionPublicKey,
            statement.SigningPublicKey,
            DshCrypto.ComputeFingerprint(statement.EncryptionPublicKey, statement.SigningPublicKey),
            statement.KeyGeneration,
            sequence));
    }

    /// <summary>Appends a key log entry, chained as the real log chains it.</summary>
    private long AppendKeyLog(
        Guid userId,
        byte[] encryptionPublicKey,
        byte[] signingPublicKey,
        byte[] statementSignature)
    {
        var previous = keyLog.Count == 0
            ? KeyLogChain.CreateGenesisPreviousHash()
            : keyLog[^1].Hash;

        var createdAt = KeyLogChain.TruncateTimestamp(DateTimeOffset.UnixEpoch);
        var sequence = keyLog.Count + 1;

        var hash = KeyLogChain.ComputeEntryHash(
            previous, userId, 1, encryptionPublicKey, signingPublicKey, statementSignature, createdAt);

        keyLog.Add(new KeyLogRecord(
            sequence,
            userId,
            Generation: 1,
            encryptionPublicKey,
            signingPublicKey,
            statementSignature,
            previous,
            hash,
            createdAt));

        return sequence;
    }

    private void Recount(Guid teamId)
    {
        var index = teams.FindIndex(team => team.TeamId == teamId);

        if (index < 0)
        {
            return;
        }

        teams[index] = teams[index] with
        {
            MemberCount = members.GetValueOrDefault(teamId, []).Count,
            VaultCount = teamVaults.Values.Count(vault => vault.TeamId == teamId),
        };
    }
}