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

using DodoSSH.Client.Api;
using DodoSSH.Client.Auth;
using DodoSSH.Client.Session;
using DodoSSH.Client.Sync;
using DodoSSH.Contracts;

namespace DodoSSH.Client.App.Tests;

/// <summary>
/// A signed-in server, without the signing in.
/// </summary>
/// <remarks>
/// Stands in for a <c>ServerConnection</c> so the shell's state machine can be driven end to end. The
/// account half stores what it is given and reports it back, because the provisioner re-reads <c>/me</c>
/// after enrolling and a stub that echoed the request would make that check meaningless. The sync half
/// applies pushes and serves them back as a change log, which is enough for the shell — the interesting
/// conflict behaviour is covered in <c>DodoSSH.Client.Sync.Tests</c> against a server that enforces
/// version checks.
/// </remarks>
internal sealed partial class FakeVaultServer : IVaultServer, IAccountApi, ISyncApi, IKeyBindingAuthorizer
{
    private readonly List<SyncChange> log = [];

    /// <remarks>
    /// Keyed on the entity type as well as the id, as the server's tables and the client's cache both are.
    /// Ids are UUIDv7 so a collision between two types will not happen by accident — but a fake that would
    /// treat a host and a key with one id as one row is a fake that could make a real bug pass.
    /// </remarks>
    private readonly Dictionary<(SyncEntityType Type, Guid EntityId), SyncChange> rows = [];

    private KeyStatement? statement;
    private byte[]? wrappedPrivateKey;
    private KdfParameters? kdfParameters;
    private VaultSummary? personalVault;

    internal Guid UserId { get; } = Guid.Parse("0192f0c8-4444-7aaa-8bbb-dddddddddddd");

    internal int EnrollmentCount { get; private set; }

    internal int PushCount { get; private set; }

    internal bool IsEnrolled => statement is not null;

    internal int LiveRowCount => rows.Values.Count(row => row.Operation != SyncOperation.Delete);

    /// <summary>Device wraps registered after enrollment, keyed on the device public key.</summary>
    internal Dictionary<string, byte[]> RegisteredDevices { get; } = new(StringComparer.Ordinal);

    /// <summary>The id issued for each registered public key, so revocation has something to name.</summary>
    private readonly Dictionary<string, Guid> deviceIds = new(StringComparer.Ordinal);

    /// <summary>When set, the next sign-in throws — how an unreachable server is exercised.</summary>
    internal Exception? SignInFailure { get; set; }

    /// <summary>
    /// When set, every synchronisation throws.
    /// </summary>
    /// <remarks>
    /// A server that answers but fails, as distinct from no server at all. The two are handled quite
    /// differently by a background pass: one is expected and silent, the other has to not overwrite
    /// whatever the user was reading.
    /// </remarks>
    internal Exception? SyncFailure { get; set; }

    /// <inheritdoc />
    public Uri ServerUrl { get; } = new("https://dodossh.example");

    /// <inheritdoc />
    public IAccountApi Account => this;

    /// <inheritdoc />
    public ISyncApi Sync => this;

    /// <inheritdoc />
    public IKeyBindingAuthorizer KeyBinding => this;

    /// <inheritdoc />
    public SyncOptions SyncOptions => SyncOptions.Default;

    /// <summary>
    /// The refresh token this "connection" holds.
    /// </summary>
    /// <remarks>
    /// Settable, because rotation is the half of remembering a sign-in that is easy to get wrong: a shell
    /// that persisted the token it first saw would leave a rotating provider refusing the next launch. A
    /// test changes this and asserts the new value reaches the cache.
    /// </remarks>
    public string? RefreshToken { get; set; } = "refresh-token-1";

    /// <inheritdoc />
    public void Dispose()
    {
        // Nothing to release; the shell disposes this on lock and on shutdown, and both paths have to be
        // safe to run more than once.
    }

    // ---- Identity provider ----

    /// <inheritdoc />
    public Task<string> AuthorizeKeyBindingAsync(string bindingNonce, CancellationToken cancellationToken) =>
        Task.FromResult("stub-id-token");

    // ---- Account ----

    /// <inheritdoc />
    public Task<MeResponse> GetMeAsync(CancellationToken cancellationToken) =>
        Task.FromResult(new MeResponse(
            UserId,
            "https://idp.example/realms/dodossh",
            "alice",
            "alice@example.com",
            "Alice Example",
            EnrollmentRequired: !IsEnrolled,
            KeyGeneration: statement?.KeyGeneration,
            WrappedPrivateKey: wrappedPrivateKey,
            KdfParameters: kdfParameters,

            // Team vaults alongside the personal one, in the order the real /me returns them: this is
            // where a vault somebody shared arrives, and a fake that only ever reported the personal one
            // would make a refresh that admits a new vault untestable.
            Vaults: personalVault is null ? [] : [personalVault, .. teamVaults.Values]));

    /// <inheritdoc />
    public Task<EnrollmentResponse> EnrollAsync(
        EnrollmentRequest request,
        CancellationToken cancellationToken)
    {
        EnrollmentCount++;

        statement = request.Statement;
        wrappedPrivateKey = request.WrappedPrivateKey;
        kdfParameters = request.KdfParameters;

        // The enrolling account joins the directory and the key log, as it does on the real server. Both
        // are what a later share reads: this client verifies its own entry as part of verifying anyone's.
        RegisterSelf(request.Statement, request.StatementSignature);

        personalVault = new VaultSummary(
            request.PersonalVault.VaultId,
            request.PersonalVault.Name,
            IsPersonal: true,
            TeamId: null,
            KeyGeneration: 1,
            Permissions: 31,
            request.PersonalVault.WrappedVaultKey,
            RekeyRequired: false);

        return Task.FromResult(new EnrollmentResponse(
            UserId,
            KeyGeneration: 1,
            Fingerprint: new byte[32],
            request.PersonalVault.VaultId,
            DeviceId: null,
            KeyLogSequence: 1));
    }

    /// <inheritdoc />
    /// <remarks>
    /// Records the wrap so a test can assert it reached the server, and refuses before enrollment as the
    /// real endpoint's <c>Auth.EnrolledPolicy</c> does.
    /// </remarks>
    public Task<RegisterDeviceResponse> RegisterDeviceAsync(
        RegisterDeviceRequest request,
        CancellationToken cancellationToken)
    {
        if (!IsEnrolled)
        {
            throw new DodoSshApiException(
                System.Net.HttpStatusCode.Forbidden,
                ProblemCodes.EnrollmentRequired,
                "This account has no identity key yet.");
        }

        var key = Convert.ToHexString(request.PublicKey);

        RegisteredDevices[key] = request.WrappedPrivateKey;

        // One id per public key, as the real service issues, so a revocation can name the device that was
        // actually registered rather than one this fake invented on the way past.
        if (!deviceIds.TryGetValue(key, out var deviceId))
        {
            deviceId = Guid.CreateVersion7();
            deviceIds[key] = deviceId;
        }

        return Task.FromResult(new RegisterDeviceResponse(deviceId, DateTimeOffset.UnixEpoch));
    }

    /// <inheritdoc />
    public Task<bool> RevokeDeviceAsync(Guid deviceId, CancellationToken cancellationToken)
    {
        var key = deviceIds.FirstOrDefault(entry => entry.Value == deviceId).Key;

        if (key is null)
        {
            return Task.FromResult(false);
        }

        deviceIds.Remove(key);

        // With its wrap, as the foreign key's cascade does on the real server.
        RegisteredDevices.Remove(key);

        return Task.FromResult(true);
    }

    // ---- Sync ----

    /// <inheritdoc />
    public Task<SyncPullResponse> SyncPullAsync(
        Guid vaultId,
        SyncPullRequest request,
        CancellationToken cancellationToken)
    {
        if (SyncFailure is { } failure)
        {
            return Task.FromException<SyncPullResponse>(failure);
        }

        var after = request.Cursor is null
            ? 0
            : long.Parse(request.Cursor.AsSpan("app-v1:".Length), provider: null);

        var page = log.Where(change => change.ChangeSequence > after).ToList();
        var next = page.Count > 0 ? page[^1].ChangeSequence : after;

        return Task.FromResult(new SyncPullResponse(
            page,
            $"app-v1:{next}",
            HasMore: false,
            ServerTime: DateTimeOffset.FromUnixTimeSeconds(1_750_000_000),
            CurrentKeyGeneration: 1));
    }

    /// <inheritdoc />
    public Task<SyncPushResponse> SyncPushAsync(
        Guid vaultId,
        SyncPushRequest request,
        CancellationToken cancellationToken)
    {
        PushCount++;

        var results = new List<SyncPushResult>(request.Operations.Count);

        foreach (var operation in request.Operations)
        {
            results.Add(Apply(operation));
        }

        return Task.FromResult(new SyncPushResponse(results, $"app-v1:{log.Count}"));
    }

    private SyncPushResult Apply(SyncPushOperation operation)
    {
        rows.TryGetValue((operation.EntityType, operation.EntityId), out var existing);

        var current = existing?.Operation == SyncOperation.Delete ? null : existing;

        if (operation.ExpectedVersion != current?.Version)
        {
            return new SyncPushResult(
                operation.OperationId,
                SyncOperationStatus.Conflict,
                current?.Version,
                current?.ChangeSequence,
                current,
                null);
        }

        var sequence = log.Count + 1;

        var change = new SyncChange(
            operation.EntityType,
            operation.EntityId,
            operation.Operation,
            Version: (current?.Version ?? 0) + 1,
            ChangeSequence: sequence,
            Payload: operation.Operation == SyncOperation.Delete ? null : operation.Payload,
            PlaintextFields: operation.Operation == SyncOperation.Delete
                ? null
                : operation.PlaintextFields,
            UpdatedAt: DateTimeOffset.FromUnixTimeSeconds(1_750_000_000 + sequence));

        rows[(operation.EntityType, operation.EntityId)] = change;
        log.Add(change);

        return new SyncPushResult(
            operation.OperationId,
            SyncOperationStatus.Applied,
            change.Version,
            sequence,
            null,
            null);
    }
}