DodoSSH/src/DodoSSH.Api/Features/Identity/EnrollmentService.cs

using System.Security.Cryptography;
using System.Text.Json;
using DodoSSH.Api.Setup;
using DodoSSH.Contracts;
using DodoSSH.Crypto;
using DodoSSH.Domain;
using DodoSSH.Infrastructure;
using Microsoft.EntityFrameworkCore;
using Npgsql;

namespace DodoSSH.Api.Features.Identity;

/// <summary>
/// Publishes a user's first identity key and creates their personal vault.
/// </summary>
/// <remarks>
/// <para>
/// One transaction produces the key, its wraps, the device, the key log entry, the vault and the
/// vault key grant. Nothing here is separable: a key with no vault leaves a user unable to store
/// anything, and a vault with no grant is a container nobody can ever open — including its owner,
/// since only the client can wrap the key and it has already moved on.
/// </para>
/// <para>
/// Enrollment is idempotent on retry. A client whose request timed out after the server committed
/// re-sends the identical body and receives the identical response, because the vault id and the
/// keys are all client-chosen. Without that, a lost response would leave the user permanently
/// enrolled against a vault they never learned the id of.
/// </para>
/// </remarks>
internal sealed class EnrollmentService(
    DodoDbContext database,
    IIdentityBindingVerifier bindingVerifier,
    TimeProvider clock,
    ILogger<EnrollmentService> logger)
{
    /// <summary>
    /// Lock key serialising key log appends across the deployment.
    /// </summary>
    /// <remarks>
    /// A constant string rather than a per-user key, because the chain is global: two concurrent
    /// enrollments reading the same head would produce two entries claiming the same predecessor,
    /// which is indistinguishable from the fork the chain exists to detect. Enrollment happens once
    /// per user, so serialising it costs nothing.
    /// </remarks>
    private const string KeyLogLockName = "dodossh:key_log";

    /// <summary>Enrolls the caller's first identity key.</summary>
    /// <exception cref="EnrollmentInvalidException">The request is structurally unacceptable.</exception>
    /// <exception cref="IdentityBindingInvalidException">The provider token did not bind the keys.</exception>
    /// <exception cref="AlreadyEnrolledException">A different current key already exists.</exception>
    internal async Task<EnrollmentResponse> EnrollAsync(
        UserAccount user,
        EnrollmentRequest request,
        CancellationToken cancellationToken)
    {
        EnrollmentValidation.Validate(request, user);

        var statement = request.Statement;
        var canonical = KeyStatementCodec.Encode(ToFields(statement));

        // Order matters only for cost: the self-signature is local and cheap, the binding check may
        // hit the provider's JWKS endpoint. Both are mandatory and neither substitutes for the
        // other — the self-signature proves possession of the private key, the binding proves whose
        // key it is.
        if (!DshSignatures.VerifyKeyStatement(
                statement.SigningPublicKey, canonical, request.StatementSignature))
        {
            IdentityLog.StatementSignatureRejected(logger, user.Id);
            throw new EnrollmentInvalidException(
                "The key statement self-signature did not verify against the statement's own signing key.");
        }

        var verification = await bindingVerifier.VerifyAsync(
                request.IdentityProviderToken,
                statement.Issuer,
                statement.Subject,
                KeyStatementCodec.ToNonce(KeyStatementCodec.ComputeBinding(canonical)),
                cancellationToken)
            .ConfigureAwait(false);

        if (!verification.Verified)
        {
            IdentityLog.BindingRejected(logger, user.Id);
            throw new IdentityBindingInvalidException(verification.Failure!);
        }

        var fingerprint = DshCrypto.ComputeFingerprint(
            statement.EncryptionPublicKey,
            statement.SigningPublicKey);

        // Cheap path for the common retry. The authoritative check runs inside the transaction.
        var current = await FindCurrentKeyAsync(user.Id, cancellationToken).ConfigureAwait(false);
        if (current is not null)
        {
            return await ResolveExistingAsync(user, request, current, fingerprint, cancellationToken)
                .ConfigureAwait(false);
        }

        return await CommitAsync(user, request, verification, fingerprint, cancellationToken)
            .ConfigureAwait(false);
    }

    private async Task<EnrollmentResponse> CommitAsync(
        UserAccount user,
        EnrollmentRequest request,
        BindingVerification verification,
        byte[] fingerprint,
        CancellationToken cancellationToken)
    {
        return await database.ExecuteFreshAsync([user], async () =>
        {
            var transaction = await database.Database
                .BeginTransactionAsync(cancellationToken)
                .ConfigureAwait(false);
            await using var _ = transaction.ConfigureAwait(false);

            // First statement in the transaction, as with a sync push. The key log's sequence is an
            // identity column, so it is assigned before commit: without serialising appends, two
            // transactions can interleave their sequences and the hash chain no longer matches the
            // order a reader sees. See ADR 0003 for the same hazard in the vault change log.
            await AcquireKeyLogLockAsync(cancellationToken).ConfigureAwait(false);

            var current = await FindCurrentKeyAsync(user.Id, cancellationToken).ConfigureAwait(false);
            if (current is not null)
            {
                // Two concurrent first enrollments. Fall back to the same resolution the fast path
                // uses, so a duplicate submission still ends in an idempotent success.
                return await ResolveExistingAsync(user, request, current, fingerprint, cancellationToken)
                    .ConfigureAwait(false);
            }

            await RequireVaultIdIsFreeAsync(request.PersonalVault.VaultId, cancellationToken)
                .ConfigureAwait(false);

            var written = await WriteAsync(user, request, verification, fingerprint, cancellationToken)
                .ConfigureAwait(false);

            await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);

            IdentityLog.Enrolled(logger, user.Id, written.Generation, written.KeyLogSequence);

            return new EnrollmentResponse(
                UserId: user.Id,
                KeyGeneration: written.Generation,
                Fingerprint: fingerprint,
                PersonalVaultId: request.PersonalVault.VaultId,
                DeviceId: written.DeviceId,
                KeyLogSequence: written.KeyLogSequence);
        }).ConfigureAwait(false);
    }

    /// <summary>What the write produced, for the response.</summary>
    [System.Runtime.InteropServices.StructLayout(System.Runtime.InteropServices.LayoutKind.Auto)]
    private readonly record struct WrittenEnrollment(int Generation, Guid? DeviceId, long KeyLogSequence);

    /// <summary>Writes every row enrollment produces, inside the caller's transaction.</summary>
    private async Task<WrittenEnrollment> WriteAsync(
        UserAccount user,
        EnrollmentRequest request,
        BindingVerification verification,
        byte[] fingerprint,
        CancellationToken cancellationToken)
    {
        // Truncated to the precision the key log hashes, so the stored row can reproduce its own
        // hash after a round trip through PostgreSQL's microsecond timestamps.
        var now = KeyLogChain.TruncateTimestamp(clock.GetUtcNow());

        var key = BuildKey(user, request, verification, fingerprint, now);
        database.UserKeys.Add(key);

        AddWraps(user, request, now);
        var deviceId = AddDevice(user, request, now);
        var entry = await AppendKeyLogAsync(user, request, now, cancellationToken).ConfigureAwait(false);
        AddPersonalVault(user, request, fingerprint, now);

        user.EnrolledAtUtc = now;
        user.UpdatedAtUtc = now;
        user.LastSeenAtUtc = now;

        try
        {
            await database.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
        }
        catch (DbUpdateException exception) when (IsUniqueViolation(exception))
        {
            // The advisory lock covers the key log, and the pre-checks cover the expected races.
            // Anything left is a genuine collision — most plausibly the global unique index on key
            // fingerprints — and must not be reported as success.
            throw new AlreadyEnrolledException(
                "Enrollment collided with existing identity data. Retry with freshly generated keys.");
        }

        return new WrittenEnrollment(key.Generation, deviceId, entry.Sequence);
    }

    /// <summary>
    /// Decides whether an enrollment against an already-enrolled user is a retry or a conflict.
    /// </summary>
    /// <remarks>
    /// Compared by fingerprint, which covers both public keys, rather than by any single field. The
    /// vault id must match too: the same keys with a different vault id is a client that has lost
    /// track of its own state, and creating a second vault for it would strand the first.
    /// </remarks>
    private async Task<EnrollmentResponse> ResolveExistingAsync(
        UserAccount user,
        EnrollmentRequest request,
        UserKey current,
        byte[] fingerprint,
        CancellationToken cancellationToken)
    {
        if (!CryptographicOperations.FixedTimeEquals(current.FingerprintSha256, fingerprint))
        {
            throw new AlreadyEnrolledException(
                "This account already holds a different identity key. Rotate the existing key "
                + "rather than enrolling a second one.");
        }

        var vaultExists = await database.Vaults
            .AnyAsync(
                v => v.Id == request.PersonalVault.VaultId
                    && v.OwnerUserId == user.Id
                    && v.DeletedAtUtc == null,
                cancellationToken)
            .ConfigureAwait(false);

        if (!vaultExists)
        {
            throw new AlreadyEnrolledException(
                "This account is already enrolled with these keys, but not against the supplied "
                + "vault id. Read /api/v1/me to recover the current state.");
        }

        var sequence = await database.KeyLog
            .Where(e => e.UserId == user.Id && e.Generation == current.Generation)
            .Select(e => e.Sequence)
            .FirstOrDefaultAsync(cancellationToken)
            .ConfigureAwait(false);

        var deviceId = request.DevicePublicKey is null
            ? (Guid?)null
            : await database.Devices
                .Where(d => d.UserId == user.Id && d.PublicKey == request.DevicePublicKey)
                .Select(d => (Guid?)d.Id)
                .FirstOrDefaultAsync(cancellationToken)
                .ConfigureAwait(false);

        IdentityLog.EnrollmentReplayed(logger, user.Id);

        return new EnrollmentResponse(
            UserId: user.Id,
            KeyGeneration: current.Generation,
            Fingerprint: current.FingerprintSha256,
            PersonalVaultId: request.PersonalVault.VaultId,
            DeviceId: deviceId,
            KeyLogSequence: sequence);
    }

    private static UserKey BuildKey(
        UserAccount user,
        EnrollmentRequest request,
        BindingVerification verification,
        byte[] fingerprint,
        DateTimeOffset now) =>
        new()
        {
            Id = Guid.CreateVersion7(),
            UserId = user.Id,
            Generation = request.Statement.KeyGeneration,
            EncryptionPublicKey = request.Statement.EncryptionPublicKey,
            SigningPublicKey = request.Statement.SigningPublicKey,
            FingerprintSha256 = fingerprint,

            // Stored verbatim so a client can recompute the canonical encoding and check the
            // binding for itself. The hash is over the canonical encoding, not over this JSON, so
            // the serialiser's formatting choices cannot invalidate anything.
            Statement = JsonSerializer.Serialize(
                request.Statement,
                DodoSshJsonContext.ResponseOptions),

            StatementSignature = request.StatementSignature,
            IdentityProviderBinding = verification.Evidence,
            IsCurrent = true,
            CreatedAtUtc = now,
        };

    private void AddWraps(UserAccount user, EnrollmentRequest request, DateTimeOffset now)
    {
        database.UserKeyWraps.Add(BuildKdfWrap(
            user.Id,
            UserKeyWrapKind.Passphrase,
            request.WrappedPrivateKey,
            request.KdfParameters,
            now));

        if (request.RecoveryWrappedPrivateKey is not null && request.RecoveryKdfParameters is not null)
        {
            database.UserKeyWraps.Add(BuildKdfWrap(
                user.Id,
                UserKeyWrapKind.Recovery,
                request.RecoveryWrappedPrivateKey,
                request.RecoveryKdfParameters,
                now));
        }
    }

    private static UserKeyWrap BuildKdfWrap(
        Guid userId,
        UserKeyWrapKind kind,
        byte[] wrap,
        KdfParameters parameters,
        DateTimeOffset now) =>
        new()
        {
            Id = Guid.CreateVersion7(),
            UserId = userId,
            Kind = kind,
            Wrap = wrap,
            WrapVersion = 1,
            KdfAlgorithm = parameters.Algorithm,
            KdfSalt = parameters.Salt,
            KdfMemoryKibibytes = parameters.MemoryKibibytes,
            KdfPasses = parameters.Passes,
            KdfParallelism = parameters.Parallelism,
            CreatedAtUtc = now,
        };

    /// <remarks>
    /// The platform is left unreported: the request carries a device name but no platform, and
    /// guessing one from a user agent would be a display-only field populated with a lie. The
    /// devices endpoint sets it properly when it lands.
    /// </remarks>
    private Guid? AddDevice(UserAccount user, EnrollmentRequest request, DateTimeOffset now)
    {
        if (request.DevicePublicKey is null || request.DeviceWrappedPrivateKey is null)
        {
            return null;
        }

        var device = new Device
        {
            Id = Guid.CreateVersion7(),
            UserId = user.Id,
            Name = request.Statement.DeviceName,
            Platform = DevicePlatform.Unspecified,
            PublicKey = request.DevicePublicKey,
            EnrolledAtUtc = now,
            LastSeenAtUtc = now,
        };

        database.Devices.Add(device);

        database.UserKeyWraps.Add(new UserKeyWrap
        {
            Id = Guid.CreateVersion7(),
            UserId = user.Id,
            Kind = UserKeyWrapKind.Device,
            DeviceId = device.Id,
            Wrap = request.DeviceWrappedPrivateKey,
            WrapVersion = 1,
            CreatedAtUtc = now,
        });

        return device.Id;
    }

    private async Task<KeyLogEntry> AppendKeyLogAsync(
        UserAccount user,
        EnrollmentRequest request,
        DateTimeOffset now,
        CancellationToken cancellationToken)
    {
        // Ordered by sequence rather than by timestamp: the sequence is what the chain follows, and
        // two entries can share a millisecond.
        var previousHash = await database.KeyLog
            .OrderByDescending(e => e.Sequence)
            .Select(e => e.Hash)
            .FirstOrDefaultAsync(cancellationToken)
            .ConfigureAwait(false)
            ?? KeyLogChain.CreateGenesisPreviousHash();

        var entry = new KeyLogEntry
        {
            UserId = user.Id,
            Generation = request.Statement.KeyGeneration,
            EncryptionPublicKey = request.Statement.EncryptionPublicKey,
            SigningPublicKey = request.Statement.SigningPublicKey,
            StatementSignature = request.StatementSignature,
            PreviousHash = previousHash,
            CreatedAtUtc = now,
        };

        entry.Hash = KeyLogChain.ComputeEntryHash(
            previousHash,
            entry.UserId,
            entry.Generation,
            entry.EncryptionPublicKey,
            entry.SigningPublicKey,
            entry.StatementSignature,
            entry.CreatedAtUtc);

        database.KeyLog.Add(entry);

        return entry;
    }

    /// <remarks>
    /// The grant carries no key log head. A self-grant has no third party whose key could have been
    /// substituted, and the entry that would supply the head is being written in this very
    /// transaction — so the client could not have signed over it. Recording the server's own
    /// observation instead would produce a row whose signature verifies against nothing.
    /// </remarks>
    private void AddPersonalVault(
        UserAccount user,
        EnrollmentRequest request,
        byte[] fingerprint,
        DateTimeOffset now)
    {
        var personal = request.PersonalVault;

        database.Vaults.Add(new Vault
        {
            Id = personal.VaultId,
            Name = personal.Name,
            OwnerKind = VaultOwnerKind.Personal,
            OwnerUserId = user.Id,
            KeyGeneration = 1,
            CreatedAtUtc = now,
            UpdatedAtUtc = now,
        });

        database.VaultKeyGrants.Add(new VaultKeyGrant
        {
            Id = Guid.CreateVersion7(),
            VaultId = personal.VaultId,
            KeyGeneration = 1,
            Kind = GrantKind.Member,
            RecipientUserId = user.Id,
            RecipientKeyFingerprint = fingerprint,
            WrappedKey = personal.WrappedVaultKey,
            GranterUserId = user.Id,
            GranterKeyFingerprint = fingerprint,
            KeyLogHead = null,
            Signature = personal.GrantSignature,
            State = GrantState.Active,
            CreatedAtUtc = now,
        });
    }

    /// <remarks>
    /// Checked rather than left to the primary key, so the failure is a clear 400 about the
    /// caller's own request instead of a constraint violation surfacing as a 500. The residual
    /// disclosure — that some vault with this id exists — is accepted: vault ids are unguessable
    /// random values, so confirming one exists tells a caller nothing they did not already supply.
    /// </remarks>
    private async Task RequireVaultIdIsFreeAsync(Guid vaultId, CancellationToken cancellationToken)
    {
        var taken = await database.Vaults
            .AnyAsync(v => v.Id == vaultId, cancellationToken)
            .ConfigureAwait(false);

        if (taken)
        {
            throw new EnrollmentInvalidException(
                "The supplied personal vault id is already in use. Generate a new UUIDv7 and retry.");
        }
    }

    private Task<UserKey?> FindCurrentKeyAsync(Guid userId, CancellationToken cancellationToken) =>
        database.UserKeys.SingleOrDefaultAsync(
            k => k.UserId == userId && k.IsCurrent,
            cancellationToken);

    private Task<int> AcquireKeyLogLockAsync(CancellationToken cancellationToken) =>
        database.Database.ExecuteSqlAsync(
            $"SELECT pg_advisory_xact_lock(hashtextextended({KeyLogLockName}, 0))",
            cancellationToken);

    private static KeyStatementFields ToFields(KeyStatement statement) =>
        new(
            statement.Version,
            statement.Issuer,
            statement.Subject,
            statement.Email,
            statement.EncryptionPublicKey,
            statement.SigningPublicKey,
            statement.KeyGeneration,
            statement.CreatedAt,
            statement.DeviceName);

    private static bool IsUniqueViolation(DbUpdateException exception) =>
        string.Equals(
            (exception.InnerException as PostgresException)?.SqlState,
            PostgresErrorCodes.UniqueViolation,
            StringComparison.Ordinal);
}