Simplified Role based Authorization for HotChocolate GraphQL v13 in Azure Functions

AzFuncGraphQLRoleAuthorizationFramework.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Reflection;
using System.Threading.Tasks;
using HotChocolate;
using HotChocolate.Resolvers;
using HotChocolate.Data.Projections.Context;
using HotChocolate.Execution;
using HotChocolate.AspNetCore.Serialization;
using HotChocolate.Execution.Configuration;
using Microsoft.Extensions.DependencyInjection;
using System.Security.Claims;
using HotChocolate.Types;
using HotChocolate.Execution.Instrumentation;
using Microsoft.AspNetCore.Authentication;
using Microsoft.AspNetCore.Http;

namespace GraphQL.HotChocolate.AzureFunctions.Authorization
{
    ///This Framework was taken directly from the GitHub Gist located Here:
    /// https://gist.github.com/cajuncoding/4c037cb881a31da4aada061fde7f95f7
    /// 
    /// <summary>
    /// BBernard / CajunCoding
    /// A simplified framework for Role based authorization of HotChocolate GraphQL in Azure Functions.
    ///
    /// To implement:
    /// 1) Call the builder.AddAzFuncRoleBasedAuthorization() extension to configure the framework in the Startup process.
    ///    1a) The claims principal is automatically retrieved from the IHttpContextAccessor so you do also need to register that
    ///         in you DI via services.AddHttpContextAccessor().
    /// 2) Add [AuthorizeAzFuncGraphQL("Role1")] attributes to any Resolver or Mutation you wish to secure.
    /// 3) Ensure that your User ClaimsPrincipal has the matching roles specified as Claims (e.g. ClaimTypes.Role).
    ///     - When using Azure Function App Keys (aka Tokens) you can read the App Key Name from the default Claims Principal
    ///         easily with the principal.GetCurrentAzFuncAppTokenKeyName() custom extension provided.
    ///     - You may then map that to a Role that may be configuration based, etc.
    ///     - The Claims of the current User Principal may be altered/updated to include the Role via a custom HotChocolate HttpRequestInterceptor.
    ///     - For More info see: https://chillicream.com/docs/hotchocolate/v13/security/authorization/#modifying-the-claimsprincipal
    /// 
    /// Additional Info:
    ///     Unfortunately as of HC v13, we are unable to use the AspNetCore AddAuthorization() extensions
    ///         within Azure Functions because it breaks the Azure Function execution due to replacing/overriding
    ///         the implementations needed by the Functions Host! And HC's default AddAuthorization() has sealed
    ///         internal classes and forces the use of the default AddAuthorization() extension so there is no workaround.
    ///     Therefore a different (and completely separated) approach is needed currently. So this simplified
    ///         framework, while certainly less robust, solves the gap by providing support for Role based authorization
    ///         via custom Attribute combined with Field Middleware and ResponseFormatter (to correctly set the HttpStatusCodes).
    /// </summary>
    [AttributeUsage(AttributeTargets.Method | AttributeTargets.Property | AttributeTargets.Class)]
    public class AuthorizeAzFuncGraphQLAttribute : Attribute
    {
        public string[] Roles { get; set; }

        public AuthorizeAzFuncGraphQLAttribute(params string[] roles)
        {
            Roles = roles;
        }
    }

    public class AzFuncGraphQLRoleAuthorizationEventListener : ExecutionDiagnosticEventListener
    {
        private readonly IServiceProvider _serviceProvider;

        public AzFuncGraphQLRoleAuthorizationEventListener(IServiceProvider serviceProvider)
        {
            _serviceProvider = serviceProvider;
        }

        /// <summary>
        /// The ExecuteOperation will run when HC is starting to execute the GraphQL operation and its resolvers.
        /// So the GraphQL Query Document will be fully parsed and available at this time for Logging.
        /// </summary>
        public override IDisposable ExecuteRequest(IRequestContext context)
        {
            var httpContext = _serviceProvider.GetRequiredService<IHttpContextAccessor>()?.HttpContext;
            if (httpContext != null)
            {
                var claimsTransformer = _serviceProvider.GetService<IClaimsTransformation>();
                if (claimsTransformer != null)
                {
                    var task = Task.Run(async () => await claimsTransformer.TransformAsync(httpContext.User).ConfigureAwait(false));
                    httpContext.User = task.GetAwaiter().GetResult();
                    context.ContextData[nameof(ClaimsPrincipal)] = httpContext.User;
                }
            }

            return base.ExecuteRequest(context);
        }
    }

    public class AzFuncGraphQLRoleAuthorizationResponseFormatter : DefaultHttpResponseFormatter
    {
        protected override HttpStatusCode OnDetermineStatusCode(IQueryResult queryResult, FormatInfo format, HttpStatusCode? proposedStatusCode)
        {
            if ((queryResult?.Errors?.Count ?? 0) <= 0) 
                return base.OnDetermineStatusCode(queryResult, format, proposedStatusCode);
            
            var authError = queryResult.Errors.FirstOrDefault(e => e.Code is ErrorCodes.Authentication.NotAuthenticated or ErrorCodes.Authentication.NotAuthorized);
            return authError?.Code switch
            {
                ErrorCodes.Authentication.NotAuthenticated => HttpStatusCode.Unauthorized,
                ErrorCodes.Authentication.NotAuthorized => HttpStatusCode.Forbidden,
                _ => base.OnDetermineStatusCode(queryResult, format, proposedStatusCode)
            };
        }
    }

    public class AzFuncGraphQLRoleAuthorizationFieldMiddleware
    {
        public const string AzFuncGraphQLRoleAuthorizationContextCacheKey = "AzFuncGraphQLRoleAuthorizationFieldMiddleware.AuthorizationCache";
        private static readonly ConcurrentDictionary<MemberInfo, Lazy<AuthorizeAzFuncGraphQLAttribute>> ResolverAuthAttrRegistry = new();
        private readonly FieldDelegate _next;

        public AzFuncGraphQLRoleAuthorizationFieldMiddleware(FieldDelegate next) { _next = next; }

        public async Task InvokeAsync(IMiddlewareContext context)
        {
            var field = context.GetSelectedField().Field;

            //NOTE: For Performance we only check Authorization if an AzFuncAuthorizationAttribute exists,
            //  and then if it does we leverage caching (static cache of Attributes & context cache of authorization results)
            //  for every call to the same Field resolver....
            if(TryGetResolverAuthorizeAttribute(field, out var azFuncAuthAttribute) 
               && !IsUserAuthorizedForField(context, field, azFuncAuthAttribute))
                    throw new GraphQLException(ErrorBuilder.New()
                        .SetMessage($"The current Azure Function App Key/Token does not have access to the field [{field.Name}] of the GraphQL request.")
                        .SetCode(ErrorCodes.Authentication.NotAuthorized)
                        .Build()
                    );

            await _next(context).ConfigureAwait(false);
        }

        private static bool IsUserAuthorizedForField(IResolverContext context, IObjectField field, AuthorizeAzFuncGraphQLAttribute azFuncAuthAttribute)
        {
            var contextData = context.ContextData;
            if (!(contextData.TryGetValue(AzFuncGraphQLRoleAuthorizationContextCacheKey, out var obj) && obj is Dictionary<MemberInfo, bool> authorizationCache))
                contextData[AzFuncGraphQLRoleAuthorizationContextCacheKey] = authorizationCache = new();

            if (field.ResolverMember is not { } resolverMember)
                return true;

            if(authorizationCache.TryGetValue(resolverMember, out bool isAuthorized))
               return isAuthorized;

            var roles = azFuncAuthAttribute.Roles;
            var user = context.GetUser();
            isAuthorized = roles != null && user != null && roles.Any(r => user.IsInRole(r));

            authorizationCache[resolverMember] = isAuthorized;
            return isAuthorized;
        }

        private static bool TryGetResolverAuthorizeAttribute(IObjectField field, out AuthorizeAzFuncGraphQLAttribute authorizeAttr)
        {
            authorizeAttr = null;
            if (field.ResolverMember is not { } resolverMemberInfo)
                return false;

            var doesResolverNeedParamsContextLazy = ResolverAuthAttrRegistry.GetOrAdd(
                resolverMemberInfo,
                new Lazy<AuthorizeAzFuncGraphQLAttribute>(() => 
                    resolverMemberInfo.GetCustomAttribute<AuthorizeAzFuncGraphQLAttribute>()
                        ?? resolverMemberInfo.DeclaringType?.GetCustomAttribute<AuthorizeAzFuncGraphQLAttribute>()
                )
            );

            return (authorizeAttr = doesResolverNeedParamsContextLazy.Value) != null;
        }
    }

    public static class AzFuncGraphQLRoleBasedAuthorizationExtensions
    {
        /// <summary>
        /// Adds a simplified Role based Authorization support for Azure Functions. This will automatically register an HttpResponseFormatter,
        ///     Authorization Field Middleware (to enforce the Authorization Security) and a GraphQL Authorization Diagnostic Event Listener
        ///     (that provides built in support for normal IClaimsTransformer compatible with default AspNet Core Authorization documentation
        ///     for initializing and mutating the Claims of the User Principal).
        /// NOTE: This is needed as of HC v13 because the standard AspNetCore AddAuthorization() does not work with Azure Functions.
        /// uses Microsoft.AspNetCore.Authorization.
        /// </summary>
        /// <param name="builder">The <see cref="IRequestExecutorBuilder"/>.</param>
        /// <returns>Returns the <see cref="IRequestExecutorBuilder"/> for chaining in more configurations.</returns>
        public static IRequestExecutorBuilder AddAzFuncRoleBasedAuthorization(this IRequestExecutorBuilder builder)
        {
            if (builder == null)
                throw new ArgumentNullException(nameof(builder));

            builder.Services.AddHttpResponseFormatter<AzFuncGraphQLRoleAuthorizationResponseFormatter>();
            builder.AddDiagnosticEventListener<AzFuncGraphQLRoleAuthorizationEventListener>();
            builder.UseField<AzFuncGraphQLRoleAuthorizationFieldMiddleware>();
            return builder;
        }

        /// <summary>
        /// Read the current Azure Function App Key Name (the name of the token as configured in azure portal).
        /// </summary>
        /// <param name="userPrincipal"></param>
        /// <returns></returns>
        public static string GetCurrentAzFuncAppTokenKeyName(this ClaimsPrincipal userPrincipal)
        {
            const string AzFuncKeyClaimUri = "http://schemas.microsoft.com/2017/07/functions/claims/keyid";
            var claim = userPrincipal?.Claims.FirstOrDefault(c => c.Type == AzFuncKeyClaimUri);
            return claim?.Value;
        }

        /// <summary>
        /// Safely get the Name of the current User Principal.
        /// </summary>
        /// <param name="userPrincipal"></param>
        /// <returns></returns>
        public static string GetName(this ClaimsPrincipal userPrincipal)
            => userPrincipal?.Identities?.FirstOrDefault(i => !string.IsNullOrWhiteSpace(i.Name))?.Name;

        /// <summary>
        /// Safely get the Name of the current User Principal.
        /// </summary>
        /// <param name="userPrincipal"></param>
        /// <returns></returns>
        public static string[] GetRoles(this ClaimsPrincipal userPrincipal)
            => userPrincipal?.Identities?.SelectMany(i => i.Claims?.Where(c => c.Type == i.RoleClaimType).Select(c => c.Value)).ToArray();
    }
}