jobtrackingapp/JobTrackerApi.Tests/TwoFactorControllerTests.cs

using JobTrackerApi.Controllers;
using JobTrackerApi.Models;
using JobTrackerApi.Services;
using JobTrackerApi.Tests.TestSupport;
using Microsoft.AspNetCore.DataProtection;
using Microsoft.AspNetCore.Http;
using Microsoft.AspNetCore.Identity;
using Microsoft.AspNetCore.Mvc;
using Microsoft.EntityFrameworkCore;
using Moq;
using OtpNet;
using Xunit;

namespace JobTrackerApi.Tests;

public sealed class TwoFactorControllerTests
{
    private static TwoFactorController BuildController(Mock<UserManager<ApplicationUser>> userManager, JobTrackerApi.Data.JobTrackerContext db, ITwoFactorPendingTokenService? pending = null, ApplicationUser? currentUser = null)
    {
        if (currentUser is not null)
        {
            userManager.Setup(x => x.GetUserAsync(It.IsAny<System.Security.Claims.ClaimsPrincipal>())).ReturnsAsync(currentUser);
        }

        var tokenService = new Mock<ITokenService>();
        tokenService.Setup(x => x.CreateAccessTokenAsync(It.IsAny<ApplicationUser>(), It.IsAny<CancellationToken>())).ReturnsAsync("app-token");

        var controller = new TwoFactorController(
            userManager.Object,
            tokenService.Object,
            db,
            pending ?? new TwoFactorPendingTokenService(new Microsoft.Extensions.Caching.Memory.MemoryCache(new Microsoft.Extensions.Caching.Memory.MemoryCacheOptions())),
            new EphemeralDataProtectionProvider())
        {
            ControllerContext = new ControllerContext { HttpContext = new DefaultHttpContext() }
        };
        return controller;
    }

    [Fact]
    public async Task Setup_then_verify_with_correct_code_enables_2fa_and_returns_recovery_codes()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "correct-password")).ReturnsAsync(true);
        userManager.Setup(x => x.UpdateAsync(It.IsAny<ApplicationUser>())).ReturnsAsync(IdentityResult.Success);

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, currentUser: user);

        var setupResult = Assert.IsType<OkObjectResult>(await controller.Setup(new TwoFactorController.PasswordConfirmRequest("correct-password"), CancellationToken.None));
        var setup = Assert.IsType<TwoFactorController.SetupResult>(setupResult.Value);
        Assert.NotNull(user.TotpPendingSecretEncrypted);
        Assert.False(user.TwoFactorEnabled);

        var code = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();

        var verifyResult = Assert.IsType<OkObjectResult>(await controller.VerifySetup(new TwoFactorController.VerifySetupRequest(code), CancellationToken.None));
        var verify = Assert.IsType<TwoFactorController.VerifySetupResult>(verifyResult.Value);

        Assert.True(verify.Enabled);
        Assert.Equal(10, verify.RecoveryCodes.Count);
        Assert.True(user.TwoFactorEnabled);
        Assert.Null(user.TotpPendingSecretEncrypted);
        Assert.NotNull(user.TotpSecretEncrypted);
        Assert.NotNull(user.TotpEnabledAtUtc);
    }

    [Fact]
    public async Task Verify_setup_with_wrong_code_is_rejected_and_does_not_enable_2fa()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "correct-password")).ReturnsAsync(true);
        userManager.Setup(x => x.UpdateAsync(It.IsAny<ApplicationUser>())).ReturnsAsync(IdentityResult.Success);

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, currentUser: user);

        await controller.Setup(new TwoFactorController.PasswordConfirmRequest("correct-password"), CancellationToken.None);

        var result = await controller.VerifySetup(new TwoFactorController.VerifySetupRequest("000000"), CancellationToken.None);

        Assert.IsType<UnauthorizedResult>(result);
        Assert.False(user.TwoFactorEnabled);
    }

    [Fact]
    public async Task Disable_with_wrong_password_is_rejected()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com", TwoFactorEnabled = true, TotpSecretEncrypted = "irrelevant" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "wrong-password")).ReturnsAsync(false);

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, currentUser: user);

        var result = await controller.Disable(new TwoFactorController.PasswordConfirmRequest("wrong-password"), CancellationToken.None);

        Assert.IsType<BadRequestObjectResult>(result);
        Assert.True(user.TwoFactorEnabled);
        Assert.NotNull(user.TotpSecretEncrypted);
    }

    [Fact]
    public async Task Challenge_with_valid_totp_code_completes_sign_in()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "correct-password")).ReturnsAsync(true);
        userManager.Setup(x => x.UpdateAsync(It.IsAny<ApplicationUser>())).ReturnsAsync(IdentityResult.Success);
        userManager.Setup(x => x.FindByIdAsync("user-1")).ReturnsAsync(user);

        var pending = new TwoFactorPendingTokenService(new Microsoft.Extensions.Caching.Memory.MemoryCache(new Microsoft.Extensions.Caching.Memory.MemoryCacheOptions()));

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, pending, currentUser: user);

        var setupResult = Assert.IsType<OkObjectResult>(await controller.Setup(new TwoFactorController.PasswordConfirmRequest("correct-password"), CancellationToken.None));
        var setup = Assert.IsType<TwoFactorController.SetupResult>(setupResult.Value);
        var code = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();
        await controller.VerifySetup(new TwoFactorController.VerifySetupRequest(code), CancellationToken.None);

        var pendingToken = pending.IssuePendingToken("user-1", rememberMe: false);
        var challengeCode = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();

        var challengeResult = Assert.IsType<OkObjectResult>(await controller.Challenge(new TwoFactorController.ChallengeRequest(pendingToken, challengeCode), CancellationToken.None));
        var session = Assert.IsType<AuthController.AuthSessionResult>(challengeResult.Value);
        Assert.True(session.Authenticated);

        // The pending token is single-use.
        var reuse = await controller.Challenge(new TwoFactorController.ChallengeRequest(pendingToken, challengeCode), CancellationToken.None);
        Assert.IsType<UnauthorizedResult>(reuse);
    }

    [Fact]
    public async Task Challenge_with_recovery_code_consumes_it_and_rejects_reuse()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "correct-password")).ReturnsAsync(true);
        userManager.Setup(x => x.UpdateAsync(It.IsAny<ApplicationUser>())).ReturnsAsync(IdentityResult.Success);
        userManager.Setup(x => x.FindByIdAsync("user-1")).ReturnsAsync(user);

        var pending = new TwoFactorPendingTokenService(new Microsoft.Extensions.Caching.Memory.MemoryCache(new Microsoft.Extensions.Caching.Memory.MemoryCacheOptions()));

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, pending, currentUser: user);

        var setupResult = Assert.IsType<OkObjectResult>(await controller.Setup(new TwoFactorController.PasswordConfirmRequest("correct-password"), CancellationToken.None));
        var setup = Assert.IsType<TwoFactorController.SetupResult>(setupResult.Value);
        var code = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();
        var verifyResult = Assert.IsType<OkObjectResult>(await controller.VerifySetup(new TwoFactorController.VerifySetupRequest(code), CancellationToken.None));
        var verify = Assert.IsType<TwoFactorController.VerifySetupResult>(verifyResult.Value);
        var recoveryCode = verify.RecoveryCodes[0];

        var pendingToken1 = pending.IssuePendingToken("user-1", rememberMe: false);
        var challengeResult = Assert.IsType<OkObjectResult>(await controller.Challenge(new TwoFactorController.ChallengeRequest(pendingToken1, recoveryCode), CancellationToken.None));
        Assert.True(Assert.IsType<AuthController.AuthSessionResult>(challengeResult.Value).Authenticated);

        // Same recovery code can't be used a second time, even against a fresh pending token.
        var pendingToken2 = pending.IssuePendingToken("user-1", rememberMe: false);
        var reuse = await controller.Challenge(new TwoFactorController.ChallengeRequest(pendingToken2, recoveryCode), CancellationToken.None);
        Assert.IsType<UnauthorizedResult>(reuse);
    }

    [Fact]
    public async Task Challenge_with_trust_device_true_creates_trusted_device_and_sets_cookie()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "correct-password")).ReturnsAsync(true);
        userManager.Setup(x => x.UpdateAsync(It.IsAny<ApplicationUser>())).ReturnsAsync(IdentityResult.Success);
        userManager.Setup(x => x.FindByIdAsync("user-1")).ReturnsAsync(user);

        var pending = new TwoFactorPendingTokenService(new Microsoft.Extensions.Caching.Memory.MemoryCache(new Microsoft.Extensions.Caching.Memory.MemoryCacheOptions()));

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, pending, currentUser: user);

        var setupResult = Assert.IsType<OkObjectResult>(await controller.Setup(new TwoFactorController.PasswordConfirmRequest("correct-password"), CancellationToken.None));
        var setup = Assert.IsType<TwoFactorController.SetupResult>(setupResult.Value);
        var code = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();
        await controller.VerifySetup(new TwoFactorController.VerifySetupRequest(code), CancellationToken.None);

        var pendingToken = pending.IssuePendingToken("user-1", rememberMe: false);
        var challengeCode = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();

        var challengeResult = Assert.IsType<OkObjectResult>(await controller.Challenge(new TwoFactorController.ChallengeRequest(pendingToken, challengeCode, TrustDevice: true), CancellationToken.None));
        Assert.True(Assert.IsType<AuthController.AuthSessionResult>(challengeResult.Value).Authenticated);

        var device = Assert.Single(db.TrustedDevices.IgnoreQueryFilters().Where(x => x.UserId == "user-1"));
        Assert.NotNull(device.TokenHash);

        var setCookieHeaders = controller.Response.Headers["Set-Cookie"];
        var trustedDeviceCookie = Assert.Single(setCookieHeaders, h => h!.StartsWith($"{AuthSessionOptions.TrustedDeviceCookieName}=", StringComparison.Ordinal))!;
        Assert.Contains("httponly", trustedDeviceCookie, StringComparison.OrdinalIgnoreCase);
        Assert.Contains("samesite=strict", trustedDeviceCookie, StringComparison.OrdinalIgnoreCase);
    }

    [Fact]
    public async Task Challenge_without_trust_device_does_not_create_trusted_device()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();
        userManager.Setup(x => x.CheckPasswordAsync(user, "correct-password")).ReturnsAsync(true);
        userManager.Setup(x => x.UpdateAsync(It.IsAny<ApplicationUser>())).ReturnsAsync(IdentityResult.Success);
        userManager.Setup(x => x.FindByIdAsync("user-1")).ReturnsAsync(user);

        var pending = new TwoFactorPendingTokenService(new Microsoft.Extensions.Caching.Memory.MemoryCache(new Microsoft.Extensions.Caching.Memory.MemoryCacheOptions()));

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, pending, currentUser: user);

        var setupResult = Assert.IsType<OkObjectResult>(await controller.Setup(new TwoFactorController.PasswordConfirmRequest("correct-password"), CancellationToken.None));
        var setup = Assert.IsType<TwoFactorController.SetupResult>(setupResult.Value);
        var code = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();
        await controller.VerifySetup(new TwoFactorController.VerifySetupRequest(code), CancellationToken.None);

        var pendingToken = pending.IssuePendingToken("user-1", rememberMe: false);
        var challengeCode = new Totp(Base32Encoding.ToBytes(setup.ManualEntryKey)).ComputeTotp();

        await controller.Challenge(new TwoFactorController.ChallengeRequest(pendingToken, challengeCode), CancellationToken.None);

        Assert.Empty(db.TrustedDevices.IgnoreQueryFilters().Where(x => x.UserId == "user-1"));
    }

    [Fact]
    public async Task Trusted_devices_can_be_listed_revoked_and_revoked_in_bulk_scoped_to_owner()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var otherUser = new ApplicationUser { Id = "user-2", Email = "other@example.com", UserName = "other@example.com" };
        var userManager = TestHostFactory.CreateUserManager();

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        db.TrustedDevices.Add(new TrustedDevice { UserId = "user-1", TokenHash = "hash-1", DeviceLabel = "Chrome on Windows", CreatedAtUtc = DateTimeOffset.UtcNow, LastSeenAtUtc = DateTimeOffset.UtcNow, ExpiresAtUtc = DateTimeOffset.UtcNow.AddDays(30) });
        db.TrustedDevices.Add(new TrustedDevice { UserId = "user-1", TokenHash = "hash-2", DeviceLabel = "Safari on Mac", CreatedAtUtc = DateTimeOffset.UtcNow, LastSeenAtUtc = DateTimeOffset.UtcNow, ExpiresAtUtc = DateTimeOffset.UtcNow.AddDays(30) });
        db.TrustedDevices.Add(new TrustedDevice { UserId = "user-2", TokenHash = "hash-3", DeviceLabel = "Someone else's device", CreatedAtUtc = DateTimeOffset.UtcNow, LastSeenAtUtc = DateTimeOffset.UtcNow, ExpiresAtUtc = DateTimeOffset.UtcNow.AddDays(30) });
        db.SaveChanges();

        var otherDeviceId = db.TrustedDevices.IgnoreQueryFilters().Single(x => x.UserId == "user-2").Id;

        var controller = BuildController(userManager, db, currentUser: user);

        var listResult = Assert.IsType<OkObjectResult>(await controller.ListTrustedDevices(CancellationToken.None));
        var list = Assert.IsType<List<TwoFactorController.TrustedDeviceDto>>(listResult.Value);
        Assert.Equal(2, list.Count);

        // Can't revoke another user's device, even by guessing its id.
        var revokeOther = await controller.RevokeTrustedDevice(otherDeviceId, CancellationToken.None);
        Assert.IsType<NotFoundResult>(revokeOther);
        Assert.NotNull(db.TrustedDevices.IgnoreQueryFilters().SingleOrDefault(x => x.Id == otherDeviceId));

        var ownDeviceId = list[0].Id;
        var revokeOwn = await controller.RevokeTrustedDevice(ownDeviceId, CancellationToken.None);
        Assert.IsType<NoContentResult>(revokeOwn);
        Assert.Null(db.TrustedDevices.IgnoreQueryFilters().SingleOrDefault(x => x.Id == ownDeviceId));

        var revokeAll = await controller.RevokeAllTrustedDevices(CancellationToken.None);
        Assert.IsType<NoContentResult>(revokeAll);
        Assert.Empty(db.TrustedDevices.IgnoreQueryFilters().Where(x => x.UserId == "user-1"));
        Assert.NotNull(db.TrustedDevices.IgnoreQueryFilters().SingleOrDefault(x => x.UserId == "user-2"));
    }

    [Fact]
    public async Task Challenge_with_unknown_pending_token_is_rejected()
    {
        var user = new ApplicationUser { Id = "user-1", Email = "person@example.com", UserName = "person@example.com" };
        var userManager = TestHostFactory.CreateUserManager();

        using var db = TestHostFactory.CreateInMemoryDb("user-1");
        var controller = BuildController(userManager, db, currentUser: user);

        var result = await controller.Challenge(new TwoFactorController.ChallengeRequest("not-a-real-token", "123456"), CancellationToken.None);

        Assert.IsType<UnauthorizedResult>(result);
    }
}