missions/tests/Azaion.Missions.E2E.Tests/Tests/Security/JwksRotationTests.cs

using System.Net;
using System.Net.Http.Headers;
using System.Net.Http.Json;
using System.Text.Json;
using Azaion.Missions.E2E.Fixtures;
using Azaion.Missions.E2E.Helpers;
using Xunit;

namespace Azaion.Missions.E2E.Tests.Security;

/// <summary>
/// NFT-SEC-11 — security-shaped view of JWKS rotation. Verifies the kid-cache
/// mechanics + grace-window timing; the resilience-shaped variant
/// (no-restart) lives in <c>Tests/Resilience/JwksRotationTests.cs</c>.
/// Traces: AC-5.7.
/// </summary>
/// <remarks>
/// Owns the <c>JwksRotation</c> xUnit collection because rotating the mock
/// changes the active kid for every subsequent test that holds a stale
/// token. After running, the next test class in any collection mints a
/// fresh token, so it picks up the new kid on its next JWKS refresh.
/// </remarks>
[Collection("JwksRotation")]
[Trait("Category", "Sec")]
[Trait("db_access", "seed-or-assert-only")]
public sealed class JwksRotationTests : TestBase, IClassFixture<DbResetFixture>
{
    [Fact(Timeout = 130_000)]
    [Trait("Traces", "AC-5.7")]
    [Trait("max_ms", "120000")]
    public async Task NFT_SEC_11_unknown_kid_rotation_completes_within_120s_honouring_grace()
    {
        // Arrange — warm up: confirm the active key works before rotation.
        DbResetFixture.ResetDatabase(TestEnvironment.DbSideChannel);
        Seeds.Apply(Seeds.OneDefaultVehicle.Sql);

        var t1 = await Tokens.MintDefaultAsync();
        var kidV1 = t1.Kid;
        using (var resp = await CallVehiclesAsync(t1.Jwt))
            await HttpAssertions.AssertStatusAsync(resp, HttpStatusCode.OK);

        var rotationStart = DateTime.UtcNow;

        // Act 1: Rotate the mock. After this call, kid_v2 is active and
        // kid_v1 is retained for OLD_KEY_GRACE_SECONDS=5.
        var kidV2 = await RotateMockAsync();
        Assert.NotEqual(kidV1, kidV2);

        // Mint T2 with the brand-new active key.
        var t2 = await Tokens.MintDefaultAsync();
        Assert.Equal(kidV2, t2.Kid);

        // Assert AC-5.7.1 — T2 is rejected BEFORE missions refreshes its JWKS
        // cache (the new kid is not yet in the cache). We probe immediately
        // and require at least one 401 — once missions refreshes, subsequent
        // calls should succeed.
        using (var resp = await CallVehiclesAsync(t2.Jwt))
            await HttpAssertions.AssertStatusAsync(resp, HttpStatusCode.Unauthorized);

        // Assert AC-5.7.3 — during the 5s grace window, the OLD-kid token T1
        // is still accepted (missions' cache still contains kid_v1 from the
        // initial bootstrap fetch; the cache hasn't refreshed yet).
        using (var resp = await CallVehiclesAsync(t1.Jwt))
            await HttpAssertions.AssertStatusAsync(resp, HttpStatusCode.OK);

        // Act 2: Wait for JWKS refresh — poll T2 every 3s, up to 90s.
        var refreshDeadline = DateTime.UtcNow.AddSeconds(90);
        var refreshed = false;
        while (DateTime.UtcNow < refreshDeadline)
        {
            using var resp = await CallVehiclesAsync(t2.Jwt);
            if (resp.StatusCode == HttpStatusCode.OK)
            {
                refreshed = true;
                break;
            }
            await Task.Delay(TimeSpan.FromSeconds(3));
        }
        Assert.True(refreshed,
            "JWKS refresh did not propagate to missions within 90s (max-age=60s + auto-refresh=30s)");

        // Assert AC-5.7.4 — after the 5s grace window, the mock refuses to
        // sign with the old kid. Wait until grace certainly expired.
        var graceExpiry = rotationStart.AddSeconds(7);
        var until = graceExpiry - DateTime.UtcNow;
        if (until > TimeSpan.Zero)
            await Task.Delay(until);

        using var http = new HttpClient { Timeout = TimeSpan.FromSeconds(10) };
        var signUrl = new Uri(TestEnvironment.JwksMockSignUrl);
        using var signResponse = await http.PostAsJsonAsync(
            signUrl,
            new { kid_override = kidV1, permissions = "FL" });
        Assert.Equal(HttpStatusCode.BadRequest, signResponse.StatusCode);
        var body = await signResponse.Content.ReadFromJsonAsync<JsonElement>();
        Assert.True(body.TryGetProperty("error", out _),
            "mock refusal must include 'error' field");
    }

    private async Task<HttpResponseMessage> CallVehiclesAsync(string jwt)
    {
        var req = new HttpRequestMessage(HttpMethod.Get, "/vehicles");
        req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", jwt);
        return await Missions.SendAsync(req);
    }

    private static async Task<string> RotateMockAsync()
    {
        using var http = new HttpClient { Timeout = TimeSpan.FromSeconds(10) };
        var rotateUrl = new Uri(new Uri(TestEnvironment.JwksMockBaseUrl), "/rotate-key");
        using var resp = await http.PostAsync(rotateUrl, content: null);
        resp.EnsureSuccessStatusCode();
        var body = await resp.Content.ReadFromJsonAsync<JsonElement>();
        return body.GetProperty("kid").GetString()
            ?? throw new InvalidOperationException("mock /rotate-key returned no kid");
    }
}