MoShiro/backend/agent/src/MoShiro.Agent/Tasks/TargetRunners.cs

using System.Net;
using System.Text.Json;
using Fleet.Core;
using MoShiro.Measurement;
using MoShiro.Protocol;

namespace MoShiro.Agent;

/// <summary>
/// 目标任务（Agent 自己探测一个地址，没有对端）：用和线路任务同一套积木（<see cref="BlockCatalog"/> 里的拨号方，例如 icmp_ping、connect），
/// 只是列块不交给 Agent 合并，而是自己盖章、10 秒一块直接上传。目标可以是 IP、域名或（http）URL；域名在 Agent 上解析，
/// 每个解析出的地址都必须在本地 <c>allowed_target_cidrs</c> 内（IP 目标由 <see cref="BudgetGuard"/> 在接受任务之前检查，域名在这里检查）。
/// </summary>
public sealed class TargetRunners
{
    private static readonly TimeSpan s_chunkSpan = TimeSpan.FromSeconds(ProtocolLimits.ChunkSeconds);

    private static readonly TimeSpan s_resolveRetry = TimeSpan.FromMinutes(1);

    private readonly Dictionary<string, Running> _running = [];
    private readonly object _gate = new();
    private readonly ChunkAssembler _chunks;
    private readonly NodeId _agent;
    private readonly MonoClock _clock;
    private readonly TimeProvider _time;
    private readonly INodeLog _log;
    private readonly RateChangeLog? _rateChanges;

    public TargetRunners(ChunkAssembler chunks, NodeId agent, MonoClock clock, TimeProvider time, INodeLog log, RateChangeLog? rateChanges = null)
    {
        _rateChanges = rateChanges;
        _chunks = chunks;
        _agent = agent;
        _clock = clock;
        _time = time;
        _log = log;
    }

    /// <summary>让运行中的任务与期望集合一致：参数没有变化的保持，其余启动或停止。</summary>
    public void Reconcile(IReadOnlyList<(TaskSpec Spec, TaskDecision Decision)> desired, Func<IPAddress, bool>? allow = null)
    {
        lock (_gate)
        {
            var wanted = desired.ToDictionary(d => d.Spec.Id.Value, d => d);
            foreach ((string id, Running running) in _running.ToList())
            {
                if (!wanted.TryGetValue(id, out (TaskSpec Spec, TaskDecision Decision) entry) || Signature(entry.Spec) != running.Signature)
                {
                    running.Cancel.Cancel();
                    _running.Remove(id);
                }
            }

            foreach ((string id, (TaskSpec spec, TaskDecision decision)) in wanted)
            {
                if (_running.ContainsKey(id))
                {
                    continue;
                }

                var cts = new CancellationTokenSource();
                _running[id] = new Running(Signature(spec), cts);
                _ = Task.Run(() => RunAsync(spec, decision, allow, cts.Token), CancellationToken.None);
            }
        }
    }

    private static string Signature(TaskSpec spec) =>
        spec.Kind + "|" + spec.Target + "|" + string.Join(',', spec.Params.OrderBy(p => p.Key, StringComparer.Ordinal).Select(p => p.Key + "=" + p.Value.GetRawText()));

    private async Task RunAsync(TaskSpec spec, TaskDecision decision, Func<IPAddress, bool>? allow, CancellationToken ct)
    {
        try
        {
            if (!TargetHost.TryParse(spec.Target, out string host, out int targetPort, out _))
            {
                _log.Warn("target is not valid", ("task", spec.Id.Value));
                return;
            }

            // http 自己解析并检查每个地址；其余 Kind 在这里解析出一个允许的地址作为对端
            IPAddress? address = null;
            if (spec.Kind != TargetHost.HttpKind)
            {
                address = await ResolveAllowedAsync(spec, host, allow, ct);
            }

            int port = targetPort != 0 ? targetPort : spec.Kind == "tcp_connect" ? (int)decision.Resolved!.GetInt("port") : 0;
            var env = new BlockEnv
            {
                Clock = _clock,
                Log = _log,
                Session = () => null,
                PeerAddress = address,
                Target = spec.Target,
                AllowAddress = allow,
                OnRateChange = _rateChanges is null ? null : _rateChanges.Add,
                Fragments = new LocalSink(this, spec, address?.AddressFamily == System.Net.Sockets.AddressFamily.InterNetworkV6 ? IpVersions.V6 : IpVersions.V4),
            };
            IBlock block = BlockCatalog.Find(spec.Kind)!.Dialer(env, new BlockSpec(decision.Kind!, decision.Resolved!, 0, port, new FragmentScope(spec.Id, null, null)));
            using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            Task run = block.RunAsync(cts.Token);
            try
            {
                using var timer = new PeriodicTimer(s_chunkSpan, _time);
                while (await timer.WaitForNextTickAsync(ct))
                {
                    block.Flush(FlushMode.Settled);
                }
            }
            finally
            {
                await cts.CancelAsync();
                await run;
                block.Flush(FlushMode.All);
            }
        }
        catch (OperationCanceledException)
        {
            // 任务被撤销
        }
        catch (Exception ex) when (ex is System.Net.Sockets.SocketException or InvalidOperationException or BlockFaultException)
        {
            _log.Warn("target task stopped", ("task", spec.Id.Value), ("err", ex.Message));
        }
    }

    /// <summary>解析目标，取第一个允许的地址；解析失败或没有允许的地址就每分钟重试（域名的解析结果会变）。</summary>
    private async Task<IPAddress> ResolveAllowedAsync(TaskSpec spec, string host, Func<IPAddress, bool>? allow, CancellationToken ct)
    {
        while (true)
        {
            try
            {
                IPAddress[] addresses = IPAddress.TryParse(host, out IPAddress? literal) ? [literal] : await Dns.GetHostAddressesAsync(host, ct);
                IPAddress? chosen = addresses.FirstOrDefault(a => allow?.Invoke(a) ?? true);
                if (chosen is not null)
                {
                    return chosen;
                }

                _log.Warn("no address of the target is allowed", ("task", spec.Id.Value), ("host", host));
            }
            catch (System.Net.Sockets.SocketException ex)
            {
                _log.Warn("target did not resolve", ("task", spec.Id.Value), ("host", host), ("err", ex.Message));
            }

            await Task.Delay(s_resolveRetry, _time, ct);
        }
    }

    /// <summary>目标任务的行没有拨号方一侧可以合并：积木交出的就是完整的行，Agent 自己盖章（目标、时钟锚点）并放进发件箱。</summary>
    private sealed class LocalSink : IFragmentSink
    {
        private readonly TargetRunners _owner;
        private readonly TaskSpec _spec;
        private readonly string _ipVersion;

        public LocalSink(TargetRunners owner, TaskSpec spec, string ipVersion)
        {
            _owner = owner;
            _spec = spec;
            _ipVersion = ipVersion;
        }

        public void Ship(string kind, FragmentScope scope, IReadOnlyList<RawStream> streams, IReadOnlyDictionary<string, JsonElement>? parameters, long startUtcMs, int durationMs, int port)
        {
            ClockAnchor anchor = _owner._clock.Anchor();
            var header = new RawChunkHeader
            {
                ChunkId = ChunkId.New(_owner._time),
                Kind = kind,
                TaskId = scope.Task,
                AgentId = _owner._agent,
                Target = _spec.Target,
                IpVersion = _ipVersion,
                Params = parameters is null ? [] : new Dictionary<string, JsonElement>(parameters),
                Port = port,
                StartUtcMs = startUtcMs,
                DurationMs = Math.Clamp(durationMs, 1, ChunkFinalizer.MaxDurationMs),
                Clock = new ChunkClock { ListenerMonoUs = anchor.ListenerMonoUs, UtcMs = anchor.UtcMs, DialerOffsetUs = 0 },
            };
            _owner._chunks.Emit(new RawChunk(header, streams));
        }
    }

    private sealed record Running(string Signature, CancellationTokenSource Cancel);
}