ShrinkNetwork/Runtime/Routing/ShrinkNetworkRouter.cs

#nullable enable

using System;
using System.Collections.Generic;
using Cysharp.Threading.Tasks;

namespace ShrinkNetwork
{
    public sealed class ShrinkNetworkRouter
    {
        private sealed class MessageHandlerRegistration
        {
            public ShrinkNetworkPermissionRequirement Requirement;
            public Func<ShrinkNetworkContext, object, UniTask> Handler = null!;
        }

        private sealed class RequestHandlerRegistration
        {
            public Type ResponseType = null!;
            public ShrinkNetworkPermissionRequirement Requirement;
            public Func<ShrinkNetworkContext, object, UniTask<object?>> Handler = null!;
        }

        private readonly Dictionary<Type, MessageHandlerRegistration> _messageHandlers = new();
        private readonly Dictionary<Type, RequestHandlerRegistration> _requestHandlers = new();

        public void RegisterHandler<TMessage>(Func<ShrinkNetworkContext, TMessage, UniTask> handler,
            ShrinkNetworkPermissionRequirement requirement = default)
            where TMessage : IShrinkNetworkMessage
        {
            RegisterHandler(typeof(TMessage), (context, message) => handler(context, (TMessage)message), requirement);
        }

        public void RegisterHandler(Type messageType, Func<ShrinkNetworkContext, object, UniTask> handler,
            ShrinkNetworkPermissionRequirement requirement = default)
        {
            if (messageType == null)
                throw new ArgumentNullException(nameof(messageType));
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));
            if (_messageHandlers.ContainsKey(messageType) || _requestHandlers.ContainsKey(messageType))
                throw new InvalidOperationException($"Handler already exists for {messageType.FullName}.");

            _messageHandlers.Add(messageType, new MessageHandlerRegistration
            {
                Requirement = requirement,
                Handler = handler
            });
        }

        public void RegisterRequestHandler<TRequest, TResponse>(Func<ShrinkNetworkContext, TRequest, UniTask<TResponse>> handler,
            ShrinkNetworkPermissionRequirement requirement = default)
            where TRequest : IShrinkNetworkRequest
            where TResponse : class, IShrinkNetworkResponse
        {
            RegisterRequestHandler(typeof(TRequest), typeof(TResponse),
                async (context, message) => await handler(context, (TRequest)message), requirement);
        }

        public void RegisterRequestHandler(Type requestType, Type responseType,
            Func<ShrinkNetworkContext, object, UniTask<object?>> handler,
            ShrinkNetworkPermissionRequirement requirement = default)
        {
            if (requestType == null)
                throw new ArgumentNullException(nameof(requestType));
            if (responseType == null)
                throw new ArgumentNullException(nameof(responseType));
            if (handler == null)
                throw new ArgumentNullException(nameof(handler));
            if (_messageHandlers.ContainsKey(requestType) || _requestHandlers.ContainsKey(requestType))
                throw new InvalidOperationException($"Handler already exists for {requestType.FullName}.");

            _requestHandlers.Add(requestType, new RequestHandlerRegistration
            {
                ResponseType = responseType,
                Requirement = requirement,
                Handler = handler
            });
        }

        public async UniTask<bool> DispatchAsync(ShrinkNetworkContext context, object message, Type messageType)
        {
            if (context.Packet.Kind == ShrinkNetworkPacketKind.Request &&
                _requestHandlers.TryGetValue(messageType, out var requestHandler))
            {
                if (!ShrinkNetworkPermissionValidator.IsAllowed(context.Session, requestHandler.Requirement))
                {
                    context.Service.ReportPermissionDenied(messageType);
                    var denied = CreatePermissionDeniedResponse(requestHandler.ResponseType, requestHandler.Requirement);
                    await context.Service.SendResponseAsync(context.Session, denied, requestHandler.ResponseType,
                        context.Packet.RequestToken, context.Packet.Route);
                    return true;
                }

                try
                {
                    var response = await requestHandler.Handler(context, message);
                    if (response is IShrinkNetworkResponse networkResponse)
                        await context.Service.SendResponseAsync(context.Session, networkResponse, requestHandler.ResponseType,
                            context.Packet.RequestToken, context.Packet.Route);
                }
                catch (Exception ex)
                {
                    context.Service.ReportHandlerException(messageType, ex);
                    var errorResponse = CreateErrorResponse(requestHandler.ResponseType, ex);
                    await context.Service.SendResponseAsync(context.Session, errorResponse, requestHandler.ResponseType,
                        context.Packet.RequestToken, context.Packet.Route);
                }

                return true;
            }

            if (_messageHandlers.TryGetValue(messageType, out var messageHandler))
            {
                if (!ShrinkNetworkPermissionValidator.IsAllowed(context.Session, messageHandler.Requirement))
                {
                    context.Service.ReportPermissionDenied(messageType);
                    ShrinkNetworkLogger.Warn($"[ShrinkNetwork] Permission denied for message {messageType.FullName} on session {context.Session.SessionId}.");
                    return true;
                }

                await messageHandler.Handler(context, message);
                return true;
            }

            return false;
        }

        private static IShrinkNetworkResponse CreateErrorResponse(Type responseType, Exception ex)
        {
            if (Activator.CreateInstance(responseType) is not IShrinkNetworkResponse response)
                throw new InvalidOperationException($"Response type {responseType.FullName} cannot be instantiated.", ex);

            response.ErrorCode = ShrinkRpcErrorCode.HandlerException;
            response.ErrorMessage = ex.Message;
            return response;
        }

        private static IShrinkNetworkResponse CreatePermissionDeniedResponse(Type responseType,
            ShrinkNetworkPermissionRequirement requirement)
        {
            if (Activator.CreateInstance(responseType) is not IShrinkNetworkResponse response)
                throw new InvalidOperationException($"Response type {responseType.FullName} cannot be instantiated.");

            response.ErrorCode = ShrinkRpcErrorCode.PermissionDenied;
            response.ErrorMessage = string.IsNullOrEmpty(requirement.Permission)
                ? $"Permission denied. Authority={requirement.Authority}"
                : $"Permission denied. Authority={requirement.Authority}, Permission={requirement.Permission}";
            return response;
        }
    }
}