AdPlatform-Server/Management/Security/ClientAuthMiddleware.cs

using Management.Data;
using Microsoft.IdentityModel.Protocols;
using Microsoft.IdentityModel.Protocols.OpenIdConnect;
using Microsoft.IdentityModel.Tokens;
using System.IdentityModel.Tokens.Jwt;
using System.Security.Claims;
using System.Text.Json;

namespace Management.Security;

/// <summary>
/// Authentication middleware for Management API.
/// 
/// Auth paths:
/// - /api/onboarding/*  → JWT (Entra, any staff role)
/// - /api/monitoring/*  → JWT (Entra, Staff.Admin or Staff.Tech role)
/// - /api/admin/*       → JWT (Entra, Staff.Admin role only)
/// - /api/registration/*→ JWT (Entra, Staff.Admin role only)
/// - /api/staff/*       → JWT (Entra, Staff.Admin role only)
/// - /api/test/*        → Anonymous
/// - /api/documents/*   → JWT (Entra, Staff.Admin or Staff.Tech role)
/// - /api/help/*        → Anonymous
/// 
/// App Role values (defined in Entra portal on the staff app registration):
///   Staff.Admin → full platform access
///   Staff.Tech  → monitoring/health only
/// </summary>
public sealed class ClientAuthMiddleware
{
    private readonly RequestDelegate _next;
    private readonly ILogger<ClientAuthMiddleware> _logger;
    private readonly IConfiguration _config;

    private static readonly HashSet<string> _anonymousExact = new(StringComparer.OrdinalIgnoreCase)
    {
        "/", "/health"
    };

    private static readonly string[] _anonymousPrefixes = { "/swagger", "/api/test", "/api/help" };
    private static readonly string[] _jwtOnlyPrefixes = { "/api/onboarding" };
    private static readonly string[] _staffRequiredPrefixes = { "/api/monitoring", "/api/documents" };  // Admin or Tech
    private static readonly string[] _adminRequiredPrefixes = { "/api/admin", "/api/registration", "/api/staff" }; // Admin only

    private static ConfigurationManager<OpenIdConnectConfiguration>? _oidcConfigManager;
    private static readonly object _oidcLock = new();

    public ClientAuthMiddleware(RequestDelegate next, ILogger<ClientAuthMiddleware> logger, IConfiguration config)
    {
        _next = next;
        _logger = logger;
        _config = config;
    }

    public async Task InvokeAsync(HttpContext context, ClientContext clientContext, SqlService sql)
    {
        var path = (context.Request.Path.Value ?? "").ToLowerInvariant();
        var corrId = EnsureCorrelationId(context);

        // Anonymous paths
        if (IsAnonymousPath(path))
        {
            await _next(context);
            return;
        }

        // Dev bypass
        if (TryDevBypass(context, clientContext))
        {
            await _next(context);
            return;
        }

        // JWT-only paths (onboarding)
        if (IsJwtOnlyPath(path))
        {
            if (await TryJwtAuthAsync(context, clientContext))
            {
                await _next(context);
                return;
            }

            context.Response.StatusCode = 401;
            await context.Response.WriteAsJsonAsync(new { ok = false, error = "Valid Entra authentication required" });
            return;
        }

        // Staff-required paths (Admin or Tech role) — monitoring/health
        // Try session auth (client session tokens for CIAM users).
        // then fall back to direct JWT for service-to-service calls.
        if (IsStaffRequiredPath(path))
        {
            var staffAuthed = await TrySessionAuthAsync(context, clientContext, sql)
                           || await TryJwtAuthAsync(context, clientContext);
            if (staffAuthed)
            {
                if (!clientContext.IsStaff)
                {
                    context.Response.StatusCode = 403;
                    await context.Response.WriteAsJsonAsync(new { ok = false, error = "Staff access required" });
                    return;
                }
                await _next(context);
                return;
            }

            context.Response.StatusCode = 401;
            await context.Response.WriteAsJsonAsync(new { ok = false, error = "Valid staff authentication required" });
            return;
        }

        // Admin-required paths (Admin role only)
        if (IsAdminRequiredPath(path))
        {
            var adminAuthed = await TrySessionAuthAsync(context, clientContext, sql)
                           || await TryJwtAuthAsync(context, clientContext);
            if (adminAuthed)
            {
                if (!clientContext.IsAdmin)
                {
                    context.Response.StatusCode = 403;
                    await context.Response.WriteAsJsonAsync(new { ok = false, error = "Admin access required" });
                    return;
                }
                await _next(context);
                return;
            }

            context.Response.StatusCode = 401;
            await context.Response.WriteAsJsonAsync(new { ok = false, error = "Valid admin authentication required" });
            return;
        }

        // Default: require session
        if (await TrySessionAuthAsync(context, clientContext, sql))
        {
            await _next(context);
            return;
        }

        context.Response.StatusCode = 401;
        await context.Response.WriteAsJsonAsync(new { ok = false, error = "Valid session required" });
    }

    private static bool IsAnonymousPath(string path) =>
        _anonymousExact.Contains(path) || _anonymousPrefixes.Any(p => path.StartsWith(p, StringComparison.OrdinalIgnoreCase));

    private static bool IsJwtOnlyPath(string path) =>
        _jwtOnlyPrefixes.Any(p => path.StartsWith(p, StringComparison.OrdinalIgnoreCase));

    private static bool IsStaffRequiredPath(string path) =>
        _staffRequiredPrefixes.Any(p => path.StartsWith(p, StringComparison.OrdinalIgnoreCase));

    private static bool IsAdminRequiredPath(string path) =>
        _adminRequiredPrefixes.Any(p => path.StartsWith(p, StringComparison.OrdinalIgnoreCase));

    private static string EnsureCorrelationId(HttpContext context)
    {
        if (!context.Request.Headers.TryGetValue("X-Correlation-Id", out var existing) || string.IsNullOrWhiteSpace(existing.FirstOrDefault()))
        {
            var id = Guid.NewGuid().ToString("N");
            context.Request.Headers["X-Correlation-Id"] = id;
            return id;
        }
        return existing.First()!;
    }

    private bool TryDevBypass(HttpContext context, ClientContext clientContext)
    {
        var env = _config["ASPNETCORE_ENVIRONMENT"] ?? Environment.GetEnvironmentVariable("ASPNETCORE_ENVIRONMENT");
        var allowBypass = _config.GetValue<bool>("Auth:AllowDevBypass");

        if (!string.Equals(env, "Development", StringComparison.OrdinalIgnoreCase) && !allowBypass)
            return false;

        if (!context.Request.Headers.TryGetValue("X-Dev-ClientId", out var devClientId))
            return false;

        var clientId = devClientId.FirstOrDefault();
        if (string.IsNullOrWhiteSpace(clientId))
            return false;

        clientContext.ClientId = clientId;
        clientContext.IsDevBypass = true;
        clientContext.Role = "Admin";
        return true;
    }

    private async Task<bool> TrySessionAuthAsync(HttpContext context, ClientContext clientContext, SqlService sql)
    {
        string? token = context.Request.Headers["X-Session-Token"].FirstOrDefault();

        if (string.IsNullOrWhiteSpace(token))
        {
            var authHeader = context.Request.Headers["Authorization"].FirstOrDefault();
            if (!string.IsNullOrWhiteSpace(authHeader) && authHeader.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase))
                token = authHeader["Bearer ".Length..].Trim();
        }

        if (string.IsNullOrWhiteSpace(token))
            return false;

        try
        {
            var rqst        = JsonSerializer.Serialize(new { sessionToken = token });
            var validateProc = "dbo.spClientSession"; // Staff use JWT Bearer; only client sessions exist
            var resp         = await sql.ExecProcAsync(validateProc, "validate", rqst, ct: context.RequestAborted);

            if (string.IsNullOrWhiteSpace(resp))
                return false;

            using var doc = JsonDocument.Parse(resp);
            var root = doc.RootElement;

            if (root.TryGetProperty("ok", out var okProp) && okProp.ValueKind == JsonValueKind.True)
            {
                var data = root.TryGetProperty("data", out var dataProp) ? dataProp : root;

                clientContext.SessionId = data.TryGetProperty("sessionId", out var sid) ? sid.GetString() : null;

                // Admin sessions return adminId; client sessions return clientId — handle both
                var clientId = data.TryGetProperty("clientId", out var cid) ? cid.GetString() : null;
                if (string.IsNullOrWhiteSpace(clientId))
                    clientId = data.TryGetProperty("adminId", out var aid) ? aid.GetString() : null;
                clientContext.ClientId         = clientId;
                clientContext.PlatformClientId = clientId;

                clientContext.ClientName = data.TryGetProperty("clientName", out var cn) ? cn.GetString() : null;
                clientContext.UserId     = data.TryGetProperty("userId",     out var uid) ? uid.GetString() : null;
                clientContext.Email      = data.TryGetProperty("userEmail",  out var ue)  ? ue.GetString()  : null;
                clientContext.Role       = data.TryGetProperty("role",       out var role) ? role.GetString() : null;

                // IsStaff is computed from Role in ClientContext — no assignment needed

                return clientContext.IsAuthenticated;
            }

            return false;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Session validation error");
            return false;
        }
    }

    private async Task<bool> TryJwtAuthAsync(HttpContext context, ClientContext clientContext)
    {
        var authHeader = context.Request.Headers["Authorization"].FirstOrDefault();
        if (string.IsNullOrWhiteSpace(authHeader) || !authHeader.StartsWith("Bearer ", StringComparison.OrdinalIgnoreCase))
            return false;

        var token = authHeader["Bearer ".Length..].Trim();
        if (string.IsNullOrWhiteSpace(token))
            return false;

        var tenantId = _config["Auth:Staff:TenantId"];
        var clientId = _config["Auth:Staff:ClientId"];
        var instance = _config["Auth:Staff:Instance"] ?? "https://login.microsoftonline.com/";

        if (string.IsNullOrWhiteSpace(tenantId) || string.IsNullOrWhiteSpace(clientId))
            return false;

        try
        {
            var handler = new JwtSecurityTokenHandler();
            // Disable default claim type remapping so JWT claim names (roles, oid, etc.)
            // are preserved as-is. Without this, "roles" is remapped to ClaimTypes.Role
            // and FindAll("roles") returns empty — causing IsAdmin to be false.
            handler.InboundClaimTypeMap.Clear();

            var authority       = $"{instance.TrimEnd('/')}/{tenantId}/v2.0";
            var metadataAddress = $"{authority}/.well-known/openid-configuration";

            var mgr          = GetOrCreateConfigManager(metadataAddress);
            var openIdConfig = await mgr.GetConfigurationAsync(context.RequestAborted);

            var validationParams = new TokenValidationParameters
            {
                ValidateIssuer    = true,
                ValidIssuers      = new[]
                {
                    $"https://login.microsoftonline.com/{tenantId}/v2.0",
                    $"https://sts.windows.net/{tenantId}/"
                },
                ValidateAudience  = true,
                ValidAudiences    = new[] { clientId, $"api://{clientId}" },
                ValidateLifetime  = true,
                IssuerSigningKeys = openIdConfig.SigningKeys,
                ClockSkew         = TimeSpan.FromMinutes(5)
            };

            var principal = handler.ValidateToken(token, validationParams, out _);

            // Map Entra App Role values to internal role names.
            // Users with no recognized role get null — middleware rejects them.
            var roles = principal.FindAll("roles").Select(c => c.Value).ToList();
            var role  = roles.Contains("Staff.Admin") ? "Admin"
                      : roles.Contains("Staff.Tech")  ? "Tech"
                      : null; // no valid role assigned — reject

            clientContext.ClientId   = principal.FindFirstValue("oid")
                                    ?? principal.FindFirstValue(ClaimTypes.NameIdentifier);
            clientContext.Email      = principal.FindFirstValue("preferred_username")
                                    ?? principal.FindFirstValue(ClaimTypes.Email);
            clientContext.ClientName = principal.FindFirstValue("name")
                                    ?? principal.FindFirstValue(ClaimTypes.Name);
            clientContext.Role       = role;

            return clientContext.IsAuthenticated;
        }
        catch (Exception ex)
        {
            _logger.LogWarning("JWT validation failed: {Message}", ex.Message);
            return false;
        }
    }

    private static ConfigurationManager<OpenIdConnectConfiguration> GetOrCreateConfigManager(string metadataAddress)
    {
        lock (_oidcLock)
        {
            _oidcConfigManager ??= new ConfigurationManager<OpenIdConnectConfiguration>(
                metadataAddress, new OpenIdConnectConfigurationRetriever());
            return _oidcConfigManager;
        }
    }
}