justme
48072737c9
A probe now links to the backend like an agent (probe.json carries the server address). The backend sends the agent list as desired state and the probe reports its platform and build in its own hello. The agent no longer relays the list or the probe platform. Existing probes must be re-created.
457 lines
20 KiB
C#
457 lines
20 KiB
C#
using System.Collections.Concurrent;
|
||
using System.Diagnostics;
|
||
using System.Net;
|
||
using System.Net.Sockets;
|
||
using System.Security.Cryptography;
|
||
using System.Text.Json;
|
||
using Fleet.Core;
|
||
using Fleet.Identity;
|
||
using Fleet.Link;
|
||
using Fleet.Rpc;
|
||
using Monitor.Protocol;
|
||
|
||
namespace Monitor.TestSupport;
|
||
|
||
/// <summary>
|
||
/// 假 Probe：通过真实的点对点握手（TLS 固定证书 + 签名 + ECDH）连接真实的 Agent，接收 <c>line_task_set</c>，
|
||
/// 在自己的 UDP 套接字上发回显包、接收单向数据流、估计时钟偏移，并把 Probe 一侧的列块经发件箱通道 17 交给 Agent。
|
||
/// 它实现 Probe 的协议职责，是将来真实 Probe 的参照。
|
||
/// </summary>
|
||
public sealed class ProbeSimulator : IAsyncDisposable
|
||
{
|
||
private readonly NodeIdentity _identity;
|
||
private readonly RpcRouter<ProbeSimulator> _router;
|
||
private readonly ConcurrentDictionary<Guid, TaskCompletionSource> _acks = new();
|
||
private readonly ConcurrentDictionary<uint, EchoCollector> _echoes = new();
|
||
private readonly ConcurrentDictionary<uint, TrainCollector> _trains = new();
|
||
private readonly ConcurrentQueue<LineTaskSet> _taskSets = new();
|
||
private readonly long _origin = TimeProvider.System.GetTimestamp();
|
||
private readonly CancellationTokenSource _stop = new();
|
||
private NodeId _agentId;
|
||
private RpcSession<ProbeSimulator>? _rpc;
|
||
private UdpClient? _udp;
|
||
private IPEndPoint? _agentUdp;
|
||
private PeerSessionInfo? _info;
|
||
private Task? _receiver;
|
||
private Task? _run;
|
||
|
||
public ProbeSimulator(string probeId, NodeIdentity identity)
|
||
{
|
||
ProbeId = probeId;
|
||
_identity = identity;
|
||
_router = BuildRouter();
|
||
}
|
||
|
||
public string ProbeId { get; }
|
||
|
||
/// <summary>同一个身份的另一个假 Probe（模拟 Probe 重启：新的进程、新的会话，序号从 0 重新开始）。</summary>
|
||
public ProbeSimulator Restarted() => new(ProbeId, _identity);
|
||
|
||
/// <summary>true 时拒绝按需测试的阶段（用来测试 Agent 对 Probe 拒绝的处理）。</summary>
|
||
public bool RefusePhases { get; set; }
|
||
|
||
/// <summary>收到 phase_end 之后不交出列块，只报告 0 个（用来测试 Agent 的等待逻辑）。</summary>
|
||
public bool DropFragments { get; set; }
|
||
|
||
public long ClockOffsetUs { get; private set; }
|
||
|
||
public IReadOnlyCollection<LineTaskSet> TaskSets => [.. _taskSets];
|
||
|
||
public ConcurrentQueue<TestPhaseStartParams> Phases { get; } = new();
|
||
|
||
public ConcurrentQueue<TestId> Aborts { get; } = new();
|
||
|
||
public bool IsConnected => _rpc is not null && !_stop.IsCancellationRequested;
|
||
|
||
public long NowUs() => (TimeProvider.System.GetTimestamp() - _origin) * 1_000_000 / TimeProvider.System.TimestampFrequency;
|
||
|
||
public async Task ConnectAsync(string host, int controlPort, byte[] agentSpkiSha256, int udpPort, NodeId agentId, TimeProvider? time = null, CancellationToken ct = default)
|
||
{
|
||
_agentId = agentId;
|
||
(PeerSessionInfo info, TlsFrameTransport transport) = await PeerConnector.ConnectAsync(host, controlPort, agentSpkiSha256, _identity, ProbeId, "0.0.0-sim", 8 << 20, time ?? TimeProvider.System, ct);
|
||
_info = info;
|
||
_udp = new UdpClient(new IPEndPoint(IPAddress.Loopback, 0));
|
||
_agentUdp = new IPEndPoint(IPAddress.Parse(host), udpPort);
|
||
_rpc = new RpcSession<ProbeSimulator>(transport, _router, this, time ?? TimeProvider.System);
|
||
_rpc.OnBinary = (_, _) => ValueTask.CompletedTask;
|
||
_run = _rpc.RunAsync(_stop.Token);
|
||
_receiver = Task.Run(ReceiveLoopAsync, _stop.Token);
|
||
await SyncClockAsync();
|
||
}
|
||
|
||
public async Task<LineTaskSet> WaitForTasksAsync(Func<LineTaskSet, bool> predicate, TimeSpan? timeout = null)
|
||
{
|
||
long started = Stopwatch.GetTimestamp();
|
||
TimeSpan limit = timeout ?? TimeSpan.FromSeconds(15);
|
||
while (Stopwatch.GetElapsedTime(started) < limit)
|
||
{
|
||
LineTaskSet? found = _taskSets.ToArray().LastOrDefault(predicate);
|
||
if (found is not null)
|
||
{
|
||
return found;
|
||
}
|
||
|
||
await Task.Delay(20);
|
||
}
|
||
|
||
throw new TimeoutException("the agent did not send the expected line task set");
|
||
}
|
||
|
||
/// <summary>NTP 式的时钟同步：偏移 = Agent 时钟 − (往返中点的) Probe 时钟。取最小往返的一次。</summary>
|
||
public async Task SyncClockAsync(int rounds = 5)
|
||
{
|
||
long bestRoundTrip = long.MaxValue;
|
||
for (int i = 0; i < rounds; i++)
|
||
{
|
||
long t0 = NowUs();
|
||
Result<ClockResult> reply = await _rpc!.RequestAsync(MonitorMethods.Clock, new ClockParams { ProbeMonoUs = t0 }, MonitorJsonContext.Default.ClockParams, MonitorJsonContext.Default.ClockResult, CancellationToken.None);
|
||
long t1 = NowUs();
|
||
long roundTrip = t1 - t0;
|
||
if (reply.IsSuccess && roundTrip < bestRoundTrip)
|
||
{
|
||
bestRoundTrip = roundTrip;
|
||
ClockOffsetUs = reply.Value.AgentMonoUs - ((t0 + t1) / 2);
|
||
}
|
||
}
|
||
}
|
||
|
||
/// <summary>对一个 udp_echo 线路任务发 <paramref name="count"/> 个回显包，等回复，然后交出一个列块。返回交给 Agent 的列块的序号范围。</summary>
|
||
public async Task<(int Sent, int Received)> RunUdpEchoAsync(LineTask task, int count, TimeSpan interval, CancellationToken ct = default)
|
||
{
|
||
int size = (int)task.Params["size"].GetInt64();
|
||
var collector = new EchoCollector();
|
||
_echoes[task.StreamId] = collector;
|
||
long startUtc = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds();
|
||
byte[] packet = new byte[size];
|
||
for (uint seq = 0; seq < count; seq++)
|
||
{
|
||
long t1 = NowUs();
|
||
collector.Sent[seq] = t1;
|
||
UdpPacket.Write(packet, new UdpHeader(UdpPacketType.EchoRequest, _info!.Credentials.SessionIdValue, task.StreamId, seq, t1), _info.Credentials.Key);
|
||
await _udp!.SendAsync(packet, packet.Length, _agentUdp!);
|
||
await Task.Delay(interval, ct);
|
||
}
|
||
|
||
await Task.Delay(300, ct); // 等最后的回复
|
||
_echoes.TryRemove(task.StreamId, out _);
|
||
|
||
var msg = new StreamBuilder(new StreamSpec("msg", [Col.Seq("seq"), Col.U16("size"), Col.TsProbe("t1"), Col.TsProbe("t4", nullable: true)]));
|
||
foreach ((uint seq, long t1) in collector.Sent.OrderBy(p => p.Key))
|
||
{
|
||
msg.Add(seq, size, t1, collector.Received.TryGetValue(seq, out long t4) ? t4 : null);
|
||
}
|
||
|
||
RawChunkHeader header = Header("udp_echo", task.Id, null, null, startUtc, (int)Math.Max(1, count * interval.TotalMilliseconds), task.Params);
|
||
await ShipAsync(new RawChunk(header, [msg.Build()]), ct: ct);
|
||
return (count, collector.Received.Count);
|
||
}
|
||
|
||
/// <summary>对一个 tcp_echo 线路任务建立一条 TCP（可选 TLS）连接，发 <paramref name="count"/> 条回显消息，然后交出一个列块。</summary>
|
||
public async Task<(int Sent, int Received)> RunTcpEchoAsync(LineTask task, int tcpPort, int count, TimeSpan interval, bool tls, CancellationToken ct = default)
|
||
{
|
||
using var client = new System.Net.Sockets.TcpClient();
|
||
await client.ConnectAsync(IPAddress.Loopback, tcpPort, ct);
|
||
Stream stream = client.GetStream();
|
||
if (tls)
|
||
{
|
||
#pragma warning disable CA5359 // 测试：Agent 的证书是自签名的，这里只验证数据面
|
||
var ssl = new System.Net.Security.SslStream(stream, false, (_, _, _, _) => true);
|
||
#pragma warning restore CA5359
|
||
await ssl.AuthenticateAsClientAsync(new System.Net.Security.SslClientAuthenticationOptions { TargetHost = "agent" }, ct);
|
||
stream = ssl;
|
||
}
|
||
|
||
byte[] auth = new byte[TcpWire.AuthBytes];
|
||
TcpWire.WriteAuth(auth, _info!.Credentials.SessionIdValue, task.StreamId, RandomNumberGenerator.GetBytes(TcpWire.NonceBytes), _info.Credentials.Key);
|
||
await stream.WriteAsync(auth, ct);
|
||
|
||
long startUtc = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds();
|
||
var msg = new StreamBuilder(new StreamSpec("msg", [Col.Seq("seq"), Col.U16("size"), Col.TsProbe("t1"), Col.TsProbe("t4", nullable: true)]));
|
||
int received = 0;
|
||
for (uint seq = 0; seq < count; seq++)
|
||
{
|
||
byte[] body = new byte[48];
|
||
long t1 = NowUs();
|
||
TcpWire.WriteEchoBody(body, seq, t1);
|
||
await TcpWire.WriteMessageAsync(stream, body, ct);
|
||
await stream.FlushAsync(ct);
|
||
byte[]? reply = await TcpWire.ReadMessageAsync(stream, ct);
|
||
long? t4 = reply is not null && TcpWire.TryReadEchoBody(reply, out uint echoed, out _) && echoed == seq ? NowUs() : null;
|
||
received += t4 is null ? 0 : 1;
|
||
msg.Add(seq, body.Length, t1, t4);
|
||
await Task.Delay(interval, ct);
|
||
}
|
||
|
||
RawChunkHeader header = Header("tcp_echo", task.Id, null, null, startUtc, (int)Math.Max(1, count * interval.TotalMilliseconds), task.Params);
|
||
await ShipAsync(new RawChunk(header, [msg.Build()]), ct: ct);
|
||
return (count, received);
|
||
}
|
||
|
||
/// <summary>发送任意字节到 Agent 的 UDP 端口（用来测试 Agent 对垃圾包的处理）。</summary>
|
||
public Task SendRawUdpAsync(byte[] datagram) => _udp!.SendAsync(datagram, datagram.Length, _agentUdp!);
|
||
|
||
/// <summary>用伪造的密钥发一个回显包。</summary>
|
||
public Task SendForgedEchoAsync(LineTask task, uint seq, byte[]? key = null)
|
||
{
|
||
byte[] packet = new byte[64];
|
||
UdpPacket.Write(packet, new UdpHeader(UdpPacketType.EchoRequest, _info!.Credentials.SessionIdValue, task.StreamId, seq, NowUs()), key ?? new byte[32]);
|
||
return _udp!.SendAsync(packet, packet.Length, _agentUdp!);
|
||
}
|
||
|
||
/// <summary>重放：用同一个序号、同样的密钥再发一次。</summary>
|
||
public Task SendEchoAsync(LineTask task, uint seq)
|
||
{
|
||
byte[] packet = new byte[64];
|
||
UdpPacket.Write(packet, new UdpHeader(UdpPacketType.EchoRequest, _info!.Credentials.SessionIdValue, task.StreamId, seq, NowUs()), _info.Credentials.Key);
|
||
return _udp!.SendAsync(packet, packet.Length, _agentUdp!);
|
||
}
|
||
|
||
public int EchoRepliesSeen(uint streamId) => _echoes.TryGetValue(streamId, out EchoCollector? c) ? c.Received.Count : 0;
|
||
|
||
/// <summary>交出一个已经构造好的列块并等待 Agent 确认。</summary>
|
||
public async Task ShipAsync(RawChunk fragment, bool waitForAck = true, CancellationToken ct = default)
|
||
{
|
||
Guid id = fragment.Header.ChunkId.Value;
|
||
var ack = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
|
||
_acks[id] = ack;
|
||
await _rpc!.SendBinaryAsync(OutboxFrame.Build(MonitorChannels.ColumnBlock, id, RawChunkCodec.Encode(fragment)), ct);
|
||
if (waitForAck)
|
||
{
|
||
await ack.Task.WaitAsync(TimeSpan.FromSeconds(10), ct);
|
||
}
|
||
}
|
||
|
||
public async ValueTask DisposeAsync()
|
||
{
|
||
await _stop.CancelAsync();
|
||
_udp?.Dispose();
|
||
if (_rpc is not null)
|
||
{
|
||
await _rpc.DisposeAsync();
|
||
}
|
||
|
||
foreach (Task? task in new[] { _run, _receiver })
|
||
{
|
||
if (task is not null)
|
||
{
|
||
try
|
||
{
|
||
await task.WaitAsync(TimeSpan.FromSeconds(2));
|
||
}
|
||
catch (Exception ex) when (ex is OperationCanceledException or TimeoutException or ObjectDisposedException or SocketException or IOException)
|
||
{
|
||
// 关闭中
|
||
}
|
||
}
|
||
}
|
||
|
||
_stop.Dispose();
|
||
}
|
||
|
||
private RawChunkHeader Header(string kind, TaskId? task, TestId? test, int? phase, long startUtcMs, int durationMs, IReadOnlyDictionary<string, JsonElement>? extra)
|
||
{
|
||
var header = new RawChunkHeader
|
||
{
|
||
ChunkId = new ChunkId(Guid.CreateVersion7()),
|
||
Kind = kind,
|
||
AgentId = _agentId,
|
||
TaskId = task,
|
||
TestId = test,
|
||
Phase = phase,
|
||
ProbeId = new ProbeId(ProbeId),
|
||
StartUtcMs = startUtcMs,
|
||
DurationMs = durationMs,
|
||
Clock = new ChunkClock { ProbeOffsetUs = ClockOffsetUs },
|
||
};
|
||
if (extra is not null)
|
||
{
|
||
foreach ((string key, JsonElement value) in extra)
|
||
{
|
||
header.Params[key] = value;
|
||
}
|
||
}
|
||
|
||
return header;
|
||
}
|
||
|
||
private async Task ReceiveLoopAsync()
|
||
{
|
||
while (!_stop.IsCancellationRequested)
|
||
{
|
||
UdpReceiveResult result;
|
||
try
|
||
{
|
||
result = await _udp!.ReceiveAsync(_stop.Token);
|
||
}
|
||
catch (Exception ex) when (ex is OperationCanceledException or ObjectDisposedException or SocketException)
|
||
{
|
||
if (_stop.IsCancellationRequested || ex is ObjectDisposedException)
|
||
{
|
||
return;
|
||
}
|
||
|
||
continue;
|
||
}
|
||
|
||
long now = NowUs();
|
||
byte[] data = result.Buffer;
|
||
if (!UdpPacket.TryParseHeader(data, out UdpHeader header) || !UdpPacket.VerifyMac(data, _info!.Credentials.Key))
|
||
{
|
||
continue;
|
||
}
|
||
|
||
if (header.Type == UdpPacketType.EchoReply && _echoes.TryGetValue(header.StreamId, out EchoCollector? echo))
|
||
{
|
||
echo.Received.TryAdd(header.Seq, now);
|
||
}
|
||
else if (header.Type == UdpPacketType.Data && _trains.TryGetValue(header.StreamId, out TrainCollector? train))
|
||
{
|
||
train.Received.TryAdd(header.Seq, now);
|
||
train.Size = data.Length;
|
||
}
|
||
}
|
||
}
|
||
|
||
private RpcRouter<ProbeSimulator> BuildRouter() =>
|
||
new RpcRouter<ProbeSimulator>()
|
||
.On(
|
||
FleetMethods.StateSync,
|
||
FleetRpcContext.Default.StateSyncParams,
|
||
FleetRpcContext.Default.StateSyncResult,
|
||
(probe, sync, _) =>
|
||
{
|
||
if (sync.Type == LineTaskSet.StateType)
|
||
{
|
||
probe._taskSets.Enqueue(sync.Doc.Deserialize(MonitorJsonContext.Default.LineTaskSet)!);
|
||
}
|
||
|
||
return ValueTask.FromResult(Result<StateSyncResult>.Ok(new StateSyncResult { AppliedVersion = sync.Version }));
|
||
})
|
||
.On(
|
||
MonitorMethods.TestPhase,
|
||
MonitorJsonContext.Default.TestPhaseStartParams,
|
||
MonitorJsonContext.Default.TestPhaseStartResult,
|
||
async (probe, start, ct) => Result<TestPhaseStartResult>.Ok(new TestPhaseStartResult { Accepted = await probe.StartPhaseAsync(start, ct) }))
|
||
.On(
|
||
MonitorMethods.TestPhaseEnd,
|
||
MonitorJsonContext.Default.TestPhaseEndParams,
|
||
MonitorJsonContext.Default.TestPhaseEndResult,
|
||
async (probe, end, ct) => Result<TestPhaseEndResult>.Ok(new TestPhaseEndResult { Fragments = await probe.EndPhaseAsync(end, ct) }))
|
||
.OnNotify(
|
||
MonitorMethods.TestAbort,
|
||
MonitorJsonContext.Default.TestAbortParams,
|
||
(probe, abort, _) =>
|
||
{
|
||
probe.Aborts.Enqueue(abort.TestId);
|
||
return ValueTask.CompletedTask;
|
||
})
|
||
.OnNotify(
|
||
FleetMethods.OutboxAck,
|
||
FleetRpcContext.Default.OutboxAckParams,
|
||
(probe, ack, _) =>
|
||
{
|
||
foreach (string id in ack.Ids.Concat(ack.Duplicates))
|
||
{
|
||
if (probe._acks.TryRemove(Guid.Parse(id), out TaskCompletionSource? waiter))
|
||
{
|
||
waiter.TrySetResult();
|
||
}
|
||
}
|
||
|
||
return ValueTask.CompletedTask;
|
||
});
|
||
|
||
private readonly ConcurrentDictionary<uint, TestPhaseStartParams> _activePhases = new();
|
||
|
||
private async Task<bool> StartPhaseAsync(TestPhaseStartParams start, CancellationToken ct)
|
||
{
|
||
Phases.Enqueue(start);
|
||
if (RefusePhases)
|
||
{
|
||
return false;
|
||
}
|
||
|
||
_activePhases[start.StreamId] = start;
|
||
if (start.Kind == "udp_train")
|
||
{
|
||
_trains[start.StreamId] = new TrainCollector { StartUtcMs = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds() };
|
||
// 先发几个探活包，让 Agent 学到（最新的）源地址，再回复 accepted
|
||
byte[] hello = new byte[64];
|
||
for (uint seq = 0; seq < 3; seq++)
|
||
{
|
||
UdpPacket.Write(hello, new UdpHeader(UdpPacketType.Data, _info!.Credentials.SessionIdValue, start.StreamId, seq, NowUs()), _info.Credentials.Key);
|
||
await _udp!.SendAsync(hello.AsMemory(), _agentUdp!, ct);
|
||
}
|
||
}
|
||
|
||
return true;
|
||
}
|
||
|
||
private async Task<int> EndPhaseAsync(TestPhaseEndParams end, CancellationToken ct)
|
||
{
|
||
_activePhases.TryRemove(end.StreamId, out TestPhaseStartParams? start);
|
||
if (start is null || DropFragments)
|
||
{
|
||
_trains.TryRemove(end.StreamId, out _);
|
||
return 0;
|
||
}
|
||
|
||
if (start.Kind == "udp_train" && _trains.TryRemove(end.StreamId, out TrainCollector? train))
|
||
{
|
||
await Task.Delay(200, ct); // 最后一批包还在路上
|
||
// 单块的记录数有上限（ProtocolLimits.MaxRecordsPerChunk）：按序号范围切成多个列块，Agent 一侧按同样的范围取行
|
||
const int PerFragment = 40_000;
|
||
uint total = (uint)end.SentCount;
|
||
int fragments = 0;
|
||
for (uint from = 0; from < total || fragments == 0; from += PerFragment)
|
||
{
|
||
uint to = (uint)Math.Min((long)from + PerFragment, total);
|
||
var pkt = new StreamBuilder(new StreamSpec("pkt", [Col.Seq("seq"), Col.TsProbe("t_recv", nullable: true)]));
|
||
foreach ((uint seq, long recv) in train.Received.Where(p => p.Key >= from && p.Key < to).OrderBy(p => p.Key))
|
||
{
|
||
pkt.Add(seq, recv);
|
||
}
|
||
|
||
var streams = new List<RawStream> { pkt.Build() };
|
||
if (fragments == 0)
|
||
{
|
||
var step = new StreamBuilder(Kinds.UdpTrain.FindStream("step")!);
|
||
step.Add(0, (long)start.Params["mbps"].GetDouble(), NowUs() - (start.DurationSeconds * 1_000_000L), NowUs());
|
||
streams.Add(step.Build());
|
||
}
|
||
|
||
RawChunkHeader header = Header("udp_train", null, start.TestId, start.Phase, train.StartUtcMs, start.DurationSeconds * 1000, start.Params);
|
||
header.Params[FragmentKeys.SeqFrom] = JsonSerializer.SerializeToElement(from);
|
||
header.Params[FragmentKeys.SeqTo] = JsonSerializer.SerializeToElement(to);
|
||
// 不能在这里等确认：确认是 Agent 的通知，由同一个入站队列顺序处理，而当前正处在 phase_end 的处理函数里
|
||
await ShipAsync(new RawChunk(header, streams), waitForAck: false, ct: ct);
|
||
fragments++;
|
||
if (total == 0)
|
||
{
|
||
break;
|
||
}
|
||
}
|
||
|
||
return fragments;
|
||
}
|
||
|
||
return 0;
|
||
}
|
||
|
||
private sealed class EchoCollector
|
||
{
|
||
public ConcurrentDictionary<uint, long> Sent { get; } = new();
|
||
|
||
public ConcurrentDictionary<uint, long> Received { get; } = new();
|
||
}
|
||
|
||
private sealed class TrainCollector
|
||
{
|
||
public ConcurrentDictionary<uint, long> Received { get; } = new();
|
||
|
||
public int Size { get; set; }
|
||
|
||
public long StartUtcMs { get; init; }
|
||
}
|
||
}
|