Function metadata available in middleware/FunctionContext about azure-functions-dotnet-worker HOT 3 OPEN

file static class FunctionContextExtensions
{
    public static MethodInfo GetTargetFunctionMethod(this FunctionContext context)
    {
        var entryPoint = context.FunctionDefinition.EntryPoint;
        var assemblyPath = context.FunctionDefinition.PathToAssembly;

        var assembly = Assembly.LoadFrom(assemblyPath);
        var typeName = entryPoint[..entryPoint.LastIndexOf('.')];
        var type = assembly.GetType(typeName);
        var methodName = entryPoint[(entryPoint.LastIndexOf('.') + 1)..];
        var method = type?.GetMethod(methodName);

        return method ?? throw new Exception("Could not get target function from the context");
    }
}

public class FunctionAuthRules
{
    private readonly HashSet<string> _anonymousFunctions = new(){"RenderSwaggerUI","RenderSwaggerDocument"};
    private readonly Dictionary<string, string> _scopedFunctions = new();

    public FunctionAuthRules() : this(GetAppDomainMethods())
    {
    }

    protected FunctionAuthRules(IEnumerable<MethodInfo> methods)
    {
        foreach (var method in methods)
        {
            var functionName = method.GetCustomAttribute<FunctionAttribute>()?.Name;
     
            if (functionName is null)
            {
                continue;
            }

            if (method.GetCustomAttribute<AllowAnonymousAccessAttribute>() is not null)
            {
                _anonymousFunctions.Add(functionName);
            }

            if (method.GetCustomAttribute<RequireScopeAttribute>() is { } scoped)
            {
                _scopedFunctions.Add(functionName, scoped.RequiredScope); 
            }
        }
    }

    public bool AllowAnonymousAccess(FunctionContext context)
    {
        return _anonymousFunctions.Contains(context.FunctionDefinition.Name);
    }

    public string? RequiredScope(FunctionContext context)
    {
        var hasScope = _scopedFunctions.TryGetValue(context.FunctionDefinition.Name, out var scope);

        return hasScope ? scope : null;
    }

    private static IEnumerable<MethodInfo> GetAppDomainMethods()
    {
        return AppDomain.CurrentDomain.GetAssemblies()
            .SelectMany(a => a.GetTypes())
            .SelectMany(t => t.GetMethods());
    }
}

[AttributeUsage(AttributeTargets.Method, AllowMultiple = true)]
public class RequireScopeAttribute : Attribute
{
    public string RequiredScope { get; }

    public RequireScopeAttribute(string requiredScope)
    {
        RequiredScope = requiredScope;
    }
}

public class JwtBearerAuthenticationMiddleware : IFunctionsWorkerMiddleware
{
    private readonly IJwtValidationService _jwtValidationService;
    private readonly FunctionAuthRules _rules;
    private readonly ILogger _logger;

    public JwtBearerAuthenticationMiddleware(IJwtValidationService jwtValidationService, FunctionAuthRules rules, ILogger logger)
    {
        _jwtValidationService = jwtValidationService;
        _rules = rules;
        _logger = logger;
    }

    public async Task Invoke(FunctionContext context, FunctionExecutionDelegate next)
    {
        var requestData = await context.GetHttpRequestDataAsync();

        if (requestData is null)
        {
            throw new InvalidOperationException("Unable to get request data");
        }

        if (_rules.AllowAnonymousAccess(context))
        {
            await next(context);

            return;
        }

        if (!requestData.TryGetJwt(out var jwt))
        {
            context.GetInvocationResult().Value = requestData.CreateResponse(HttpStatusCode.Unauthorized);
            _logger.Warning("Authorization header not found");

            return;
        }

        var jwtValidationResponse = await _jwtValidationService.ValidateJwt(jwt!);

        if (!jwtValidationResponse.IsValid)
        {
            context.GetInvocationResult().Value = requestData.CreateResponse(HttpStatusCode.Unauthorized);
            _logger.Warning("Unauthorized: {Request_Uri}", requestData.Url);

            return;
        }

        context.SetApiClaimsPrincipal(new ApiClaimsPrincipal(jwtValidationResponse.ToUserPrincipal()));

        await next(context);
    }
}

public class RequireScopeMiddleware : IFunctionsWorkerMiddleware
{
    private readonly FunctionAuthRules _rules;
    private readonly ILogger _logger;

    public RequireScopeMiddleware(FunctionAuthRules rules, ILogger logger)
    {
        _rules = rules;
        _logger = logger;
    }

    public async Task Invoke(FunctionContext context, FunctionExecutionDelegate next)
    {
        var requestData = await context.GetHttpRequestDataAsync();

        if (requestData is null)
        {
            throw new InvalidOperationException("Unable to get request data");
        }

        if (_rules.RequiredScope(context) is { } requiredScope)
        {
            var principal = context.GetApiClaimsPrincipal();

            if (principal is null)
            {
                context.GetInvocationResult().Value = requestData.CreateResponse(HttpStatusCode.Unauthorized);
                _logger.Warning("Unauthorized: {Request_Uri}", requestData.Url);

                return;
            }

            if (!principal.HasScope(requiredScope))
            {
                context.GetInvocationResult().Value = requestData.CreateResponse(HttpStatusCode.Forbidden);
                _logger
                    .ForContext(nameof(principal.ClientId), principal.ClientId)
                    .ForContext(nameof(principal.Scopes), principal.Scopes)
                    .ForContext("RequiredScope", requiredScope)
                    .Warning("Forbidden: {Request_Uri}", requestData.Url);

                return;
            }
        }

        await next(context);
    }
}