MoShiro/backend/src/MoShiro.Backend.Application/Agents/MoShiroFleetHooks.cs

using System.Text.Json;
using Fleet.Core;
using Fleet.Rpc;
using Fleet.Server;
using MoShiro.Protocol;

namespace MoShiro.Backend.Application;

/// <summary>MoShiro 在 Fleet 之上的接入点：注册之后补全公开地址与端口（不监听的节点没有）。</summary>
public sealed class MoShiroEnrollmentHook : IEnrollmentHook
{
    private readonly IAgentStore _agents;
    private readonly AgentService _service;

    public MoShiroEnrollmentHook(IAgentStore agents, AgentService service)
    {
        _agents = agents;
        _service = service;
    }

    public async ValueTask OnEnrolledAsync(NodeRecord node, EnrollmentContext context, CancellationToken ct)
    {
        NodeEnrollMeta meta = new();
        if (context.ProductMeta is { } raw)
        {
            try
            {
                meta = raw.Deserialize(MoShiroJsonContext.Default.NodeEnrollMeta) ?? meta;
            }
            catch (JsonException)
            {
                // 元数据格式不对就用默认值，节点之后可以通过 hello 补全
            }
        }

        // 注册请求带了 public_host 就用它；否则来源地址是公网地址时用来源地址（私网或隧道出口不用，等 hello 里的本机地址）
        string host = meta.TcpPort <= 0 ? ""
            : !string.IsNullOrWhiteSpace(meta.PublicHost) ? meta.PublicHost
            : System.Net.IPAddress.TryParse(context.RemoteIp, out System.Net.IPAddress? remote) && PublicAddress.IsPublicIpv4(remote) ? remote.ToString() : "";
        await _agents.ApplyEnrollmentAsync(node.Id, new AgentEnrollment(host, meta.TcpPort, meta.QuicPorts, meta.Kinds), ct);
        if (meta.TcpPort > 0 && await _agents.GetAsync(node.Id, ct) is { } enrolled)
        {
            (_, string? observedV6) = AgentAddressPolicy.Choose(enrolled, [], context.RemoteIp);
            await _agents.ApplyAutoAddressesAsync(node.Id, null, observedV6, ct);
        }

        if (node.Status == NodeStatus.Active)
        {
            await _service.AfterAgentSetChangedAsync(ct);
        }
    }
}

/// <summary>节点上线、下线、hello、心跳：更新运行时状态，补全节点信息，并推送给 SSE。</summary>
public sealed class MoShiroNodeEventSink : INodeEventSink
{
    private readonly IAgentStore _agents;
    private readonly AgentRuntime _runtime;
    private readonly IEventStore _events;
    private readonly LiveHub _live;
    private readonly AgentService _service;
    private readonly NodeInfoService _nodes;
    private readonly ReleaseService _releases;
    private readonly ScheduleService _schedule;
    private readonly NodeOwnerService _owners;
    private readonly LinkRuntime _linkRuntime;
    private readonly NodeRuntimeLogService _nodeLogs;

    public MoShiroNodeEventSink(IAgentStore agents, AgentRuntime runtime, IEventStore events, LiveHub live, AgentService service, NodeInfoService nodes, ReleaseService releases, ScheduleService schedule, NodeOwnerService owners, LinkRuntime linkRuntime, NodeRuntimeLogService nodeLogs)
    {
        _nodeLogs = nodeLogs;
        _owners = owners;
        _linkRuntime = linkRuntime;
        _releases = releases;
        _schedule = schedule;
        _nodes = nodes;
        _agents = agents;
        _runtime = runtime;
        _events = events;
        _live = live;
        _service = service;
    }

    public async ValueTask OnEventAsync(NodeEvent nodeEvent, CancellationToken ct)
    {
        NodeId id = nodeEvent.Node.Id;
        switch (nodeEvent.Kind)
        {
            case NodeEventKind.Online:
                _runtime.Get(id).MarkOnline();
                _live.Publish(new NodeSignal("agent", id.Value, true, nodeEvent.AtMs));
                await _events.AddAsync(id, "agent_online", nodeEvent.AtMs, null, ct);
                break;
            case NodeEventKind.Offline:
                _runtime.Get(id).MarkOffline(nodeEvent.Reason);
                _runtime.MarkOffline(id);
                _releases.RecordStatus(id, null);
                _linkRuntime.MarkDialerOffline(id.Value, nodeEvent.AtMs);
                _live.Publish(new NodeSignal("agent", id.Value, false, nodeEvent.AtMs));
                await _events.AddAsync(id, "agent_offline", nodeEvent.AtMs, null, ct);
                break;
            case NodeEventKind.Hello when nodeEvent.Hello is { } hello:
                await ApplyHelloAsync(id, hello, nodeEvent.RemoteIp, ct);
                break;
            case NodeEventKind.Heartbeat when nodeEvent.Extra is { } extra:
                await ApplyHeartbeatAsync(id, nodeEvent, extra, ct);
                break;
        }
    }

    private async ValueTask ApplyHeartbeatAsync(NodeId id, NodeEvent nodeEvent, System.Text.Json.JsonElement extra, CancellationToken ct)
    {
        await _nodes.SetVersionAsync(id.Value, nodeEvent.Node.NodeVersion, ct);
        if (_nodes.ShouldRetryGeo(id.Value) && await _agents.GetAsync(id, ct) is { } unlocated)
        {
            _nodes.ObserveAgentInBackground(id.Value, GeoHostOf(unlocated, nodeEvent.RemoteIp), null);
        }

        try
        {
            AgentRuntimeState state = _runtime.Get(id);
            NodeHeartbeatExtra? beat = extra.Deserialize(MoShiroJsonContext.Default.NodeHeartbeatExtra);
            state.LastHeartbeat = beat;
            state.LastHeartbeatAtMs = nodeEvent.AtMs;
            _releases.RecordStatus(id, beat?.Update);
            if (beat is null)
            {
                return;
            }

            _nodeLogs.Append(id, beat.RuntimeLogs, beat.DroppedRuntimeLogs ?? 0);
            state.ReplaceDialers(beat.Dialers.Select(d => new ConnectedDialerInfo(d.Id.Value, d.Version, d.Remote, d.SinceUtcMs, d.IpVersion)));
            foreach (ConnectedDialer dialer in beat.Dialers)
            {
                _nodes.ObserveDialerInBackground(dialer.Id.Value, dialer.Remote);
            }

            if (beat.Links is { } links)
            {
                _linkRuntime.Replace(id.Value, nodeEvent.AtMs, links);
            }

            if (beat.Ipv6 is { } ipv6)
            {
                await _agents.SetIpv6DiagnosticsAsync(id, JsonSerializer.Serialize(ipv6, MoShiroJsonContext.Default.Ipv6Diagnostic), ct);
            }
        }
        catch (JsonException)
        {
            // 忽略格式不正确的附加字段
        }
    }

    /// <summary>查国家和运营商用哪个地址：对外地址优先，其次 IPv6 地址，都没有（NAT 后面、或者没有上报公网地址）才用这条连接的来源地址。</summary>
    private static string? GeoHostOf(AgentRecord? agent, string? remoteIp) =>
        !string.IsNullOrEmpty(agent?.PublicHost) ? agent.PublicHost
        : !string.IsNullOrEmpty(agent?.Ipv6Host) ? agent.Ipv6Host
        : remoteIp;

    private async ValueTask ApplyHelloAsync(NodeId id, HelloParams hello, string? remoteIp, CancellationToken ct)
    {
        var kinds = new List<string>();
        foreach (string capability in hello.Capabilities)
        {
            if (capability.StartsWith(MoShiroCapabilities.KindsPrefix, StringComparison.Ordinal))
            {
                kinds.AddRange(capability[MoShiroCapabilities.KindsPrefix.Length..].Split(',', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries));
            }
        }

        NodeHelloMeta meta = new();
        if (hello.Meta is { } raw)
        {
            try
            {
                meta = raw.Deserialize(MoShiroJsonContext.Default.NodeHelloMeta) ?? meta;
            }
            catch (JsonException)
            {
                // 忽略
            }
        }

        _owners.Report(id.Value, meta.Owners);
        if (meta.Platform is { } platform)
        {
            platform.Version = hello.NodeVersion;
        }

        await _nodes.SetVersionAsync(id.Value, hello.NodeVersion, ct);
        AgentRecord? before = await _agents.GetAsync(id, ct);
        string? limitsJson = meta.LocalLimits is null ? null : JsonSerializer.Serialize(meta.LocalLimits, LocalLimitsJson.Default.LocalLimits);
        string? ipv6Json = meta.Ipv6 is null ? null : JsonSerializer.Serialize(meta.Ipv6, MoShiroJsonContext.Default.Ipv6Diagnostic);
        await _agents.ApplyHelloAsync(id, new AgentHello(kinds, hello.NodeVersion, limitsJson, meta.Listen, ipv6Json), ct);

        // 自动识别对外地址：监听的节点上报的公网地址；手填的不动
        if (meta.Listen?.Addresses is not null && await _agents.GetAsync(id, ct) is { } reported)
        {
            (string? publicHost, string? ipv6Host) = AgentAddressPolicy.Choose(reported, meta.Listen.Addresses, null);
            await _agents.ApplyAutoAddressesAsync(id, publicHost, ipv6Host, ct);
        }

        AgentRecord? after = await _agents.GetAsync(id, ct);
        _nodes.ObserveAgentInBackground(id.Value, GeoHostOf(after, remoteIp), meta.Platform);
        await _releases.OfferToNodeAsync(id, ct);
        await _schedule.SendToNodeAsync(id, ct);

        // 监听状态、端口或对外地址变了，拨号方需要新的目标列表
        if (before is not null && after is not null && after.Status == NodeStatus.Active
            && (before.Listens != after.Listens || before.TcpPort != after.TcpPort || !before.QuicPorts.SequenceEqual(after.QuicPorts) || before.Ipv6Supported != after.Ipv6Supported || before.PublicHost != after.PublicHost || before.Ipv6Host != after.Ipv6Host))
        {
            await _service.AfterAgentSetChangedAsync(ct);
        }
    }
}

[System.Text.Json.Serialization.JsonSerializable(typeof(LocalLimits))]
internal sealed partial class LocalLimitsJson : System.Text.Json.Serialization.JsonSerializerContext;