Connected.Framework/Connected.Services/Authorization/ServiceAuthorizationContext.cs

using Connected.Interop;
using Connected.Middleware;
using Connected.Security.Authorization;
using Connected.ServiceModel;
using Connected.Services.Annotations;
using Connected.Threading;
using System.Collections.Immutable;
using System.Reflection;

namespace Connected.Services.Authorization;

internal sealed class ServiceAuthorizationContext : IAuthorizationContext
{
	public ServiceAuthorizationContext(IAuthorizationService authorization, IMiddlewareService middleware)
	{
		Authorization = authorization;

		Middleware = new AsyncLazy<ImmutableList<IServiceAuthorizationMiddleware>>(middleware.Query<IServiceAuthorizationMiddleware>());
	}

	private IAuthorizationService Authorization { get; }
	private AsyncLazy<ImmutableList<IServiceAuthorizationMiddleware>> Middleware { get; }
	public AuthorizationContextState State { get; private set; } = AuthorizationContextState.Pending;

	public async Task<IAuthorizationResult> Authorize(AuthorizationArgs args)
	{
		/*
		 * Authorization is performed only if the context does not have a grant permission
		 * for the execution. That means only the first authorization request is actually
		 * performed. Once the caller has been granted access no authorization takes place
		 * afterwards.
		 */
		if (State == AuthorizationContextState.Granted)
			return AuthorizationResult.OK();

		var result = await Authorization.Authorize(args);
		/*
		 * Authorization completed. Set correct state which means no 
		 * authorization is needed for any subsequent calls.
		 */
		if (result.Success)
			State = AuthorizationContextState.Granted;
		else
			State = AuthorizationContextState.Revoked;

		return result;
	}

	public async Task Authorize<TArgs>(ICallerContext context, TArgs args) where TArgs : IDto
	{
		/*
		 * We allow only one authorizaton process in the same context at the same time.
		 * We would cause a potential stack overflow otherwise.
		 */
		if (State == AuthorizationContextState.Authorizing)
			return;

		if (State == AuthorizationContextState.Granted)
			return;

		State = AuthorizationContextState.Authorizing;

		try
		{
			if (context.Sender is null)
				return;

			if (ResolveAttributes(context, args) is not List<ServiceAuthorizationAttribute> attributes || !attributes.Any())
				return;

			var staging = attributes.Where(f => f.Stage == AuthorizationStage.Init);

			if (!staging.Any())
				return;

			Exception? firstFail = null;
			bool onePassed = false;

			foreach (var attribute in staging)
			{
				try
				{
					if (attribute.Behavior == AuthorizationPolicyBehavior.Optional && onePassed)
						continue;

					var claims = attribute.Claims.ToImmutableArray();

					foreach (var middleware in await Middleware.Value)
						claims = await middleware.ResolveClaims(claims);

					var middlewareArgs = new ServiceAuthorizationMiddlewareArgs<TArgs>(context, claims, args);

					foreach (var middleware in await Middleware.Value)
						await middleware.Authorize(middlewareArgs);

					onePassed = true;
				}
				catch (Exception ex)
				{
					if (attribute.Behavior == AuthorizationPolicyBehavior.Mandatory)
						throw;

					firstFail = ex;
				}
			}

			if (!onePassed && firstFail is not null)
				throw firstFail;

			State = AuthorizationContextState.Granted;
		}
		catch
		{
			State = AuthorizationContextState.Revoked;

			throw;
		}
	}

	/// <summary>
	/// Results are always authorized, event if <see cref="State"/> is <see cref="AuthorizationContextState.Granted"/>.
	/// </summary>
	/// <typeparam name="TArgs"></typeparam>
	/// <typeparam name="TComponent"></typeparam>
	/// <param name="context"></param>
	/// <param name="args"></param>
	/// <param name="component"></param>
	/// <returns></returns>
	public async Task<TComponent> Authorize<TArgs, TComponent>(ICallerContext context, TArgs args, TComponent component)
		where TArgs : IDto
	{
		if (context.Sender is null)
			return component;

		if (ResolveAttributes(context, args) is not List<ServiceAuthorizationAttribute> attributes || !attributes.Any())
			return component;

		var staging = attributes.Where(f => f.Stage == AuthorizationStage.Result);

		if (!staging.Any())
			return component;

		Exception? firstFail = null;
		bool onePassed = false;
		var result = component;

		foreach (var attribute in staging)
		{
			try
			{
				if (attribute.Behavior == AuthorizationPolicyBehavior.Optional && onePassed)
					continue;

				var claims = attribute.Claims.ToImmutableArray();

				foreach (var middleware in await Middleware.Value)
					claims = await middleware.ResolveClaims(claims);

				var middlewareArgs = new ServiceAuthorizationMiddlewareArgs<TArgs>(context, claims, args);

				foreach (var middleware in await Middleware.Value)
					result = await middleware.Authorize(middlewareArgs, result);

				onePassed = true;
			}
			catch (Exception ex)
			{
				if (attribute.Behavior == AuthorizationPolicyBehavior.Mandatory)
					throw;

				firstFail = ex;
			}
		}

		if (!onePassed && firstFail is not null)
			throw firstFail;

		return result;
	}

	public void Revoke()
	{
		if (State == AuthorizationContextState.Granted)
			State = AuthorizationContextState.Revoked;
	}

	private static List<ServiceAuthorizationAttribute>? ResolveAttributes<TArgs>(ICallerContext context, TArgs args) where TArgs : IDto
	{
		if (context.Sender is null || string.IsNullOrEmpty(context.Method))
			return default;

		if (context.Sender.GetType().ResolveMethod(context.Method, null, new Type[1] { args.GetType() }) is not MethodInfo method)
			throw new InvalidOperationException($"{SR.ErrMethodResolve} ({context.Sender.GetType().Name}.{context.Method})");

		var attributes = method.GetCustomAttributes(typeof(ServiceAuthorizationAttribute), false);

		if (attributes is null || !attributes.Any())
			return null;

		var result = new List<ServiceAuthorizationAttribute>();

		foreach (var attribute in attributes)
			result.Add((ServiceAuthorizationAttribute)attribute);

		result.Sort((left, right) =>
		{
			if (left.Priority < right.Priority)
				return -1;
			else if (left.Priority == right.Priority)
				return 0;
			else
				return 1;
		});

		return result;
	}
}