diff --git a/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/GraphFeatureDetectorTests.cs b/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/GraphFeatureDetectorTests.cs index 0c52ff10..2a25bbee 100644 --- a/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/GraphFeatureDetectorTests.cs +++ b/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/GraphFeatureDetectorTests.cs @@ -313,6 +313,145 @@ public class GraphFeatureDetectorTests Assert.Equal(2, tokenService.CallCount); } + [Fact] + public async Task InvalidateTenant_drops_underlying_token_cache_so_post_grant_refresh_sees_new_roles() + { + // Regression for the exact consent-propagation bug: a token acquired BEFORE the grant + // carries an empty roles claim, and GraphTokenService caches that raw JWT (~55 min). After + // the admin grants the role and hits Refresh (→ InvalidateTenant), the detector must mint a + // FRESH token and report the new role — NOT re-parse the stale cached JWT. This exercises + // the cross-layer invalidation wiring with a REAL GraphTokenService (not the stub), so a + // future change that forgets to clear the token cache fails here. + const string tenantId = "66666666-6666-6666-6666-666666666666"; + + var handler = new SwappableTokenHandler( + BuildJwt(Array.Empty(), expires: DateTime.UtcNow.AddHours(1))); // pre-grant: no roles + var cache = new MemoryCache(new MemoryCacheOptions()); + var tokenService = new GraphTokenService( + NullLogger.Instance, + new SingleHandlerFactory(handler), + cache, + new ConfigurationBuilder() + .AddInMemoryCollection(new Dictionary + { + ["EntraId:ClientId"] = "test-client", + ["EntraId:ClientSecret"] = "test-secret", + }) + .Build()); + var telemetry = new TelemetryClient(new TelemetryConfiguration { DisableTelemetry = true }); + var detector = new GraphFeatureDetector( + tokenService, cache, NullLogger.Instance, telemetry); + + // Pre-grant read: empty roles, and the raw JWT is now cached in GraphTokenService. + var before = await detector.GetSnapshotAsync(tenantId); + Assert.Empty(before.GrantedRoles); + Assert.Equal(1, handler.Requests); + + // Admin grants the role; the wire now returns a JWT WITH it. + handler.Body = BuildJwt( + new[] { new Claim("roles", GraphAppPermissions.DeviceManagementScriptsReadAll) }, + expires: DateTime.UtcNow.AddHours(1)); + + // The fresh-read trigger. Without the cross-layer wiring the stale empty-roles JWT would + // still be served and this would keep reporting "not granted". + detector.InvalidateTenant(tenantId); + + var after = await detector.GetSnapshotAsync(tenantId); + Assert.Contains(GraphAppPermissions.DeviceManagementScriptsReadAll, after.GrantedRoles); + Assert.Equal(2, handler.Requests); // fresh token minted, not served from cache + } + + [Fact] + public async Task InvalidateTenant_during_in_flight_acquire_does_not_cache_stale_roles() + { + // Write-after-invalidate race on the DETECTOR's own cache: an acquire that started before + // InvalidateTenant must not leave a stale parsed-roles ctx behind. The token acquire is + // parked in flight; the admin grants + Refreshes (InvalidateTenant) while it is parked; on + // release the detector writes a (stale-generation-tagged) ctx. A fresh read MUST reject it + // and re-acquire, seeing the newly granted role. + const string tenantId = "77777777-7777-7777-7777-777777777777"; + var tokenService = new GatedStubGraphTokenService( + firstToken: BuildJwt(Array.Empty(), expires: DateTime.UtcNow.AddHours(1)), // pre-grant: no roles + laterToken: BuildJwt(new[] { new Claim("roles", GraphAppPermissions.DeviceManagementScriptsReadAll) }, + expires: DateTime.UtcNow.AddHours(1))); + var cache = new MemoryCache(new MemoryCacheOptions()); + var telemetry = new TelemetryClient(new TelemetryConfiguration { DisableTelemetry = true }); + var detector = new GraphFeatureDetector( + tokenService, cache, NullLogger.Instance, telemetry); + + // D1 starts; its token acquire parks in flight (genAtStart captured at method entry). + var inflight = detector.GetSnapshotAsync(tenantId); + await WaitUntilAsync(() => tokenService.CallCount >= 1, TimeSpan.FromSeconds(2)); + + // Admin grants + Refresh WHILE D1 is mid-acquire. + detector.InvalidateTenant(tenantId); + + // D1 completes → writes a CachedVerdict tagged with the pre-invalidation generation. + tokenService.Release(); + var before = await inflight; + Assert.Empty(before.GrantedRoles); // D1 itself saw the pre-grant token + + // Fresh read must reject the stale-tagged entry (gen mismatch) and re-acquire → new role. + var after = await detector.GetSnapshotAsync(tenantId); + Assert.Contains(GraphAppPermissions.DeviceManagementScriptsReadAll, after.GrantedRoles); + Assert.Equal(2, tokenService.CallCount); // a real re-acquire happened (no stale cache hit) + } + + [Fact] + public async Task InvalidateTenant_token_layer_first_prevents_stale_detector_write_in_overlap() + { + // Cross-layer ordering race: InvalidateTenant must invalidate the token layer BEFORE bumping + // its own generation. We pause INSIDE the token-layer invalidation and slip a concurrent + // GetSnapshotAsync into the gap. With the correct (token-first) order, that overlapping read + // cannot leave an accepted stale ctx behind; with the buggy (gen-first) order it would, and + // the post-refresh read below would still report the old (no-roles) verdict. + const string tenantId = "88888888-8888-8888-8888-888888888888"; + var handler = new SwappableTokenHandler( + BuildJwt(Array.Empty(), expires: DateTime.UtcNow.AddHours(1))); // pre-grant: no roles + var cache = new MemoryCache(new MemoryCacheOptions()); + + var entered = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var release = new System.Threading.ManualResetEventSlim(false); + var tokenService = new BarrieredInvalidateTokenService( + new SingleHandlerFactory(handler), cache, + onInvalidate: () => { entered.TrySetResult(); release.Wait(TimeSpan.FromSeconds(5)); }); + var telemetry = new TelemetryClient(new TelemetryConfiguration { DisableTelemetry = true }); + var detector = new GraphFeatureDetector( + tokenService, cache, NullLogger.Instance, telemetry); + + // Prime: cache the old (no-roles) token + parsed ctx on both layers. + var primed = await detector.GetSnapshotAsync(tenantId); + Assert.Empty(primed.GrantedRoles); + + // Admin grants the role on the wire. + handler.Body = BuildJwt( + new[] { new Claim("roles", GraphAppPermissions.DeviceManagementScriptsReadAll) }, + expires: DateTime.UtcNow.AddHours(1)); + + // Admin Refresh → InvalidateTenant, blocked inside the token-layer invalidation step. + var invalidateTask = Task.Run(() => detector.InvalidateTenant(tenantId)); + await entered.Task; + + // A concurrent read slips into the gap, then we let InvalidateTenant finish. + var overlap = await detector.GetSnapshotAsync(tenantId); + release.Set(); + await invalidateTask; + + // Post-refresh read MUST reflect the granted role — no stale entry survived the overlap. + var afterRefresh = await detector.GetSnapshotAsync(tenantId); + Assert.Contains(GraphAppPermissions.DeviceManagementScriptsReadAll, afterRefresh.GrantedRoles); + _ = overlap; // the overlapping read itself may legitimately observe pre-refresh state + } + + private static async Task WaitUntilAsync(Func condition, TimeSpan timeout) + { + var deadline = DateTime.UtcNow + timeout; + while (!condition() && DateTime.UtcNow < deadline) + { + await Task.Delay(10); + } + } + // ── Test seam: token service stub + detector factory ────────────────────────────── /// @@ -349,6 +488,85 @@ public class GraphFeatureDetectorTests } } + /// + /// Token-service stub whose FIRST acquire parks until is called (modelling + /// a token fetch in flight while an InvalidateTenant runs), returning firstToken; every + /// later acquire returns laterToken immediately. Drives the detector's in-flight + /// write-after-invalidate test. CallCount is incremented BEFORE parking and read thread-safely. + /// + private sealed class GatedStubGraphTokenService : GraphTokenService + { + private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly string _firstToken; + private readonly string _laterToken; + private int _calls; + + public GatedStubGraphTokenService(string firstToken, string laterToken) + : base( + NullLogger.Instance, + new NoopHttpClientFactory(), + new MemoryCache(new MemoryCacheOptions()), + new ConfigurationBuilder() + .AddInMemoryCollection(new Dictionary + { + ["EntraId:ClientId"] = "test-client-id", + ["EntraId:ClientSecret"] = "test-secret", + }) + .Build()) + { + _firstToken = firstToken; + _laterToken = laterToken; + } + + public int CallCount => System.Threading.Volatile.Read(ref _calls); + public void Release() => _release.TrySetResult(); + + public override async Task GetAccessTokenAsync(string tenantId, CancellationToken ct = default) + { + var n = System.Threading.Interlocked.Increment(ref _calls); + if (n == 1) + { + await _release.Task.ConfigureAwait(false); + return GraphTokenResult.Success(_firstToken); + } + return GraphTokenResult.Success(_laterToken); + } + } + + /// + /// Real GraphTokenService (real token cache + generation + HTTP path) whose InvalidateTenant is + /// wrapped with a test-controlled barrier, run BEFORE delegating to the base. Lets a test pause + /// inside the token-layer invalidation and interleave a concurrent acquire, to prove the + /// cross-layer ordering in GraphFeatureDetector.InvalidateTenant. + /// + private sealed class BarrieredInvalidateTokenService : GraphTokenService + { + private readonly Action _onInvalidate; + + public BarrieredInvalidateTokenService( + System.Net.Http.IHttpClientFactory factory, IMemoryCache cache, Action onInvalidate) + : base( + NullLogger.Instance, + factory, + cache, + new ConfigurationBuilder() + .AddInMemoryCollection(new Dictionary + { + ["EntraId:ClientId"] = "test-client-id", + ["EntraId:ClientSecret"] = "test-secret", + }) + .Build()) + { + _onInvalidate = onInvalidate; + } + + public override void InvalidateTenant(string tenantId) + { + _onInvalidate(); + base.InvalidateTenant(tenantId); + } + } + /// Placeholder factory — the detector never causes the base GraphTokenService HTTP path to run. private sealed class NoopHttpClientFactory : System.Net.Http.IHttpClientFactory { diff --git a/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/StubHttpMessageHandler.cs b/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/StubHttpMessageHandler.cs index 6b50b882..b062f871 100644 --- a/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/StubHttpMessageHandler.cs +++ b/src/Backend/AutopilotMonitor.Functions.Tests/GraphResolution/StubHttpMessageHandler.cs @@ -47,3 +47,94 @@ internal sealed class StubHttpClientFactory : System.Net.Http.IHttpClientFactory public StubHttpClientFactory(StubHttpMessageHandler handler) => _handler = handler; public HttpClient CreateClient(string name) => new(_handler, disposeHandler: false); } + +/// Factory that wires every client to one arbitrary . +internal sealed class SingleHandlerFactory : System.Net.Http.IHttpClientFactory +{ + private readonly HttpMessageHandler _handler; + public SingleHandlerFactory(HttpMessageHandler handler) => _handler = handler; + public HttpClient CreateClient(string name) => new(_handler, disposeHandler: false); +} + +/// +/// Returns a token response whose body can be swapped between calls — models AAD issuing a token +/// with a different roles claim after an admin grants a permission. Counts requests so a +/// test can assert a fresh acquire actually hit the wire (vs. being served from cache). +/// +internal sealed class SwappableTokenHandler : HttpMessageHandler +{ + private int _requests; + public string Body { get; set; } + public int Requests => System.Threading.Volatile.Read(ref _requests); + + public SwappableTokenHandler(string body) => Body = body; + + protected override Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + System.Threading.Interlocked.Increment(ref _requests); + return Task.FromResult(new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent($"{{\"access_token\":\"{Body}\",\"expires_in\":3599}}"), + }); + } +} + +/// +/// Parks the FIRST request until is called (modelling a token POST that is +/// in flight while an InvalidateTenant runs), then answers it with firstBody; every later +/// request answers immediately with laterBody (the post-grant re-mint). Lets a test prove a +/// stale in-flight write is rejected and a fresh read re-mints. +/// +internal sealed class InflightRaceHandler : HttpMessageHandler +{ + private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously); + private readonly string _firstBody; + private readonly string _laterBody; + private int _entered; + + public InflightRaceHandler(string firstBody, string laterBody) + { + _firstBody = firstBody; + _laterBody = laterBody; + } + + public int Entered => System.Threading.Volatile.Read(ref _entered); + public void Release() => _release.TrySetResult(); + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + var n = System.Threading.Interlocked.Increment(ref _entered); + if (n == 1) + { + await _release.Task.ConfigureAwait(false); + } + return new HttpResponseMessage(HttpStatusCode.OK) + { + Content = new StringContent(n == 1 ? _firstBody : _laterBody), + }; + } +} + +/// +/// Parks every request until is called, counting how many reached the wire. +/// Lets a stampede test prove that N simultaneous cache-misses collapse to a SINGLE in-flight POST +/// — which a synchronous stub cannot show, because the first call can populate the cache before the +/// others even start. +/// +internal sealed class GatedHttpMessageHandler : HttpMessageHandler +{ + private readonly string _body; + private readonly TaskCompletionSource _release = new(TaskCreationOptions.RunContinuationsAsynchronously); + private int _entered; + + public GatedHttpMessageHandler(string body) => _body = body; + public int Entered => System.Threading.Volatile.Read(ref _entered); + public void Release() => _release.TrySetResult(); + + protected override async Task SendAsync(HttpRequestMessage request, CancellationToken cancellationToken) + { + System.Threading.Interlocked.Increment(ref _entered); + await _release.Task.ConfigureAwait(false); + return new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(_body) }; + } +} diff --git a/src/Backend/AutopilotMonitor.Functions.Tests/GraphTokenServiceTests.cs b/src/Backend/AutopilotMonitor.Functions.Tests/GraphTokenServiceTests.cs new file mode 100644 index 00000000..1c95d760 --- /dev/null +++ b/src/Backend/AutopilotMonitor.Functions.Tests/GraphTokenServiceTests.cs @@ -0,0 +1,220 @@ +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 = new GraphTokenService( + NullLogger.Instance, + new StubHttpClientFactory(handler), + new MemoryCache(new MemoryCacheOptions()), + config); + + return (svc, handler); + } + + 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 = new GraphTokenService( + NullLogger.Instance, + new SingleHandlerFactory(handler), + new MemoryCache(new MemoryCacheOptions()), + config); + + // 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 = new GraphTokenService( + NullLogger.Instance, + new SingleHandlerFactory(handler), + new MemoryCache(new MemoryCacheOptions()), + config); + + 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); + } + + private static async Task WaitUntilAsync(Func condition, TimeSpan timeout) + { + var deadline = DateTime.UtcNow + timeout; + while (!condition() && DateTime.UtcNow < deadline) + { + await Task.Delay(10); + } + } +} diff --git a/src/Backend/AutopilotMonitor.Functions.Tests/IngestTelemetryFunctionTests.cs b/src/Backend/AutopilotMonitor.Functions.Tests/IngestTelemetryFunctionTests.cs index 15ff904f..0b36b5b4 100644 --- a/src/Backend/AutopilotMonitor.Functions.Tests/IngestTelemetryFunctionTests.cs +++ b/src/Backend/AutopilotMonitor.Functions.Tests/IngestTelemetryFunctionTests.cs @@ -48,53 +48,88 @@ public class IngestTelemetryFunctionTests [Fact] public async Task ReadBodyWithSizeCapAsync_accepts_body_under_cap() { - var payload = Encoding.UTF8.GetBytes("[{\"kind\":\"Event\"}]"); + var payload = Encoding.UTF8.GetBytes("[{\"Kind\":\"Event\"}]"); using var stream = new MemoryStream(payload); - var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024); + var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024); Assert.False(exceeded); - Assert.Equal("[{\"kind\":\"Event\"}]", body); + Assert.NotNull(items); + var item = Assert.Single(items!); + Assert.Equal("Event", item.Kind); + } + + [Fact] + public async Task ReadBodyWithSizeCapAsync_preserves_PayloadJson_as_raw_string() + { + // Hot-path contract: the inner payload stays an opaque escaped string on the envelope — + // the batch parse must NOT re-parse/re-serialise it (that's TelemetryPayloadParser's job, + // exactly once, downstream). Guards against an accidental round-trip that would mangle it. + var payload = Encoding.UTF8.GetBytes( + "[{\"Kind\":\"Event\",\"PayloadJson\":\"{\\\"EventType\\\":\\\"x\\\",\\\"n\\\":1}\"}]"); + using var stream = new MemoryStream(payload); + + var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024); + + Assert.False(exceeded); + var item = Assert.Single(items!); + Assert.Equal("{\"EventType\":\"x\",\"n\":1}", item.PayloadJson); } [Fact] public async Task ReadBodyWithSizeCapAsync_accepts_body_exactly_at_cap() { - // Strict greater-than semantics (matches legacy NdjsonParser): equal-to is OK. - var payload = new byte[100]; + // Strict greater-than semantics (matches legacy NdjsonParser): equal-to is OK. Pad a valid + // batch with trailing whitespace (ignored by the parser) to land exactly on the cap. + var json = "[{\"Kind\":\"Signal\",\"PayloadJson\":\"{}\"}]".PadRight(100); + var payload = Encoding.UTF8.GetBytes(json); + Assert.Equal(100, payload.Length); // ASCII → 1 byte/char using var stream = new MemoryStream(payload); - var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100); + var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100); Assert.False(exceeded); - Assert.Equal(100, Encoding.UTF8.GetByteCount(body)); + Assert.Single(items!); } [Fact] public async Task ReadBodyWithSizeCapAsync_rejects_body_over_cap_without_draining_full_stream() { - // 1 MB payload with a 100-byte cap → short-circuits after a single buffer read. + // 1 MB payload with a 100-byte cap → short-circuits after a single buffer read, before + // the deserialiser is ever invoked. var payload = new byte[1_000_000]; using var stream = new MemoryStream(payload); - var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100); + var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 100); Assert.True(exceeded); - Assert.Equal(string.Empty, body); + Assert.Null(items); // Stream position advanced past the cap but not to EOF — the helper bails early to // bound memory on hostile senders. Assert.True(stream.Position < stream.Length); } [Fact] - public async Task ReadBodyWithSizeCapAsync_empty_stream_returns_empty_body() + public async Task ReadBodyWithSizeCapAsync_empty_stream_returns_null_items() { using var stream = new MemoryStream(Array.Empty()); - var (exceeded, body) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024); + var (exceeded, items) = await IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024); Assert.False(exceeded); - Assert.Equal(string.Empty, body); + Assert.Null(items); + } + + [Fact] + public async Task ReadBodyWithSizeCapAsync_throws_on_malformed_json() + { + // Malformed JSON surfaces as JsonException so the HTTP trigger maps it to 400 — the cap + // path must not swallow it as an empty/successful parse. + var payload = Encoding.UTF8.GetBytes("[{\"Kind\":\"Event\""); + using var stream = new MemoryStream(payload); + + await Assert.ThrowsAnyAsync( + () => IngestTelemetryFunction.ReadBodyWithSizeCapAsync(stream, maxBytes: 1024)); } // ============================================================ Batch PartitionKey uniformity diff --git a/src/Backend/AutopilotMonitor.Functions/Functions/Ingest/IngestTelemetryFunction.cs b/src/Backend/AutopilotMonitor.Functions/Functions/Ingest/IngestTelemetryFunction.cs index 0451a44a..f0e46103 100644 --- a/src/Backend/AutopilotMonitor.Functions/Functions/Ingest/IngestTelemetryFunction.cs +++ b/src/Backend/AutopilotMonitor.Functions/Functions/Ingest/IngestTelemetryFunction.cs @@ -136,7 +136,18 @@ namespace AutopilotMonitor.Functions.Functions.Ingest var tenantConfig = await _configService.GetConfigurationAsync(tenantIdHeader); var maxPayloadBytes = (tenantConfig?.MaxNdjsonPayloadSizeMB ?? 5) * 1024 * 1024; - var (exceeded, body) = await ReadBodyWithSizeCapAsync(req.Body, maxPayloadBytes); + bool exceeded; + List? items; + try + { + (exceeded, items) = await ReadBodyWithSizeCapAsync(req.Body, maxPayloadBytes); + } + catch (JsonException ex) + { + _logger.LogWarning(ex, "IngestTelemetry: malformed JSON body"); + return AsOutput(await WriteErrorAsync(req, HttpStatusCode.BadRequest, "Malformed JSON body")); + } + if (exceeded) { _logger.LogWarning( @@ -147,17 +158,6 @@ namespace AutopilotMonitor.Functions.Functions.Ingest $"Payload exceeds {maxPayloadBytes / (1024 * 1024)} MB cap")); } - List? items; - try - { - items = JsonConvert.DeserializeObject>(body); - } - catch (JsonException ex) - { - _logger.LogWarning(ex, "IngestTelemetry: malformed JSON body"); - return AsOutput(await WriteErrorAsync(req, HttpStatusCode.BadRequest, "Malformed JSON body")); - } - if (items == null || items.Count == 0) { return AsOutput(await WriteErrorAsync(req, HttpStatusCode.BadRequest, "No telemetry items provided")); @@ -502,12 +502,24 @@ namespace AutopilotMonitor.Functions.Functions.Ingest } /// - /// Reads a stream into a string with a strict byte cap. Returns (true, "") as - /// soon as the cap would be exceeded — the stream isn't fully drained, which bounds - /// memory use even for a malicious sender. Matches the legacy NDJSON parser's guard - /// semantics (strict greater-than, so a payload equal to the cap is accepted). + /// Buffers the request body under a strict byte cap, then deserialises the telemetry batch + /// directly off the buffered bytes. Returns (true, null) as soon as the cap would be + /// exceeded — the stream isn't fully drained, which bounds memory use even for a malicious + /// sender (strict greater-than, so a payload equal to the cap is accepted; matches the + /// legacy NDJSON parser's guard semantics). + /// + /// Hot-path note: this deserialises through a over the buffered + /// rather than materialising a full UTF-16 string + a + /// ToArray() copy first. That removes ~3× the body size in transient allocations per + /// request on the highest-volume backend endpoint — the reader decodes UTF-8 incrementally + /// in 8 KB chunks. The cap is still applied before parsing, so an over-cap body + /// never reaches the deserialiser. + /// + /// Throws on malformed JSON (caller maps to 400). The + /// deserialiser uses the same default settings as JsonConvert.DeserializeObject, so + /// the parse semantics are unchanged — only the intermediate allocations are gone. /// - internal static async Task<(bool exceeded, string body)> ReadBodyWithSizeCapAsync( + internal static async Task<(bool exceeded, List? items)> ReadBodyWithSizeCapAsync( Stream source, int maxBytes) { using var buffered = new MemoryStream(); @@ -517,10 +529,15 @@ namespace AutopilotMonitor.Functions.Functions.Ingest while ((read = await source.ReadAsync(buffer, 0, buffer.Length)) > 0) { total += read; - if (total > maxBytes) return (true, string.Empty); + if (total > maxBytes) return (true, null); await buffered.WriteAsync(buffer, 0, read); } - return (false, System.Text.Encoding.UTF8.GetString(buffered.ToArray())); + + buffered.Position = 0; + using var textReader = new StreamReader(buffered, System.Text.Encoding.UTF8); + using var jsonReader = new JsonTextReader(textReader); + var items = JsonSerializer.CreateDefault().Deserialize>(jsonReader); + return (false, items); } /// diff --git a/src/Backend/AutopilotMonitor.Functions/Security/GraphTokenService.cs b/src/Backend/AutopilotMonitor.Functions/Security/GraphTokenService.cs index 104dfea7..7ecc857a 100644 --- a/src/Backend/AutopilotMonitor.Functions/Security/GraphTokenService.cs +++ b/src/Backend/AutopilotMonitor.Functions/Security/GraphTokenService.cs @@ -1,4 +1,5 @@ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Net.Http; using System.Threading; @@ -19,6 +20,32 @@ namespace AutopilotMonitor.Functions.Security { private static readonly TimeSpan ConsentStatusTtl = TimeSpan.FromMinutes(2); + // App-token cache tuning. TTL = expires_in - skew margin, clamped so a missing/pathological + // expires_in can neither thrash the cache nor pin a stale token. Max stays under the typical + // ~60-90 min AAD client-credentials lifetime; the 5-min margin covers clock skew between + // this worker and the Graph resource server plus the slowest downstream call in a request. + private static readonly TimeSpan TokenSkewMargin = TimeSpan.FromMinutes(5); + private static readonly TimeSpan MinTokenTtl = TimeSpan.FromSeconds(60); + private static readonly TimeSpan MaxTokenTtl = TimeSpan.FromMinutes(55); + private static readonly TimeSpan DefaultTokenLifetime = TimeSpan.FromMinutes(35); // when AAD omits expires_in + private static readonly TimeSpan TokenGateWait = TimeSpan.FromSeconds(10); + + // One coalescing gate per tenant (see GetAccessTokenAsync). A SemaphoreSlim is tiny and + // tenant cardinality is bounded + validated upstream, so these are never evicted. + private readonly ConcurrentDictionary _tokenLocks = new(); + + // Per-tenant invalidation generation. Closes the write-after-invalidate race: a token POST + // that STARTED before InvalidateTenant captures the old generation and tags its cache entry + // with it; InvalidateTenant bumps the generation, so that late stale write is rejected by + // every reader (entry.Gen != current) even though it physically lands after the Remove. + // Lock-free — the comparison happens at read time against the latest generation. + private readonly ConcurrentDictionary _tokenCacheGen = new(); + + private long CurrentGen(string tenantId) => _tokenCacheGen.TryGetValue(tenantId, out var g) ? g : 0; + + /// Cached app token tagged with the invalidation generation in effect when its acquire started. + private sealed record CachedAppToken(string Token, long Gen); + private readonly ILogger _logger; private readonly IHttpClientFactory _httpClientFactory; private readonly IMemoryCache _cache; @@ -38,7 +65,7 @@ namespace AutopilotMonitor.Functions.Security public async Task GetConsentStatusAsync(string tenantId, CancellationToken ct = default) { - var cacheKey = $"graph-consent-status:{tenantId}"; + var cacheKey = ConsentCacheKey(tenantId); if (_cache.TryGetValue(cacheKey, out GraphConsentStatusResult? cached) && cached != null) { return cached; @@ -76,6 +103,14 @@ namespace AutopilotMonitor.Functions.Security public virtual async Task GetAccessTokenAsync(string tenantId, CancellationToken ct = default) { + // Serve a still-valid cached app token for this tenant. The cache key is keyed strictly + // on the (already authenticated upstream) tenantId — the same value that drives the token + // URL below — so a token minted for one tenant can never be served to another. + if (TryGetCachedToken(tenantId, out var cachedToken)) + { + return GraphTokenResult.Success(cachedToken); + } + var clientId = _configuration["EntraId:ClientId"]; var clientSecret = _configuration["EntraId:ClientSecret"]; @@ -86,6 +121,54 @@ namespace AutopilotMonitor.Functions.Security return GraphTokenResult.PermanentFailure(); } + // Stampede guard: coalesce concurrent cache-misses for the same tenant so an enrollment + // wave doesn't fan a burst of identical token POSTs at AAD (which would trip its own + // throttling). One acquirer fetches + populates the cache; the rest hit it after the + // gate. The wait is bounded so a pathological slow acquire (e.g. a no-consent tenant on + // the long retry chain) degrades to today's parallel-fetch behaviour instead of pinning + // callers indefinitely. + var gate = _tokenLocks.GetOrAdd(tenantId, _ => new SemaphoreSlim(1, 1)); + bool gotGate; + try + { + gotGate = await gate.WaitAsync(TokenGateWait, ct); + } + catch (OperationCanceledException) when (ct.IsCancellationRequested) + { + // Caller budget exhausted while waiting — transient by definition (never cache). + return GraphTokenResult.TransientFailure(); + } + + try + { + // Double-check: another thread may have populated the cache while we waited. The + // gen-aware read also rejects a stale entry written by an acquire that overlapped an + // InvalidateTenant, so we fall through to a fresh mint instead of serving it. + if (gotGate && TryGetCachedToken(tenantId, out var tokenAfterWait)) + { + return GraphTokenResult.Success(tokenAfterWait); + } + + return await AcquireTokenViaClientCredentialsAsync(tenantId, clientId, clientSecret, ct); + } + finally + { + if (gotGate) gate.Release(); + } + } + + /// + /// Performs the client-credentials POST (with the consent-propagation retry chain) and + /// caches a successful token until shortly before its expiry. Failures are never cached, so + /// a tenant that loses consent re-acquires (and fails) on the next request after expiry. + /// + private async Task AcquireTokenViaClientCredentialsAsync( + string tenantId, string clientId, string clientSecret, CancellationToken ct) + { + // Capture the invalidation generation BEFORE the POST. If an InvalidateTenant lands + // while this POST is in flight, the generation changes and the token we cache below is + // tagged with the now-stale value — so readers reject it and re-mint fresh. + var genAtStart = CurrentGen(tenantId); var tokenUrl = $"https://login.microsoftonline.com/{tenantId}/oauth2/v2.0/token"; // Retry with backoff to handle Azure AD consent propagation delays. @@ -124,6 +207,12 @@ namespace AutopilotMonitor.Functions.Security { var tokenJson = JsonConvert.DeserializeObject(responseBody); var accessToken = tokenJson?["access_token"]?.ToString(); + if (!string.IsNullOrEmpty(accessToken)) + { + // Cache the freshly minted token (never log its value). TTL derives from + // the token's own expires_in, minus a skew margin, clamped to a safe band. + CacheToken(tenantId, accessToken!, genAtStart, (long?)tokenJson?["expires_in"]); + } return GraphTokenResult.Success(accessToken); } @@ -208,6 +297,64 @@ namespace AutopilotMonitor.Functions.Security } } } + + private static string TokenCacheKey(string tenantId) => $"graph-token:{tenantId}"; + private static string ConsentCacheKey(string tenantId) => $"graph-consent-status:{tenantId}"; + + /// + /// Drops this tenant's cached app token (and consent-status verdict) so the next acquire + /// mints a fresh JWT from AAD. Required for the consent / Graph-permission fresh-read + /// contract: after an admin grants consent or assigns app roles, the OLD token still carries + /// the OLD roles claim until it expires — serving it would make a post-grant "Refresh" + /// keep reporting "not granted" for up to the cache TTL. The single fresh-read entry point + /// GraphFeatureDetector.InvalidateTenant calls this so BOTH cache layers are cleared + /// together (the detector's parsed-roles cache AND this raw-token cache). + /// + public virtual void InvalidateTenant(string tenantId) + { + if (string.IsNullOrWhiteSpace(tenantId)) return; + // Bump the generation FIRST so any token POST already in flight (which captured the old + // generation) writes an entry that readers reject as stale — this closes the + // write-after-invalidate race, not just the already-written entry. Then drop the current + // entries to free memory and clear the coarse consent-status verdict. + _tokenCacheGen.AddOrUpdate(tenantId, 1, (_, v) => v + 1); + _cache.Remove(TokenCacheKey(tenantId)); + _cache.Remove(ConsentCacheKey(tenantId)); + } + + private bool TryGetCachedToken(string tenantId, out string token) + { + // Accept the cached token only if its generation still matches the tenant's current one. + // A mismatch means an InvalidateTenant ran since this entry's acquire started, so the + // token may carry a stale roles claim — treat it as a miss and re-mint. + if (_cache.TryGetValue(TokenCacheKey(tenantId), out CachedAppToken? cached) + && cached != null + && !string.IsNullOrEmpty(cached.Token) + && cached.Gen == CurrentGen(tenantId)) + { + token = cached.Token; + return true; + } + token = string.Empty; + return false; + } + + private void CacheToken(string tenantId, string accessToken, long genAtStart, long? expiresInSeconds) + { + // expires_in is reported by AAD over TLS, so it is trustworthy — but still clamped so a + // missing or pathological value cannot pin a stale token (cap) or cause thrash (floor). + var lifetime = expiresInSeconds is > 0 + ? TimeSpan.FromSeconds(expiresInSeconds.Value) + : DefaultTokenLifetime; + var ttl = lifetime - TokenSkewMargin; + if (ttl < MinTokenTtl) ttl = MinTokenTtl; + if (ttl > MaxTokenTtl) ttl = MaxTokenTtl; + + _cache.Set(TokenCacheKey(tenantId), new CachedAppToken(accessToken, genAtStart), new MemoryCacheEntryOptions + { + AbsoluteExpirationRelativeToNow = ttl, + }); + } } /// diff --git a/src/Backend/AutopilotMonitor.Functions/Services/GraphResolution/GraphFeatureDetector.cs b/src/Backend/AutopilotMonitor.Functions/Services/GraphResolution/GraphFeatureDetector.cs index c4b818fc..dae9b1d3 100644 --- a/src/Backend/AutopilotMonitor.Functions/Services/GraphResolution/GraphFeatureDetector.cs +++ b/src/Backend/AutopilotMonitor.Functions/Services/GraphResolution/GraphFeatureDetector.cs @@ -1,4 +1,5 @@ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.Globalization; using System.IdentityModel.Tokens.Jwt; @@ -43,6 +44,21 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector public static readonly NegativeAcquireMarker Instance = new(); } + /// + /// Cache entry tagged with the per-tenant invalidation generation in effect when its acquire + /// started. is either a (success) + /// or (definitive no-consent). Readers accept it only while + /// still matches the tenant's current generation — a stale write that lands + /// after an carries an older generation and is rejected. + /// + private sealed record CachedVerdict(object Payload, long Gen); + + // Per-tenant invalidation generation (see CachedVerdict + InvalidateTenant). Mirrors the + // GraphTokenService fix one layer down so the fresh-read contract holds on this cache too. + private readonly ConcurrentDictionary _invalidationGen = new(); + + private long CurrentGen(string tenantId) => _invalidationGen.TryGetValue(tenantId, out var g) ? g : 0; + private readonly GraphTokenService _tokenService; private readonly IMemoryCache _cache; private readonly ILogger _logger; @@ -111,15 +127,23 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector { if (string.IsNullOrWhiteSpace(tenantId)) return (null, false); + // Capture the invalidation generation BEFORE the acquire. If InvalidateTenant runs while we + // fetch/parse the token below, the generation changes and the entry we write is tagged with + // the now-stale value — so readers reject it and re-acquire. Mirrors the GraphTokenService + // fix one layer down, closing the write-after-invalidate race on this cache too. + var genAtStart = CurrentGen(tenantId); + var cacheKey = BuildCacheKey(tenantId); - if (_cache.TryGetValue(cacheKey, out object? cached) && cached != null) + if (_cache.TryGetValue(cacheKey, out CachedVerdict? cached) + && cached != null + && cached.Gen == CurrentGen(tenantId)) { - if (cached is NegativeAcquireMarker) + if (cached.Payload is NegativeAcquireMarker) { // Definitive "no consent / no SP in this tenant" verdict from a recent attempt. return (null, false); } - if (cached is GraphTenantTokenContext cachedCtx + if (cached.Payload is GraphTenantTokenContext cachedCtx && cachedCtx.ExpiresAt > _time.GetUtcNow() + ExpirySafetyMargin) { return (cachedCtx, false); @@ -145,7 +169,7 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector // revoked, or the app credentials are misconfigured. Cache the verdict briefly so // repeated cross-tenant clicks don't keep replaying the consent-propagation // retry chain for the same dead-end tenant. - CacheNegativeVerdict(cacheKey); + CacheNegativeVerdict(cacheKey, genAtStart); return (null, false); } @@ -165,7 +189,10 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector }; var ttl = ComputeTtl(expiresAt); - _cache.Set(cacheKey, ctx, new MemoryCacheEntryOptions { AbsoluteExpirationRelativeToNow = ttl }); + // Tag with the generation captured before the acquire (see top of method). A concurrent + // InvalidateTenant bumps the generation, so this write — if it lands after the invalidation — + // is rejected by readers via the gen check, closing the write-after-invalidate race. + _cache.Set(cacheKey, new CachedVerdict(ctx, genAtStart), new MemoryCacheEntryOptions { AbsoluteExpirationRelativeToNow = ttl }); _logger.LogDebug( "GraphFeatureDetector: cached token for tenant {TenantId} with {RoleCount} roles, TTL={Ttl}", @@ -201,15 +228,27 @@ public sealed class GraphFeatureDetector : IGraphFeatureDetector public void InvalidateTenant(string tenantId) { if (string.IsNullOrWhiteSpace(tenantId)) return; - // Single cache key carries both the success token AND the negative marker, so one - // Remove call clears whichever verdict is currently in flight. + // ORDER IS LOAD-BEARING. Invalidate the UNDERLYING token layer FIRST, then bump OUR + // generation, then drop our cache. Token-first is what closes the cross-layer overlap race: + // otherwise a GetSnapshotAsync starting in the gap could capture our NEW generation yet + // still read the OLD (not-yet-invalidated) GraphTokenService cache, and then write a + // stale-roles ctx tagged with the new generation — which readers would ACCEPT, re-opening + // the stale "not granted" window until our TTL. Because an acquire captures its generation + // at method entry (before its token read), invalidating the token layer first guarantees + // any acquire able to read a stale token captured its generation before this bump, so its + // write loses the gen check. (This also drops the raw-token cache so our own next read can't + // re-parse a stale JWT — the original cross-layer fix.) + _tokenService.InvalidateTenant(tenantId); + _invalidationGen.AddOrUpdate(tenantId, 1, (_, v) => v + 1); + // Single cache key carries both the success token AND the negative marker, so one Remove + // clears whichever verdict is currently in flight. _cache.Remove(BuildCacheKey(tenantId)); _logger.LogInformation("GraphFeatureDetector: cache invalidated for tenant {TenantId}", tenantId); } - private void CacheNegativeVerdict(string cacheKey) + private void CacheNegativeVerdict(string cacheKey, long gen) { - _cache.Set(cacheKey, NegativeAcquireMarker.Instance, new MemoryCacheEntryOptions + _cache.Set(cacheKey, new CachedVerdict(NegativeAcquireMarker.Instance, gen), new MemoryCacheEntryOptions { AbsoluteExpirationRelativeToNow = NegativeCacheTtl, });