fix(entitlements): address Codex review findings

- High: TenantConfigurationService.SaveConfigurationAsync now THROWS when the repository reports a failed persist (it swallows storage exceptions and returns false) - plan/trial/config writes can no longer be audited and returned as 200 without being persisted. Cache is invalidated on every attempted save (finally), covering in-place mutations of the cached instance; the now-redundant invalidate-on-failure wrapper in PlanManagementFunction is removed. AuthFunction login side-effect saves (first-login domain write, auto-re-enable) are wrapped best-effort so a storage blip cannot 500 auth/me - they self-gate and retry on next login. - Medium: MCP usage increment moved from AuthenticationMiddleware (fire-and-forget BEFORE the quota decision) into McpQuotaEnforcementMiddleware: strictly check-then-increment and only for requests actually served (GA tracked but quota-exempt; denied 403/429 requests no longer inflate counters; the exact-limit request is now deterministic - 1..limit pass, limit+1 blocks). - Low: fixed both new nullable warnings (PlanTierDefinitionParser guarded json, McpQuotaEnforcementMiddleware principal null-flow). Tests: +3 save-semantics regression tests (throw on false, cache invalidation on failed/successful save). Backend 2656/2656 green, build clean.
2026-10-24 04:32:32 +00:00 · 2026-07-07 13:38:00 +02:00
parent d8e8cd15fd
commit dfa0fe27c1
7 changed files with 156 additions and 69 deletions
@@ -100,4 +100,57 @@ public sealed class TenantConfigurationServiceReadSemanticsTests
        Assert.True(exists);
        Assert.Equal(TenantId, config.TenantId);
    }
+
+    // ── Save semantics (Codex finding 2026-07-07): the repository swallows storage exceptions
+    //    and reports failure via its bool return — the service must THROW on false so callers
+    //    (plan/trial endpoints, config PUT) can never audit + 200 a write that never persisted.
+
+    [Fact]
+    public async Task SaveConfigurationAsync_RepoReportsFalse_Throws()
+    {
+        var (service, repo) = Build();
+        repo.Setup(r => r.SaveTenantConfigurationAsync(It.IsAny<TenantConfiguration>()))
+            .ReturnsAsync(false);
+
+        await Assert.ThrowsAsync<InvalidOperationException>(
+            () => service.SaveConfigurationAsync(new TenantConfiguration { TenantId = TenantId, UpdatedBy = "alice@contoso.com" }));
+    }
+
+    [Fact]
+    public async Task SaveConfigurationAsync_FailedSave_InvalidatesCachedInstance()
+    {
+        // The plan/trial endpoints mutate the CACHED instance in place before saving. A failed
+        // save must drop that instance from the cache — otherwise the unsaved mutation is served
+        // as if persisted for up to 5 minutes.
+        var (service, repo) = Build();
+        var stored = new TenantConfiguration { TenantId = TenantId, PlanTier = "free", UpdatedBy = "x" };
+        repo.Setup(r => r.GetTenantConfigurationAsync(TenantId)).ReturnsAsync(stored);
+        repo.Setup(r => r.SaveTenantConfigurationAsync(It.IsAny<TenantConfiguration>()))
+            .ReturnsAsync(false);
+
+        var cached = await service.GetConfigurationIfExistsAsync(TenantId);
+        cached!.PlanTier = "enterprise"; // in-place mutation, as the plan endpoint does
+
+        await Assert.ThrowsAsync<InvalidOperationException>(() => service.SaveConfigurationAsync(cached));
+
+        // Next read must go back to storage (cache dropped), not serve the mutated instance.
+        await service.GetConfigurationIfExistsAsync(TenantId);
+        repo.Verify(r => r.GetTenantConfigurationAsync(TenantId), Times.Exactly(2));
+    }
+
+    [Fact]
+    public async Task SaveConfigurationAsync_Success_InvalidatesCache()
+    {
+        var (service, repo) = Build();
+        repo.Setup(r => r.GetTenantConfigurationAsync(TenantId))
+            .ReturnsAsync(new TenantConfiguration { TenantId = TenantId, UpdatedBy = "x" });
+        repo.Setup(r => r.SaveTenantConfigurationAsync(It.IsAny<TenantConfiguration>()))
+            .ReturnsAsync(true);
+
+        await service.GetConfigurationIfExistsAsync(TenantId); // prime cache
+        await service.SaveConfigurationAsync(new TenantConfiguration { TenantId = TenantId, UpdatedBy = "x" });
+        await service.GetConfigurationIfExistsAsync(TenantId); // must re-read after save
+
+        repo.Verify(r => r.GetTenantConfigurationAsync(TenantId), Times.Exactly(2));
+    }
 }
@@ -156,7 +156,10 @@ namespace AutopilotMonitor.Functions.Functions.Config
                if (changes.Count > 0)
                {
                    config.UpdatedBy = caller;
-                    await SaveInvalidatingOnFailureAsync(config);
+                    // Fail-loud: SaveConfigurationAsync throws when the write did not persist (and
+                    // invalidates the cached — possibly mutated — instance in its finally), so a
+                    // failed save can never be audited or returned as 200.
+                    await _configService.SaveConfigurationAsync(config);

                    await _maintenanceRepo.LogAuditEntryAsync(
                        requestCtx.TargetTenantId,
@@ -243,7 +246,8 @@ namespace AutopilotMonitor.Functions.Functions.Config
                config.TrialGrantedBy = caller;
                config.UpdatedBy = caller;

-                await SaveInvalidatingOnFailureAsync(config);
+                // Fail-loud: throws when the write did not persist (cache invalidated in finally).
+                await _configService.SaveConfigurationAsync(config);

                await _maintenanceRepo.LogAuditEntryAsync(
                    requestCtx.TargetTenantId,
@@ -346,24 +350,6 @@ namespace AutopilotMonitor.Functions.Functions.Config
            }
        }

-        /// <summary>
-        /// Saves via the service (which invalidates the cache on success). Because the mutated
-        /// instance IS the cached instance, a failed save would otherwise leave the cache holding
-        /// unsaved mutations for up to 5 minutes — invalidate before rethrowing.
-        /// </summary>
-        private async Task SaveInvalidatingOnFailureAsync(AutopilotMonitor.Shared.Models.TenantConfiguration config)
-        {
-            try
-            {
-                await _configService.SaveConfigurationAsync(config);
-            }
-            catch
-            {
-                _configService.InvalidateCache(config.TenantId);
-                throw;
-            }
-        }
-
        private static string FormatUtc(DateTime? value)
            => value?.ToString("yyyy-MM-ddTHH:mm:ssZ") ?? "(none)";

@@ -302,7 +302,19 @@ public class AuthFunction
        // sync jobs that mutate UpdatedBy cannot leak sentinel strings into TenantAdmins.
        if (string.IsNullOrWhiteSpace(tenantConfig.OnboardedBy))
            tenantConfig.OnboardedBy = upn;
-        await _tenantConfigService.SaveConfigurationAsync(tenantConfig);
+        try
+        {
+            await _tenantConfigService.SaveConfigurationAsync(tenantConfig);
+        }
+        catch (Exception ex)
+        {
+            // Best-effort side-effect: SaveConfigurationAsync now throws on a failed persist
+            // (fail-loud for config endpoints), but a storage blip must not 500 the login.
+            // The handler self-gates on DomainName being empty, so the next login retries —
+            // skip the notifications too so they only fire once the write actually stuck.
+            _logger.LogWarning(ex, "First-login domain write failed for tenant {TenantId} — will retry on next login", tenantId);
+            return;
+        }

        // Fire-and-forget: Telegram
        _ = _telegramNotificationService.SendNewTenantSignupAsync(tenantId, upn)
@@ -336,7 +348,16 @@ public class AuthFunction
        tenantConfig.DisabledReason = null;
        tenantConfig.DisabledUntil = null;
        tenantConfig.UpdatedBy = "System (auto-re-enable)";
-        await _tenantConfigService.SaveConfigurationAsync(tenantConfig);
+        try
+        {
+            await _tenantConfigService.SaveConfigurationAsync(tenantConfig);
+        }
+        catch (Exception ex)
+        {
+            // Best-effort: the in-memory flip above already lets THIS login through; a failed
+            // persist just means the next login re-runs the auto-re-enable. Must not 500 auth/me.
+            _logger.LogWarning(ex, "Auto-re-enable persist failed for tenant {TenantId} — will retry on next login", tenantId);
+        }
    }

    /// <summary>
@@ -10,10 +10,7 @@ using System.IdentityModel.Tokens.Jwt;
 using Microsoft.AspNetCore.Http;
 using System.Net;
 using System.Security.Claims;
-using AutopilotMonitor.Functions.Extensions;
-using AutopilotMonitor.Functions.Helpers;
 using AutopilotMonitor.Functions.Security;
-using AutopilotMonitor.Shared.DataAccess;

 namespace AutopilotMonitor.Functions.Middleware;

@@ -26,7 +23,6 @@ public class AuthenticationMiddleware : IFunctionsWorkerMiddleware
    private readonly ILogger<AuthenticationMiddleware> _logger;
    private readonly IConfiguration _configuration;
    private readonly JwtSecurityTokenHandler _tokenHandler;
-    private readonly IUserUsageRepository _userUsageRepo;

    // Cache configuration managers per tenant to avoid repeated OIDC metadata fetches
    // Bounded with LRU eviction to prevent memory exhaustion from malicious tenant ID flooding
@@ -36,12 +32,10 @@ public class AuthenticationMiddleware : IFunctionsWorkerMiddleware

    public AuthenticationMiddleware(
        ILogger<AuthenticationMiddleware> logger,
-        IConfiguration configuration,
-        IUserUsageRepository userUsageRepo)
+        IConfiguration configuration)
    {
        _logger = logger;
        _configuration = configuration;
-        _userUsageRepo = userUsageRepo;
        _tokenHandler = new JwtSecurityTokenHandler();
        _configManagerCache = new Dictionary<string, (IConfigurationManager<OpenIdConnectConfiguration> Manager, DateTime LastAccessed)>();

@@ -222,35 +216,11 @@ public class AuthenticationMiddleware : IFunctionsWorkerMiddleware
            context.Items["ClaimsPrincipal"] = principal;
            authenticated = true;

-            // Track MCP usage only (identified by X-Client-Source header, fire-and-forget, non-blocking)
-            var isMcpRequest = string.Equals(
-                httpContext.Request.Headers["X-Client-Source"].FirstOrDefault(), "mcp", StringComparison.OrdinalIgnoreCase);
-            if (isMcpRequest)
-            {
-                var oid = principal.GetObjectId();
-                if (!string.IsNullOrEmpty(oid))
-                {
-                    var upn = principal.GetUserPrincipalName() ?? "unknown";
-                    var tid = principal.GetTenantId() ?? "";
-                    var normalizedEndpoint = EndpointNormalizer.Normalize(requestPath);
-                    var mcpToolName = httpContext.Request.Headers["X-MCP-Tool-Name"].FirstOrDefault();
-                    if (!string.IsNullOrEmpty(mcpToolName))
-                        normalizedEndpoint = $"{mcpToolName}:{normalizedEndpoint}";
-                    var repo = _userUsageRepo;
-                    var logger = _logger;
-                    _ = Task.Run(async () =>
-                    {
-                        try
-                        {
-                            await repo.IncrementUsageAsync(oid, upn, tid, normalizedEndpoint);
-                        }
-                        catch (Exception ex)
-                        {
-                            logger.LogWarning(ex, "[Auth Middleware] Failed to record usage: user={UserId}, endpoint={Endpoint}", oid, normalizedEndpoint);
-                        }
-                    });
-                }
-            }
+            // NOTE: MCP usage tracking moved to McpQuotaEnforcementMiddleware (Codex finding,
+            // 2026-07-07): incrementing here — BEFORE the quota decision — raced the quota read
+            // (the exact-limit request depended on whether the fire-and-forget increment won),
+            // and denied (403/429) requests inflated the counters. The quota middleware now
+            // counts strictly check-then-increment, and only requests that are actually served.
        }
        catch (SecurityTokenValidationException ex)
        {
@@ -2,6 +2,7 @@ using System.Net;
 using AutopilotMonitor.Functions.Extensions;
 using AutopilotMonitor.Functions.Helpers;
 using AutopilotMonitor.Functions.Services;
+using AutopilotMonitor.Shared.DataAccess;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Azure.Functions.Worker;
 using Microsoft.Azure.Functions.Worker.Middleware;
@@ -13,7 +14,14 @@ namespace AutopilotMonitor.Functions.Middleware;
 /// Enforces the per-user MCP daily/monthly request quota. Runs after
 /// <see cref="UserRateLimitMiddleware"/> (per-minute burst control) — this is the budget layer on
 /// top. Applies ONLY to HTTP requests marked <c>X-Client-Source: mcp</c> with an authenticated
-/// principal (oid); everything else passes through untouched. Global Admins are exempt.
+/// principal (oid); everything else passes through untouched. Global Admins are exempt from the
+/// quota (but their usage is still tracked for the metrics surface).
+///
+/// This middleware also OWNS the usage-counter increment (moved here from AuthenticationMiddleware,
+/// Codex finding 2026-07-07): strictly check-then-increment, and only for requests that are
+/// actually served — the request that hits the limit is deterministic (requests 1..limit pass,
+/// limit+1 blocks) and denied requests (403 upstream, 429 here) no longer inflate the counters.
+/// The increment itself stays fire-and-forget (never blocks the request path).
 ///
 /// Over-quota requests get 429 with a structured body, <c>Retry-After</c>, and
 /// <c>X-MCP-Quota-*</c> headers; allowed MCP requests get the quota headers too so the MCP
@@ -23,13 +31,16 @@ namespace AutopilotMonitor.Functions.Middleware;
 public class McpQuotaEnforcementMiddleware : IFunctionsWorkerMiddleware
 {
    private readonly McpQuotaService _quotaService;
+    private readonly IUserUsageRepository _userUsageRepo;
    private readonly ILogger<McpQuotaEnforcementMiddleware> _logger;

    public McpQuotaEnforcementMiddleware(
        McpQuotaService quotaService,
+        IUserUsageRepository userUsageRepo,
        ILogger<McpQuotaEnforcementMiddleware> logger)
    {
        _quotaService = quotaService;
+        _userUsageRepo = userUsageRepo;
        _logger = logger;
    }

@@ -52,17 +63,21 @@ public class McpQuotaEnforcementMiddleware : IFunctionsWorkerMiddleware

        var principal = context.GetUser();
        var oid = principal?.GetObjectId();
-        if (string.IsNullOrEmpty(oid))
+        if (principal == null || string.IsNullOrEmpty(oid))
        {
            // Unauthenticated MCP probe — auth middleware / policy enforcement handle rejection.
            await next(context);
            return;
        }

-        // GA exempt: platform operators are not customers of the quota.
+        var upn = principal.GetUserPrincipalName();
+        var tenantId = principal.GetTenantId();
+
+        // GA exempt from the QUOTA — but their usage is still tracked (metrics surface).
        var requestContext = context.GetRequestContext();
        if (requestContext.IsGlobalAdmin)
        {
+            TrackUsage(httpContext, oid, upn, tenantId);
            await next(context);
            return;
        }
@@ -70,14 +85,13 @@ public class McpQuotaEnforcementMiddleware : IFunctionsWorkerMiddleware
        McpQuotaDecision decision;
        try
        {
-            var upn = principal!.GetUserPrincipalName();
-            var tenantId = principal.GetTenantId();
            decision = await _quotaService.CheckAsync(oid, upn, tenantId);
        }
        catch (Exception ex)
        {
            // Belt-and-braces fail-open: the quota layer must never take MCP down.
            _logger.LogError(ex, "[McpQuota] Quota check threw for oid={Oid} — allowing request (fail-open)", oid);
+            TrackUsage(httpContext, oid, upn, tenantId);
            await next(context);
            return;
        }
@@ -86,6 +100,8 @@ public class McpQuotaEnforcementMiddleware : IFunctionsWorkerMiddleware

        if (decision.Allowed)
        {
+            // Check-then-increment: the decision above reflects previously SERVED requests only.
+            TrackUsage(httpContext, oid, upn, tenantId);
            await next(context);
            return;
        }
@@ -110,6 +126,35 @@ public class McpQuotaEnforcementMiddleware : IFunctionsWorkerMiddleware
        });
    }

+    /// <summary>
+    /// Fire-and-forget usage increment (same posture as the previous AuthenticationMiddleware
+    /// tracking — never blocks or fails the request path). Endpoint is normalized and prefixed
+    /// with the X-MCP-Tool-Name when the MCP server supplies it.
+    /// </summary>
+    private void TrackUsage(HttpContext httpContext, string oid, string? upn, string? tenantId)
+    {
+        var normalizedEndpoint = EndpointNormalizer.Normalize(httpContext.Request.Path.Value ?? string.Empty);
+        var mcpToolName = httpContext.Request.Headers["X-MCP-Tool-Name"].FirstOrDefault();
+        if (!string.IsNullOrEmpty(mcpToolName))
+            normalizedEndpoint = $"{mcpToolName}:{normalizedEndpoint}";
+
+        var repo = _userUsageRepo;
+        var logger = _logger;
+        var upnValue = upn ?? "unknown";
+        var tidValue = tenantId ?? "";
+        _ = Task.Run(async () =>
+        {
+            try
+            {
+                await repo.IncrementUsageAsync(oid, upnValue, tidValue, normalizedEndpoint);
+            }
+            catch (Exception ex)
+            {
+                logger.LogWarning(ex, "[McpQuota] Failed to record usage: user={UserId}, endpoint={Endpoint}", oid, normalizedEndpoint);
+            }
+        });
+    }
+
    private static void StampQuotaHeaders(HttpContext httpContext, McpQuotaDecision decision)
    {
        // Direct-write pattern — same as UserRateLimitMiddleware's X-RateLimit-* headers.
@@ -86,7 +86,12 @@ namespace AutopilotMonitor.Functions.Services
        }

        /// <summary>
-        /// Saves configuration for a tenant
+        /// Saves configuration for a tenant. THROWS when the write did not persist — the repository
+        /// swallows storage exceptions and reports failure via its bool return, so ignoring it would
+        /// let callers audit + return 200 for a save that never happened (Codex finding, 2026-07-07).
+        /// The cached row is invalidated on EVERY attempted save (finally): on success it is stale;
+        /// on failure the caller may have mutated the cached instance in place (plan/trial endpoints),
+        /// which must not linger as phantom-saved state for the 5-minute TTL.
        /// </summary>
        public virtual async Task SaveConfigurationAsync(TenantConfiguration config)
        {
@@ -99,10 +104,12 @@ namespace AutopilotMonitor.Functions.Services
            {
                config.LastUpdated = DateTime.UtcNow;

-                await _configRepo.SaveTenantConfigurationAsync(config);
-
-                // Invalidate cache
-                _cache.Remove($"tenant-config:{config.TenantId}");
+                var saved = await _configRepo.SaveTenantConfigurationAsync(config);
+                if (!saved)
+                {
+                    throw new InvalidOperationException(
+                        $"Tenant configuration save failed for {config.TenantId} — storage write did not persist (see repository logs)");
+                }

                _logger.LogInformation($"Configuration saved for tenant {config.TenantId} by {config.UpdatedBy}");
            }
@@ -111,6 +118,11 @@ namespace AutopilotMonitor.Functions.Services
                _logger.LogError(ex, $"Error saving configuration for tenant {config.TenantId}");
                throw;
            }
+            finally
+            {
+                // Invalidate cache — success AND failure paths (see summary).
+                _cache.Remove($"tenant-config:{config.TenantId}");
+            }
        }

        /// <summary>
@@ -19,7 +19,7 @@ namespace AutopilotMonitor.Shared.Models.Config

            try
            {
-                return JsonSerializer.Deserialize<List<PlanTierDefinition>>(json, Options)
+                return JsonSerializer.Deserialize<List<PlanTierDefinition>>(json!, Options)
                       ?? new List<PlanTierDefinition>();
            }
            catch