using System; using System.Collections.Generic; using System.Linq; using System.Net; using System.Threading.Tasks; using AutopilotMonitor.Functions.Security; using AutopilotMonitor.Functions.Tests.GraphResolution; using Microsoft.Extensions.Caching.Memory; using Microsoft.Extensions.Configuration; using Microsoft.Extensions.Logging.Abstractions; namespace AutopilotMonitor.Functions.Tests; /// /// Tests for the per-tenant app-token cache on . The security /// contract: a successful client-credentials token is reused until shortly before expiry, keyed /// strictly per tenant (no cross-tenant reuse), and failures are never cached. The /// records every outbound request so we can assert exactly how /// many token POSTs hit AAD. /// public class GraphTokenServiceTests { private const string TenantA = "11111111-1111-1111-1111-111111111111"; private const string TenantB = "22222222-2222-2222-2222-222222222222"; private static (GraphTokenService svc, StubHttpMessageHandler handler) Create(StubHttpMessageHandler handler) { var config = new ConfigurationBuilder() .AddInMemoryCollection(new Dictionary { ["EntraId:ClientId"] = "test-client", ["EntraId:ClientSecret"] = "test-secret", }) .Build(); var svc = NewService(config, new StubHttpClientFactory(handler)); return (svc, handler); } /// /// Builds the service with the dual app-reg dependencies. No legacy app is configured in /// these tests ⇒ the registry always resolves the primary credentials and the tenant-config /// lookup is never consulted for the decision — a null-returning repo is fine. /// private static GraphTokenService NewService(IConfiguration config, System.Net.Http.IHttpClientFactory factory) => new GraphTokenService( NullLogger.Instance, factory, new MemoryCache(new MemoryCacheOptions()), config, new EntraAppRegistry(config, NullLogger.Instance), new AutopilotMonitor.Functions.Services.TenantConfigurationService( Moq.Mock.Of(), NullLogger.Instance, new MemoryCache(new MemoryCacheOptions()))); private static int TokenPosts(StubHttpMessageHandler handler, string tenantId) => handler.Requests.Count(u => u.Contains($"/{tenantId}/oauth2/v2.0/token")); [Fact] public async Task GetAccessTokenAsync_caches_token_across_calls() { var handler = new StubHttpMessageHandler() .When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\",\"expires_in\":3599}"); var (svc, _) = Create(handler); var first = await svc.GetAccessTokenAsync(TenantA); var second = await svc.GetAccessTokenAsync(TenantA); Assert.Equal("tok-A", first.AccessToken); Assert.Equal("tok-A", second.AccessToken); // Second call is served from cache — only one POST ever reached AAD. Assert.Equal(1, TokenPosts(handler, TenantA)); } [Fact] public async Task GetAccessTokenAsync_isolates_tokens_per_tenant() { // Cross-tenant isolation: tenant A's token must never be served for tenant B, and each // tenant caches independently. var handler = new StubHttpMessageHandler() .When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\",\"expires_in\":3599}") .When($"/{TenantB}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-B\",\"expires_in\":3599}"); var (svc, _) = Create(handler); var a1 = await svc.GetAccessTokenAsync(TenantA); var a2 = await svc.GetAccessTokenAsync(TenantA); var b1 = await svc.GetAccessTokenAsync(TenantB); Assert.Equal("tok-A", a1.AccessToken); Assert.Equal("tok-A", a2.AccessToken); Assert.Equal("tok-B", b1.AccessToken); Assert.Equal(1, TokenPosts(handler, TenantA)); Assert.Equal(1, TokenPosts(handler, TenantB)); } [Fact] public async Task GetAccessTokenAsync_does_not_cache_failure() { // A non-retryable 400 (not a consent-propagation error) returns immediately and must NOT be // cached — the next request re-acquires so a transient outage doesn't pin a failure. var handler = new StubHttpMessageHandler() .When($"/{TenantA}/oauth2", HttpStatusCode.BadRequest, "{\"error\":\"invalid_request\"}"); var (svc, _) = Create(handler); var first = await svc.GetAccessTokenAsync(TenantA); var second = await svc.GetAccessTokenAsync(TenantA); Assert.Null(first.AccessToken); Assert.Null(second.AccessToken); Assert.Equal(2, TokenPosts(handler, TenantA)); } [Fact] public async Task GetAccessTokenAsync_caches_token_without_expires_in() { // AAD always returns expires_in, but the fallback lifetime path must still cache so a quirky // response doesn't silently disable caching. var handler = new StubHttpMessageHandler() .When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\"}"); var (svc, _) = Create(handler); await svc.GetAccessTokenAsync(TenantA); await svc.GetAccessTokenAsync(TenantA); Assert.Equal(1, TokenPosts(handler, TenantA)); } [Fact] public async Task InvalidateTenant_forces_fresh_token_on_next_call() { // The fresh-read contract: after an admin grants consent/roles, the cached token still // carries the OLD roles claim. InvalidateTenant (called by GraphFeatureDetector.InvalidateTenant // on the admin Refresh / access-check path) must drop it so the next acquire mints fresh. var handler = new StubHttpMessageHandler() .When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok-A\",\"expires_in\":3599}"); var (svc, _) = Create(handler); await svc.GetAccessTokenAsync(TenantA); svc.InvalidateTenant(TenantA); await svc.GetAccessTokenAsync(TenantA); Assert.Equal(2, TokenPosts(handler, TenantA)); } [Fact] public async Task InvalidateTenant_during_in_flight_acquire_is_not_overwritten_by_stale_token() { // Write-after-invalidate race: a token POST that STARTED before InvalidateTenant must not // repopulate the cache with its now-stale token. The first POST is parked in flight; the // admin grants + Refreshes (InvalidateTenant) while it is parked; on release it writes a // (stale-generation-tagged) token. A subsequent read MUST reject it and re-mint fresh. var handler = new InflightRaceHandler( firstBody: "{\"access_token\":\"stale\",\"expires_in\":3599}", laterBody: "{\"access_token\":\"fresh\",\"expires_in\":3599}"); var config = new ConfigurationBuilder() .AddInMemoryCollection(new Dictionary { ["EntraId:ClientId"] = "test-client", ["EntraId:ClientSecret"] = "test-secret", }) .Build(); var svc = NewService(config, new SingleHandlerFactory(handler)); // R1 starts; its POST parks in flight (holding the gate, before writing to the cache). var inflight = svc.GetAccessTokenAsync(TenantA); await WaitUntilAsync(() => handler.Entered >= 1, TimeSpan.FromSeconds(2)); // Admin grants + Refresh WHILE R1's POST is in flight. svc.InvalidateTenant(TenantA); // R1 completes → writes "stale", tagged with the pre-invalidation generation. handler.Release(); var r1 = await inflight; Assert.Equal("stale", r1.AccessToken); // R1's own caller still gets the token it acquired // Fresh read must NOT serve the stale entry — the generation mismatch makes it a miss → re-mint. var r2 = await svc.GetAccessTokenAsync(TenantA); Assert.Equal("fresh", r2.AccessToken); Assert.Equal(2, handler.Entered); // a real second POST happened (no stale cache hit) } [Fact] public async Task GetAccessTokenAsync_coalesces_concurrent_misses() { // Stampede guard, proven properly: a barrier-gated handler parks the first request on the // wire so the other nine genuinely pile up at the miss BEFORE any response/cache-fill. With // the per-tenant semaphore only ONE request ever reaches the wire; without it all ten would. // (A synchronous stub could pass even without the gate, because the first call may fill the // cache before the rest start — hence this stronger construction.) var handler = new GatedHttpMessageHandler("{\"access_token\":\"tok-A\",\"expires_in\":3599}"); var config = new ConfigurationBuilder() .AddInMemoryCollection(new Dictionary { ["EntraId:ClientId"] = "test-client", ["EntraId:ClientSecret"] = "test-secret", }) .Build(); var svc = NewService(config, new SingleHandlerFactory(handler)); var tasks = Enumerable.Range(0, 10).Select(_ => svc.GetAccessTokenAsync(TenantA)).ToArray(); // Wait for the first request to reach the wire, then give the other nine ample time to block // on the semaphore. If coalescing works, exactly one request is at the handler. await WaitUntilAsync(() => handler.Entered >= 1, TimeSpan.FromSeconds(2)); await Task.Delay(300); Assert.Equal(1, handler.Entered); handler.Release(); var results = await Task.WhenAll(tasks); Assert.All(results, r => Assert.Equal("tok-A", r.AccessToken)); Assert.Equal(1, handler.Entered); } [Fact] public async Task GetAccessTokenForAppAsync_caches_under_the_supplied_app_without_cross_app_poisoning() { // Dual app-reg window: the homing-flip consent probe mints under EXPLICIT primary // credentials for a still-legacy-homed tenant. Contract: the probe token caches under the // primary key, a second probe is served from cache, and the tenant's normal (legacy- // resolved) path still mints its own token — never the probe's. var handler = new StubHttpMessageHandler() .When($"/{TenantA}/oauth2", HttpStatusCode.OK, "{\"access_token\":\"tok\",\"expires_in\":3599}"); var config = new ConfigurationBuilder() .AddInMemoryCollection(new Dictionary { ["EntraId:ClientId"] = "aaaaaaaa-0000-0000-0000-000000000001", ["EntraId:ClientSecret"] = "primary-secret", ["EntraId:LegacyClientId"] = "bbbbbbbb-0000-0000-0000-000000000002", ["EntraId:LegacyClientSecret"] = "legacy-secret", }) .Build(); var registry = new EntraAppRegistry(config, NullLogger.Instance); var svc = new GraphTokenService( NullLogger.Instance, new StubHttpClientFactory(handler), new MemoryCache(new MemoryCacheOptions()), config, registry, new AutopilotMonitor.Functions.Services.TenantConfigurationService( Moq.Mock.Of(), NullLogger.Instance, new MemoryCache(new MemoryCacheOptions()))); var probe1 = await svc.GetAccessTokenForAppAsync(TenantA, registry.Primary); var probe2 = await svc.GetAccessTokenForAppAsync(TenantA, registry.Primary); Assert.Equal("tok", probe1.AccessToken); Assert.Equal("tok", probe2.AccessToken); Assert.Equal(1, TokenPosts(handler, TenantA)); // second probe served from the primary-keyed cache // Normal path resolves LEGACY for this tenant (no homing stored) — must be a cache miss // against the probe's primary-keyed entry and mint its own token. var homed = await svc.GetAccessTokenAsync(TenantA); Assert.Equal("tok", homed.AccessToken); Assert.Equal(2, TokenPosts(handler, TenantA)); // InvalidateTenant drops BOTH apps' entries — the next probe re-mints fresh. svc.InvalidateTenant(TenantA); await svc.GetAccessTokenForAppAsync(TenantA, registry.Primary); Assert.Equal(3, TokenPosts(handler, TenantA)); } private static async Task WaitUntilAsync(Func condition, TimeSpan timeout) { var deadline = DateTime.UtcNow + timeout; while (!condition() && DateTime.UtcNow < deadline) { await Task.Delay(10); } } }