Authorization

[AttributeUsage(AttributeTargets.Method, AllowMultiple = true)]
public class RequireRoleAttribute : Attribute
{
    public string Role { get; }
    public RequireRoleAttribute(string role) => Role = role;
}

[AttributeUsage(AttributeTargets.Method)]
public class AllowAnonymousAttribute : Attribute
{
}

using Mcp.Gateway.Tools.Lifecycle;

public class AuthorizationHook : IToolLifecycleHook
{
    private readonly IHttpContextAccessor _httpContextAccessor;
    private readonly Dictionary<string, MethodInfo> _toolMethods = new();
    
    public Task OnToolInvokingAsync(string toolName, JsonRpcMessage request)
    {
        var httpContext = _httpContextAccessor.HttpContext;
        if (httpContext == null) return Task.CompletedTask;
        
        var method = GetToolMethod(toolName);
        if (method == null) return Task.CompletedTask;
        
        // Check [AllowAnonymous]
        if (method.GetCustomAttribute<AllowAnonymousAttribute>() != null)
            return Task.CompletedTask;
        
        // Get required roles
        var requiredRoles = method.GetCustomAttributes<RequireRoleAttribute>()
            .Select(attr => attr.Role)
            .ToList();
        
        if (!requiredRoles.Any()) return Task.CompletedTask;
        
        // Get user roles
        var userRoles = httpContext.Items["UserRoles"] as List<string> ?? new();
        
        // Check authorization
        if (!requiredRoles.Any(role => userRoles.Contains(role)))
        {
            throw new ToolInvalidParamsException(
                $"Insufficient permissions. Required: {string.Join(" or ", requiredRoles)}",
                toolName);
        }
        
        return Task.CompletedTask;
    }
    
    // ... other interface methods
}

var builder = WebApplication.CreateBuilder(args);

builder.Services.AddHttpContextAccessor();
builder.AddToolsService();
builder.AddToolLifecycleHook<AuthorizationHook>();

var app = builder.Build();

// Authentication middleware
app.Use(async (context, next) =>
{
    if (!context.Request.Headers.TryGetValue("Authorization", out var authHeader))
    {
        context.Response.StatusCode = 401;
        return;
    }
    
    var token = authHeader.ToString().Replace("Bearer ", "");
    var validation = await ValidateTokenAsync(token);
    
    if (!validation.IsValid)
    {
        context.Response.StatusCode = 403;
        return;
    }
    
    // Store user info for AuthorizationHook
    context.Items["UserId"] = validation.UserId;
    context.Items["UserRoles"] = validation.Roles;
    
    await next();
});

app.UseWebSockets();
app.MapStreamableHttpEndpoint("/mcp");
app.Run();

public class AdminTools
{
    [McpTool("delete_user")]
    [RequireRole("Admin")]
    public JsonRpcMessage DeleteUser(JsonRpcMessage request)
    {
        // Admin-only logic
        return ToolResponse.Success(request.Id, new { deleted = true });
    }
    
    [McpTool("create_user")]
    [RequireRole("Admin")]
    [RequireRole("Manager")]  // Admin OR Manager
    public JsonRpcMessage CreateUser(JsonRpcMessage request)
    {
        return ToolResponse.Success(request.Id, new { created = true });
    }
    
    [McpTool("get_public_info")]
    [AllowAnonymous]
    public JsonRpcMessage GetPublicInfo(JsonRpcMessage request)
    {
        return ToolResponse.Success(request.Id, new { info = "Public" });
    }
}

using Microsoft.AspNetCore.Authentication.JwtBearer;

builder.Services.AddAuthentication("Bearer")
    .AddJwtBearer("Bearer", options =>
    {
        options.Authority = "https://your-identity-server.com";
        options.Audience = "mcp-gateway-api";
    });

builder.Services.AddAuthorization();

app.UseAuthentication();
app.UseAuthorization();

// Extract claims to HttpContext.Items
app.Use(async (context, next) =>
{
    if (context.User.Identity?.IsAuthenticated == true)
    {
        context.Items["UserId"] = context.User.FindFirst("sub")?.Value;
        context.Items["UserRoles"] = context.User.FindAll(ClaimTypes.Role)
            .Select(c => c.Value)
            .ToList();
    }
    await next();
});

public class AuditLoggingHook : IToolLifecycleHook
{
    private readonly IAuditLogger _auditLogger;
    
    public async Task OnToolInvokingAsync(string toolName, JsonRpcMessage request)
    {
        await _auditLogger.LogAccessAttemptAsync(new AuditLogEntry
        {
            Timestamp = DateTime.UtcNow,
            ToolName = toolName,
            UserId = GetUserId(),
            Action = "ToolInvoking"
        });
    }
    
    // ... other methods
}

// Register both hooks
builder.AddToolLifecycleHook<AuthorizationHook>();
builder.AddToolLifecycleHook<AuditLoggingHook>();

[AttributeUsage(AttributeTargets.Method)]
public class RequireTenantAttribute : Attribute
{
    public string[] AllowedTenants { get; }
    public RequireTenantAttribute(params string[] tenants) => AllowedTenants = tenants;
}

// Usage:
[McpTool("tenant_tool")]
[RequireTenant("tenant-a", "tenant-b")]
[RequireRole("Admin")]
public JsonRpcMessage TenantTool(JsonRpcMessage request) { /* ... */ }

// ❌ BAD: Trusting client-provided role
var clientRole = request.GetParams().GetProperty("role").GetString();

// ✅ GOOD: Only trust server-validated token
var roles = httpContext.Items["UserRoles"] as List<string>;

var validationParameters = new TokenValidationParameters
{
    ValidateIssuerSigningKey = true,
    IssuerSigningKey = GetSecurityKey(),
    ValidateIssuer = true,
    ValidateAudience = true,
    ValidateLifetime = true,
    ClockSkew = TimeSpan.Zero  // No tolerance for expired tokens
};

_logger.LogWarning(
    "Authorization failed for user '{UserId}' attempting '{ToolName}'. " +
    "Required: {Required}, Has: {Has}",
    userId, toolName, requiredRoles, userRoles);

builder.Services.AddRateLimiter(options =>
{
    options.GlobalLimiter = PartitionedRateLimiter.Create<HttpContext, string>(
        context => RateLimitPartition.GetFixedWindowLimiter(
            partitionKey: context.Items["UserId"]?.ToString() ?? "anonymous",
            factory: _ => new FixedWindowRateLimiterOptions
            {
                PermitLimit = 100,
                Window = TimeSpan.FromMinutes(1)
            }));
});

app.UseRateLimiter();

[Fact]
public async Task OnToolInvokingAsync_WithAdminRole_AllowsAccess()
{
    var httpContextAccessor = CreateMockAccessor(
        userId: "admin-user",
        roles: new List<string> { "Admin" });
    
    var hook = new AuthorizationHook(httpContextAccessor);
    var request = JsonRpcMessage.CreateRequest("delete_user", "1");
    
    // Should not throw
    await hook.OnToolInvokingAsync("delete_user", request);
}

[Fact]
public async Task AdminTool_WithUserToken_ReturnsUnauthorized()
{
    var client = factory.CreateClient();
    var httpRequest = new HttpRequestMessage(HttpMethod.Post, "/rpc")
    {
        Content = JsonContent.Create(new
        {
            jsonrpc = "2.0",
            method = "tools/call",
            @params = new { name = "delete_user" }
        })
    };
    httpRequest.Headers.Authorization = 
        new AuthenticationHeaderValue("Bearer", "user-token");
    
    var response = await client.SendAsync(httpRequest);
    var content = await response.Content.ReadAsStringAsync();
    
    Assert.Contains("Insufficient permissions", content);
}