460 lines
20 KiB
C#
460 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 MoShiro.Protocol;
|
||
|
||
namespace MoShiro.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 byte[]? _agentKey;
|
||
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[] agentPublicKey, int udpPort, NodeId agentId, TimeProvider? time = null, CancellationToken ct = default)
|
||
{
|
||
_agentId = agentId;
|
||
(PeerSessionInfo info, TlsFrameTransport transport) = await PeerConnector.ConnectAsync(host, controlPort, agentPublicKey, _identity, ProbeId, "0.0.0-sim", 8 << 20, time ?? TimeProvider.System, ct);
|
||
_info = info;
|
||
_agentKey = agentPublicKey;
|
||
_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(MoShiroMethods.Clock, new ClockParams { ProbeMonoUs = t0 }, MoShiroJsonContext.Default.ClockParams, MoShiroJsonContext.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();
|
||
await stream.WriteAsync(Knock.Build(Knock.DeriveKey(_identity, _agentKey!), TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds()), ct); // 预握手敲门
|
||
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(MoShiroChannels.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>(RpcSide.Dialer, MoShiroMethodTable.Table)
|
||
.On(
|
||
FleetMethods.StateSync,
|
||
FleetRpcContext.Default.StateSyncParams,
|
||
FleetRpcContext.Default.StateSyncResult,
|
||
(probe, sync, _) =>
|
||
{
|
||
if (sync.Type == LineTaskSet.StateType)
|
||
{
|
||
probe._taskSets.Enqueue(sync.Doc.Deserialize(MoShiroJsonContext.Default.LineTaskSet)!);
|
||
}
|
||
|
||
return ValueTask.FromResult(Result<StateSyncResult>.Ok(new StateSyncResult { AppliedVersion = sync.Version }));
|
||
})
|
||
.On(
|
||
MoShiroMethods.TestPhase,
|
||
MoShiroJsonContext.Default.TestPhaseStartParams,
|
||
MoShiroJsonContext.Default.TestPhaseStartResult,
|
||
async (probe, start, ct) => Result<TestPhaseStartResult>.Ok(new TestPhaseStartResult { Accepted = await probe.StartPhaseAsync(start, ct) }))
|
||
.On(
|
||
MoShiroMethods.TestPhaseEnd,
|
||
MoShiroJsonContext.Default.TestPhaseEndParams,
|
||
MoShiroJsonContext.Default.TestPhaseEndResult,
|
||
async (probe, end, ct) => Result<TestPhaseEndResult>.Ok(new TestPhaseEndResult { Fragments = await probe.EndPhaseAsync(end, ct) }))
|
||
.OnNotify(
|
||
MoShiroMethods.TestAbort,
|
||
MoShiroJsonContext.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; }
|
||
}
|
||
}
|