MoShiro/backend/fleet/src/Fleet.Rpc/RpcRouter.cs

using System.Text.Json;
using System.Text.Json.Serialization.Metadata;
using Fleet.Core;

namespace Fleet.Rpc;

/// <summary>对端违反了协议：方法没有声明、方向不对、请求和通知用反了、或者角色不允许。</summary>
public readonly record struct RpcViolation(string Method, string Reason);

/// <summary>一次分发的结果：要回的响应（通知为空），以及对端有没有违规。</summary>
public readonly record struct RpcDispatch(RpcMessage? Response, RpcViolation? Violation);

/// <summary>
/// 方法路由，与传输方式无关，也是协议方向的<b>唯一卡口</b>：节点和服务端、Probe 和 Agent 的入站分发都经过这里。
/// <typeparamref name="TContext"/> 由传输层提供（例如调用方节点）。路由器知道自己在连接的哪一侧（<see cref="Side"/>）和声明表，
/// 注册处理函数时核对方向、类型和角色，分发时拒绝没有声明或不该由对方发起的方法，发送时核对自己有没有资格发起。
/// </summary>
public sealed class RpcRouter<TContext>
{
    private delegate ValueTask<Result<JsonElement?>> Route(TContext context, JsonElement? parameters, CancellationToken ct);

    private static readonly JsonElement s_emptyObject = JsonDocument.Parse("{}").RootElement.Clone();

    private readonly Dictionary<string, Route> _routes = new(StringComparer.Ordinal);
    private readonly Action<string, Exception>? _onUnhandled;
    private readonly string? _localRole;
    private readonly Func<TContext, string?>? _callerRole;

    /// <param name="side">本路由器在连接的哪一侧。</param>
    /// <param name="table">方法声明表。</param>
    /// <param name="localRole">本节点的角色（节点侧和对等面用）：声明限定了角色的方法，角色不符就不能注册。</param>
    /// <param name="callerRole">从调用上下文取出对端的角色（服务端用）：声明限定了角色的方法，角色不符就拒绝。</param>
    public RpcRouter(RpcSide side, MethodTable table, string? localRole = null, Func<TContext, string?>? callerRole = null, Action<string, Exception>? onUnhandled = null)
    {
        Side = side;
        Table = table;
        _localRole = localRole;
        _callerRole = callerRole;
        _onUnhandled = onUnhandled;
    }

    public RpcSide Side { get; }

    public MethodTable Table { get; }

    private RpcPlane Plane => RpcSides.PlaneOf(Side);

    public bool Has(string method) => _routes.ContainsKey(method);

    public IReadOnlyCollection<string> Methods => _routes.Keys;

    public RpcRouter<TContext> On<TParams, TResult>(
        string method,
        JsonTypeInfo<TParams> paramsInfo,
        JsonTypeInfo<TResult> resultInfo,
        Func<TContext, TParams, CancellationToken, ValueTask<Result<TResult>>> handler)
    {
        Declare(method, MethodKind.Request);
        _routes[method] = async (context, raw, ct) =>
        {
            if (!TryDeserialize(raw, paramsInfo, out TParams? p))
            {
                return Result<JsonElement?>.Fail(FleetErrors.InvalidParams, "params could not be parsed");
            }

            Result<TResult> r = await handler(context, p, ct);
            return r.TryGet(out TResult? value, out FleetError? error)
                ? Result<JsonElement?>.Ok(JsonSerializer.SerializeToElement(value, resultInfo))
                : Result<JsonElement?>.Fail(error);
        };
        return this;
    }

    public RpcRouter<TContext> OnNotify<TParams>(
        string method,
        JsonTypeInfo<TParams> paramsInfo,
        Func<TContext, TParams, CancellationToken, ValueTask> handler)
    {
        Declare(method, MethodKind.Notification);
        _routes[method] = async (context, raw, ct) =>
        {
            if (!TryDeserialize(raw, paramsInfo, out TParams? p))
            {
                return Result<JsonElement?>.Fail(FleetErrors.InvalidParams, "params could not be parsed");
            }

            await handler(context, p, ct);
            return Result<JsonElement?>.Ok(null);
        };
        return this;
    }

    /// <summary>
    /// 处理一条请求或通知。请求返回响应；通知返回 null（失败时只丢弃）。
    /// 处理函数里的异常属于程序缺陷：在这个边界捕获，回复 internal 错误，并交给 <c>onUnhandled</c> 记录。
    /// </summary>
    public async ValueTask<RpcMessage?> DispatchAsync(TContext context, RpcMessage message, CancellationToken ct) =>
        (await DispatchCheckedAsync(context, message, ct)).Response;

    /// <summary>
    /// 先核对声明（方向、类型、角色），再分发。违规的消息不会到达处理函数，返回的结果里带 <see cref="RpcViolation"/>，
    /// 由会话决定怎么办（计数、超过上限就断开连接）。声明过、方向对、但这一侧没有处理函数的方法不算违规，只是不支持。
    /// </summary>
    public async ValueTask<RpcDispatch> DispatchCheckedAsync(TContext context, RpcMessage message, CancellationToken ct)
    {
        string method = message.Method ?? "";
        if (Check(context, message) is { } violation)
        {
            return new RpcDispatch(
                message.Id is null ? null : RpcMessage.Failure(message.Id, new FleetError(FleetErrors.ProtocolViolation, violation.Reason)),
                violation);
        }

        if (!_routes.TryGetValue(method, out Route? route))
        {
            return new RpcDispatch(message.Id is null ? null : RpcMessage.Failure(message.Id, new FleetError(FleetErrors.Unsupported, $"unknown method {method}")), null);
        }

        Result<JsonElement?> result;
        try
        {
            result = await route(context, message.Params, ct);
        }
        catch (OperationCanceledException) when (ct.IsCancellationRequested)
        {
            throw;
        }
        catch (Exception ex)
        {
            _onUnhandled?.Invoke(method, ex);
            result = Result<JsonElement?>.Fail(FleetErrors.Internal, "internal error");
        }

        if (message.Id is null)
        {
            return new RpcDispatch(null, null);
        }

        return new RpcDispatch(
            result.TryGet(out JsonElement? value, out FleetError? error)
                ? new RpcMessage { Id = message.Id, Result = value ?? s_emptyObject }
                : RpcMessage.Failure(message.Id, error),
            null);
    }

    /// <summary>入站检查：这条消息是不是对方有资格发起的、类型对不对、对方的角色允不允许。</summary>
    public RpcViolation? Check(TContext context, RpcMessage message)
    {
        string method = message.Method ?? "";
        MethodSpec? spec = Table.Find(Plane, method);
        if (spec is null)
        {
            return new RpcViolation(method, $"method {method} is not declared");
        }

        if (spec.To != Side)
        {
            return new RpcViolation(method, $"{method} cannot be initiated by the {RpcSides.Opposite(Side)} side");
        }

        MethodKind kind = message.Id is null ? MethodKind.Notification : MethodKind.Request;
        if (spec.Kind != kind)
        {
            return new RpcViolation(method, $"{method} is a {spec.Kind.ToString().ToLowerInvariant()}");
        }

        if (_callerRole is not null && !spec.AllowsRole(_callerRole(context)))
        {
            return new RpcViolation(method, $"{method} is not allowed for this node role");
        }

        return null;
    }

    /// <summary>
    /// 发送前检查：本侧有没有资格发起这个方法、类型对不对、对端的角色允不允许（服务端发给节点时传入节点的角色）。
    /// 违反是程序缺陷，抛 <see cref="InvalidOperationException"/>。
    /// </summary>
    public void CheckSend(string method, MethodKind kind, string? peerRole = null)
    {
        MethodSpec? spec = Table.Find(Plane, method);
        if (spec is null)
        {
            throw new InvalidOperationException($"method {method} is not declared");
        }

        if (spec.From != Side)
        {
            throw new InvalidOperationException($"{method} cannot be initiated by the {Side} side");
        }

        if (spec.Kind != kind)
        {
            throw new InvalidOperationException($"{method} is a {spec.Kind.ToString().ToLowerInvariant()}");
        }

        if (peerRole is not null && !spec.AllowsRole(peerRole))
        {
            throw new InvalidOperationException($"{method} cannot be sent to a {peerRole}");
        }
    }

    private void Declare(string method, MethodKind kind)
    {
        MethodSpec spec = Table.Find(Plane, method) ?? throw new InvalidOperationException($"method {method} is not declared in the method table");
        if (spec.To != Side)
        {
            throw new InvalidOperationException($"{method} is initiated by the {spec.From} side and cannot be handled on the {Side} side");
        }

        if (spec.Kind != kind)
        {
            throw new InvalidOperationException($"{method} is declared as a {spec.Kind.ToString().ToLowerInvariant()}");
        }

        if (_localRole is not null && !spec.AllowsRole(_localRole))
        {
            throw new InvalidOperationException($"{method} is not for the {_localRole} role");
        }
    }

    private static bool TryDeserialize<T>(JsonElement? raw, JsonTypeInfo<T> info, [System.Diagnostics.CodeAnalysis.NotNullWhen(true)] out T? value)
    {
        value = default;
        try
        {
            value = raw is null ? JsonSerializer.Deserialize("{}", info) : raw.Value.Deserialize(info);
            return value is not null;
        }
        catch (JsonException)
        {
            return false;
        }
    }
}