470 lines
20 KiB
C#
470 lines
20 KiB
C#
using System.Collections.Concurrent;
|
|
using System.Net.WebSockets;
|
|
using System.Security.Cryptography;
|
|
using System.Text;
|
|
using System.Text.Json;
|
|
using System.Text.Json.Nodes;
|
|
using Fleet.Core;
|
|
using Fleet.Identity;
|
|
using Fleet.Rpc;
|
|
using MoShiro.Protocol;
|
|
|
|
namespace MoShiro.TestSupport;
|
|
|
|
/// <summary>
|
|
/// 一个会说 Fleet 协议的假 Agent：用加入令牌注册 → 连接 → 签名认证 → hello → 接收期望状态 → 上传原始块并等待 ACK。
|
|
/// 只用于测试和离线工具，不是真正的节点。
|
|
/// </summary>
|
|
public sealed class AgentSimulator : IAsyncDisposable
|
|
{
|
|
/// <summary>节点上保留的日志（<c>moshiro.logs.fetch</c> 按编号返回它们）和收到的拉取请求。</summary>
|
|
public System.Collections.Concurrent.ConcurrentBag<LogEntryDto> Logs { get; } = [];
|
|
|
|
public System.Collections.Concurrent.ConcurrentQueue<LogsFetchParams> LogFetches { get; } = new();
|
|
|
|
private readonly NodeIdentity _identity;
|
|
private readonly NodeIdentity _tlsKey;
|
|
private readonly ClientWebSocket _socket = new();
|
|
private readonly ConcurrentDictionary<Guid, TaskCompletionSource<bool>> _acks = new();
|
|
private readonly ConcurrentDictionary<long, TaskCompletionSource<RpcMessage>> _responses = new();
|
|
private readonly CancellationTokenSource _stop = new();
|
|
private readonly SemaphoreSlim _sendLock = new(1, 1);
|
|
private long _nextId;
|
|
private int _disposed;
|
|
private Task? _loop;
|
|
|
|
private AgentSimulator(NodeIdentity identity, NodeIdentity tlsKey, NodeId nodeId, string status, bool listens)
|
|
{
|
|
_identity = identity;
|
|
_tlsKey = tlsKey;
|
|
Listens = listens;
|
|
ExportedKey = identity.ExportPkcs8();
|
|
NodeId = nodeId;
|
|
EnrollmentStatus = status;
|
|
}
|
|
|
|
public NodeId NodeId { get; }
|
|
|
|
/// <summary>是否像监听的节点那样上报端口（false：只拨出的节点）。</summary>
|
|
public bool Listens { get; set; }
|
|
|
|
/// <summary>身份私钥（PKCS#8）：释放之后仍可用它以同一个身份重新连接。</summary>
|
|
public byte[] ExportedKey { get; }
|
|
|
|
public string EnrollmentStatus { get; }
|
|
|
|
public string AuthStatus { get; private set; } = "";
|
|
|
|
public ConcurrentQueue<TaskSet> TaskSets { get; } = new();
|
|
|
|
public ConcurrentQueue<TestRunParams> TestRuns { get; } = new();
|
|
|
|
/// <summary>面板发来的发布清单（<c>moshiro.update.offer</c>）和更新请求（<c>moshiro.update.apply</c>）。</summary>
|
|
public ConcurrentQueue<UpdateOfferParams> UpdateOffers { get; } = new();
|
|
|
|
/// <summary>面板下发的定时重启设置（<c>moshiro.schedule</c>），按到达顺序。</summary>
|
|
public ConcurrentQueue<ScheduleParams> Schedules { get; } = new();
|
|
|
|
private int _updateApplies;
|
|
|
|
public int UpdateApplies => Volatile.Read(ref _updateApplies);
|
|
|
|
public ConcurrentQueue<TestId> TestCancels { get; } = new();
|
|
|
|
public IReadOnlyList<RejectedItem> RejectNextState { get; set; } = [];
|
|
|
|
public bool Revoked { get; private set; }
|
|
|
|
public string? CloseReason { get; private set; }
|
|
|
|
/// <summary>收到 moshiro.test.run 时的应答：默认接受。</summary>
|
|
public Func<TestRunParams, (bool Accepted, string? ErrorCode)> TestRunHandler { get; set; } = _ => (true, null);
|
|
|
|
public byte[] PublicKey => _identity.PublicKeySpki;
|
|
|
|
/// <summary>注册时交给面板的 TLS 证书公钥（拨号方用它固定证书、派生敲门密钥）。</summary>
|
|
public byte[] TlsPublicKey => _tlsKey.PublicKeySpki;
|
|
|
|
/// <summary>用这个节点的私钥为认证随机数签名（可以冒充别的节点 ID：用来测试服务端的校验）。</summary>
|
|
public byte[] SignAuth(NodeId claimedNode, string nonce) => _identity.Sign(SignaturePurpose.Auth, claimedNode.Value, nonce);
|
|
|
|
public static async Task<AgentSimulator> EnrollAsync(HttpClient http, string joinToken, string name, string[]? kinds = null, string publicHost = "203.0.113.10", bool listens = true)
|
|
{
|
|
NodeIdentity identity = NodeIdentity.Generate();
|
|
NodeIdentity tls = NodeIdentity.Generate();
|
|
string pubkey = Convert.ToBase64String(identity.PublicKeySpki);
|
|
byte[] tokenHash = SHA256.HashData(Encoding.UTF8.GetBytes(joinToken));
|
|
string proof = Convert.ToBase64String(identity.Sign(SignaturePurpose.Enroll, Convert.ToHexStringLower(tokenHash), pubkey));
|
|
var body = new JsonObject
|
|
{
|
|
["join_token"] = joinToken,
|
|
["name"] = name,
|
|
["pubkey"] = pubkey,
|
|
["tls_spki"] = Convert.ToBase64String(NodeIdentity.Fingerprint(tls.PublicKeySpki)),
|
|
["tls_pubkey"] = Convert.ToBase64String(tls.PublicKeySpki),
|
|
["proof"] = proof,
|
|
["node_version"] = "0.0.1-sim",
|
|
["meta"] = new JsonObject
|
|
{
|
|
["tcp_port"] = listens ? 7443 : 0,
|
|
["quic_ports"] = listens ? new JsonArray(7443) : new JsonArray(),
|
|
["public_host"] = listens ? publicHost : "",
|
|
["kinds"] = new JsonArray((kinds ?? ["quic_ping", "tcp_ping", "quic_stream", "tcp_connect", "tls_hold", "tcp_transfer", "icmp_ping"]).Select(k => (JsonNode?)JsonValue.Create(k)).ToArray()),
|
|
},
|
|
};
|
|
|
|
using var content = new StringContent(body.ToJsonString(), Encoding.UTF8, "application/json");
|
|
using HttpResponseMessage response = await http.PostAsync("/api/agent/v1/enroll", content);
|
|
string text = await response.Content.ReadAsStringAsync();
|
|
if (!response.IsSuccessStatusCode)
|
|
{
|
|
identity.Dispose();
|
|
tls.Dispose();
|
|
throw new InvalidOperationException($"enroll failed: {(int)response.StatusCode} {text}");
|
|
}
|
|
|
|
JsonObject result = JsonNode.Parse(text)!.AsObject();
|
|
return new AgentSimulator(identity, tls, new NodeId(result["node_id"]!.GetValue<string>()), result["status"]!.GetValue<string>(), listens);
|
|
}
|
|
|
|
/// <summary>用已有的身份恢复一个模拟器（例如后端重启之后同一个节点重新连接）。</summary>
|
|
public static AgentSimulator Restore(byte[] pkcs8, NodeId nodeId, bool listens = true) => new(NodeIdentity.FromPkcs8(pkcs8), NodeIdentity.Generate(), nodeId, "active", listens);
|
|
|
|
/// <summary>连接并完成认证与 hello。认证失败时抛出异常。</summary>
|
|
public async Task ConnectAsync(Uri baseAddress, bool sendHello = true, string[]? kinds = null)
|
|
{
|
|
var uri = new UriBuilder(baseAddress) { Scheme = baseAddress.Scheme == "https" ? "wss" : "ws", Path = "/api/agent/v1/rpc" }.Uri;
|
|
await _socket.ConnectAsync(uri, _stop.Token);
|
|
|
|
RpcMessage challenge = await ReceiveMessageAsync(_stop.Token) ?? throw new InvalidOperationException("no challenge");
|
|
ChallengeParams nonce = challenge.Params!.Value.Deserialize(FleetRpcContext.Default.ChallengeParams)!;
|
|
string signature = Convert.ToBase64String(_identity.Sign(SignaturePurpose.Auth, NodeId.Value, nonce.Nonce));
|
|
await SendAsync(RpcMessage.Request(NextId(), FleetMethods.Auth, new AuthParams { NodeId = NodeId.Value, Signature = signature }, FleetRpcContext.Default.AuthParams));
|
|
RpcMessage? authReply = await ReceiveMessageAsync(_stop.Token);
|
|
if (authReply?.Error is { } error)
|
|
{
|
|
throw new InvalidOperationException($"auth failed: {error.Data?.Code} {error.Message}");
|
|
}
|
|
|
|
AuthStatus = authReply!.Result!.Value.Deserialize(FleetRpcContext.Default.AuthResult)!.Status;
|
|
_loop = Task.Run(ReceiveLoopAsync);
|
|
if (sendHello)
|
|
{
|
|
await HelloAsync(kinds);
|
|
}
|
|
}
|
|
|
|
/// <summary>hello 里上报的本机公网地址（缺省不上报）。</summary>
|
|
public string[]? Addresses { get; set; }
|
|
|
|
public async Task HelloAsync(string[]? kinds = null)
|
|
{
|
|
var meta = JsonSerializer.SerializeToElement(
|
|
new NodeHelloMeta
|
|
{
|
|
Listen = Listens ? new ListenInfo { TcpPort = 7443, QuicPorts = [7443], Addresses = Addresses is null ? null : [.. Addresses] } : null,
|
|
LocalLimits = new LocalLimits { MaxTotalMbps = 200, MaxTotalPps = 20_000, MaxConcurrentTests = 1 },
|
|
},
|
|
MoShiroJsonContext.Default.NodeHelloMeta);
|
|
var hello = new HelloParams
|
|
{
|
|
NodeVersion = "0.0.1-sim",
|
|
Capabilities = [MoShiroCapabilities.Base, MoShiroCapabilities.KindsPrefix + string.Join(',', kinds ?? ["quic_ping", "tcp_ping", "quic_stream", "tcp_connect", "tls_hold", "tcp_transfer", "icmp_ping"])],
|
|
NowUtcMs = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds(),
|
|
Meta = meta,
|
|
};
|
|
await RequestAsync(FleetMethods.Hello, JsonSerializer.SerializeToElement(hello, FleetRpcContext.Default.HelloParams));
|
|
}
|
|
|
|
/// <summary>上传一个原始块并等待服务端 ACK（持久化之后才会确认）。</summary>
|
|
/// <summary>设置之后，这个假 Agent 发出的每个块都归属于这个拨号方（夹具的默认拨号方）。</summary>
|
|
public NodeId? DialerForChunks { get; set; }
|
|
|
|
public async Task<bool> SendChunkAsync(RawChunk chunk, TimeSpan? timeout = null)
|
|
{
|
|
if (DialerForChunks is { } dialer)
|
|
{
|
|
chunk.Header.DialerId = dialer;
|
|
}
|
|
|
|
byte[] bytes = RawChunkCodec.Encode(chunk);
|
|
Guid id = chunk.Header.ChunkId.Value;
|
|
var ack = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
|
|
_acks[id] = ack;
|
|
await SendBinaryAsync(OutboxFrame.Build(MoShiroChannels.RawChunk, id, bytes));
|
|
return await ack.Task.WaitAsync(timeout ?? TimeSpan.FromSeconds(10));
|
|
}
|
|
|
|
/// <summary>发送一个任意的文本帧（用来测试服务端对非法输入的处理）。</summary>
|
|
public Task SendTextAsync(string text) => SendFrameAsync(Encoding.UTF8.GetBytes(text), WebSocketMessageType.Text);
|
|
|
|
/// <summary>发送一个任意的二进制帧（用来测试服务端对非法输入的处理）。</summary>
|
|
public Task SendBinaryAsync(byte[] frame) => SendFrameAsync(frame, WebSocketMessageType.Binary);
|
|
|
|
public Task SendEventAsync(string kind, string? detail = null) =>
|
|
SendAsync(RpcMessage.Notification(
|
|
MoShiroMethods.Event,
|
|
new AgentEventParams { Kind = kind, AtUtcMs = TimeProvider.System.GetUtcNow().ToUnixTimeMilliseconds(), Detail = detail },
|
|
MoShiroJsonContext.Default.AgentEventParams));
|
|
|
|
public Task SendHostAsync(HostSnapshot snapshot) =>
|
|
SendAsync(RpcMessage.Notification(MoShiroMethods.Host, snapshot, MoShiroJsonContext.Default.HostSnapshot));
|
|
|
|
public Task SendTestStatusAsync(TestId test, string status, int? phase = null, string? errorCode = null) =>
|
|
SendAsync(RpcMessage.Notification(
|
|
MoShiroMethods.TestStatus,
|
|
new TestStatusParams { TestId = test, Status = status, Phase = phase, ErrorCode = errorCode },
|
|
MoShiroJsonContext.Default.TestStatusParams));
|
|
|
|
/// <summary>像节点的更新器报告的那样，随心跳上报更新状态。</summary>
|
|
public Task SendUpdateStatusAsync(UpdateStatus status) =>
|
|
SendAsync(RpcMessage.Notification(
|
|
FleetMethods.Heartbeat,
|
|
new HeartbeatParams { Extra = JsonSerializer.SerializeToElement(new NodeHeartbeatExtra { Update = status }, MoShiroJsonContext.Default.NodeHeartbeatExtra) },
|
|
FleetRpcContext.Default.HeartbeatParams));
|
|
|
|
public async Task<RpcMessage> RequestRawAsync(string method, JsonElement? parameters = null) => await RequestAsync(method, parameters);
|
|
|
|
public async Task<TaskSet> WaitForTaskSetAsync(Func<TaskSet, bool> predicate, TimeSpan? timeout = null)
|
|
{
|
|
long started = TimeProvider.System.GetTimestamp();
|
|
TimeSpan limit = timeout ?? TimeSpan.FromSeconds(10);
|
|
while (TimeProvider.System.GetElapsedTime(started) < limit)
|
|
{
|
|
foreach (TaskSet set in TaskSets.ToArray().Reverse())
|
|
{
|
|
if (predicate(set))
|
|
{
|
|
return set;
|
|
}
|
|
}
|
|
|
|
await Task.Delay(25);
|
|
}
|
|
|
|
throw new TimeoutException("the expected task set never arrived");
|
|
}
|
|
|
|
public async Task WaitUntilAsync(Func<bool> condition, TimeSpan? timeout = null)
|
|
{
|
|
long started = TimeProvider.System.GetTimestamp();
|
|
TimeSpan limit = timeout ?? TimeSpan.FromSeconds(10);
|
|
while (TimeProvider.System.GetElapsedTime(started) < limit)
|
|
{
|
|
if (condition())
|
|
{
|
|
return;
|
|
}
|
|
|
|
await Task.Delay(25);
|
|
}
|
|
|
|
throw new TimeoutException("condition was not met in time");
|
|
}
|
|
|
|
public async ValueTask DisposeAsync()
|
|
{
|
|
if (Interlocked.Exchange(ref _disposed, 1) != 0)
|
|
{
|
|
return;
|
|
}
|
|
|
|
await _stop.CancelAsync();
|
|
if (_socket.State == WebSocketState.Open)
|
|
{
|
|
try
|
|
{
|
|
await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "bye", CancellationToken.None);
|
|
}
|
|
catch (WebSocketException)
|
|
{
|
|
// 已经断开
|
|
}
|
|
catch (OperationCanceledException) when (_stop.IsCancellationRequested)
|
|
{
|
|
// Canceling the receive loop may abort/dispose the transport
|
|
// between the State check and CloseOutputAsync.
|
|
}
|
|
}
|
|
|
|
if (_loop is not null)
|
|
{
|
|
try
|
|
{
|
|
await _loop;
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
// 正常结束
|
|
}
|
|
}
|
|
|
|
_socket.Dispose();
|
|
_identity.Dispose();
|
|
_stop.Dispose();
|
|
_sendLock.Dispose();
|
|
}
|
|
|
|
private long NextId() => Interlocked.Increment(ref _nextId);
|
|
|
|
private async Task<RpcMessage> RequestAsync(string method, JsonElement? parameters)
|
|
{
|
|
long id = NextId();
|
|
var waiter = new TaskCompletionSource<RpcMessage>(TaskCreationOptions.RunContinuationsAsynchronously);
|
|
_responses[id] = waiter;
|
|
await SendAsync(new RpcMessage { Id = id, Method = method, Params = parameters });
|
|
return await waiter.Task.WaitAsync(TimeSpan.FromSeconds(10));
|
|
}
|
|
|
|
private Task SendAsync(RpcMessage message) =>
|
|
SendFrameAsync(JsonSerializer.SerializeToUtf8Bytes(message, FleetRpcContext.Default.RpcMessage), WebSocketMessageType.Text);
|
|
|
|
private async Task SendFrameAsync(byte[] bytes, WebSocketMessageType type)
|
|
{
|
|
await _sendLock.WaitAsync(_stop.Token);
|
|
try
|
|
{
|
|
await _socket.SendAsync(bytes, type, endOfMessage: true, _stop.Token);
|
|
}
|
|
finally
|
|
{
|
|
_sendLock.Release();
|
|
}
|
|
}
|
|
|
|
private async Task<RpcMessage?> ReceiveMessageAsync(CancellationToken ct)
|
|
{
|
|
var buffer = new byte[16 * 1024];
|
|
using var message = new MemoryStream();
|
|
while (true)
|
|
{
|
|
WebSocketReceiveResult result = await _socket.ReceiveAsync(buffer, ct);
|
|
if (result.MessageType == WebSocketMessageType.Close)
|
|
{
|
|
CloseReason = result.CloseStatusDescription;
|
|
if (_socket.State == WebSocketState.CloseReceived)
|
|
{
|
|
// 按协议确认对端发起的关闭
|
|
try
|
|
{
|
|
await _socket.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "ack", CancellationToken.None);
|
|
}
|
|
catch (WebSocketException)
|
|
{
|
|
// 对端已经断开
|
|
}
|
|
}
|
|
|
|
return null;
|
|
}
|
|
|
|
message.Write(buffer, 0, result.Count);
|
|
if (result.EndOfMessage)
|
|
{
|
|
return JsonSerializer.Deserialize(message.ToArray(), FleetRpcContext.Default.RpcMessage);
|
|
}
|
|
}
|
|
}
|
|
|
|
private async Task ReceiveLoopAsync()
|
|
{
|
|
try
|
|
{
|
|
while (!_stop.IsCancellationRequested)
|
|
{
|
|
RpcMessage? message = await ReceiveMessageAsync(_stop.Token);
|
|
if (message is null)
|
|
{
|
|
return;
|
|
}
|
|
|
|
await HandleAsync(message);
|
|
}
|
|
}
|
|
catch (OperationCanceledException)
|
|
{
|
|
// 正常结束
|
|
}
|
|
catch (WebSocketException)
|
|
{
|
|
// 连接断开
|
|
}
|
|
}
|
|
|
|
private async Task HandleAsync(RpcMessage message)
|
|
{
|
|
if (message.IsResponse)
|
|
{
|
|
if (_responses.TryRemove(message.Id!.Value, out TaskCompletionSource<RpcMessage>? waiter))
|
|
{
|
|
waiter.TrySetResult(message);
|
|
}
|
|
|
|
return;
|
|
}
|
|
|
|
switch (message.Method)
|
|
{
|
|
case FleetMethods.StateSync:
|
|
StateSyncParams sync = message.Params!.Value.Deserialize(FleetRpcContext.Default.StateSyncParams)!;
|
|
if (sync.Type == TaskSet.StateType)
|
|
{
|
|
TaskSets.Enqueue(sync.Doc.Deserialize(MoShiroJsonContext.Default.TaskSet)!);
|
|
}
|
|
|
|
var result = new StateSyncResult { AppliedVersion = sync.Version, Rejected = [.. RejectNextState] };
|
|
await SendAsync(RpcMessage.Success(message.Id!.Value, result, FleetRpcContext.Default.StateSyncResult));
|
|
break;
|
|
case FleetMethods.OutboxAck:
|
|
OutboxAckParams ack = message.Params!.Value.Deserialize(FleetRpcContext.Default.OutboxAckParams)!;
|
|
foreach (string id in ack.Ids.Concat(ack.Duplicates))
|
|
{
|
|
if (_acks.TryRemove(Guid.Parse(id), out TaskCompletionSource<bool>? waiter))
|
|
{
|
|
waiter.TrySetResult(ack.Duplicates.Contains(id) ? false : true);
|
|
}
|
|
}
|
|
|
|
break;
|
|
case FleetMethods.Revoked:
|
|
Revoked = true;
|
|
break;
|
|
case MoShiroMethods.UpdateOffer:
|
|
UpdateOffers.Enqueue(message.Params!.Value.Deserialize(MoShiroJsonContext.Default.UpdateOfferParams)!);
|
|
break;
|
|
case MoShiroMethods.Schedule:
|
|
Schedules.Enqueue(message.Params!.Value.Deserialize(MoShiroJsonContext.Default.ScheduleParams)!);
|
|
break;
|
|
case MoShiroMethods.UpdateApply:
|
|
Interlocked.Increment(ref _updateApplies);
|
|
break;
|
|
case MoShiroMethods.TestRun:
|
|
TestRunParams run = message.Params!.Value.Deserialize(MoShiroJsonContext.Default.TestRunParams)!;
|
|
TestRuns.Enqueue(run);
|
|
(bool accepted, string? errorCode) = TestRunHandler(run);
|
|
if (accepted)
|
|
{
|
|
await SendAsync(RpcMessage.Success(message.Id!.Value, new TestRunResult { Accepted = true }, MoShiroJsonContext.Default.TestRunResult));
|
|
}
|
|
else
|
|
{
|
|
await SendAsync(RpcMessage.Failure(message.Id, new FleetError(new ErrorCode(errorCode ?? "test_busy"), "refused")));
|
|
}
|
|
|
|
break;
|
|
case MoShiroMethods.LogsFetch:
|
|
LogsFetchParams fetch = message.Params!.Value.Deserialize(MoShiroJsonContext.Default.LogsFetchParams)!;
|
|
LogFetches.Enqueue(fetch);
|
|
var page = new LogsFetchResult { LastId = Logs.IsEmpty ? 0 : Logs.Max(l => l.Id) };
|
|
page.Items.AddRange(Logs.Where(l => l.Id > fetch.AfterId).OrderBy(l => l.Id).Take(Math.Max(1, fetch.Limit)));
|
|
await SendAsync(RpcMessage.Success(message.Id!.Value, page, MoShiroJsonContext.Default.LogsFetchResult));
|
|
break;
|
|
case MoShiroMethods.TestCancel:
|
|
TestCancelParams cancel = message.Params!.Value.Deserialize(MoShiroJsonContext.Default.TestCancelParams)!;
|
|
TestCancels.Enqueue(cancel.TestId);
|
|
await SendAsync(RpcMessage.Success(message.Id!.Value, new TestRunResult { Accepted = true }, MoShiroJsonContext.Default.TestRunResult));
|
|
break;
|
|
}
|
|
}
|
|
}
|