343 lines
12 KiB
C#
343 lines
12 KiB
C#
using System.Collections.Concurrent;
|
||
using System.Net;
|
||
using System.Net.Security;
|
||
using System.Net.Sockets;
|
||
using System.Security.Authentication;
|
||
using System.Security.Cryptography;
|
||
using System.Security.Cryptography.X509Certificates;
|
||
using MoShiro.Protocol;
|
||
|
||
namespace MoShiro.Measurement;
|
||
|
||
/// <summary>拨号方的一条数据面 TCP 连接（已经完成 TLS 和首帧认证）。</summary>
|
||
public sealed class TcpLink : IDisposable
|
||
{
|
||
private readonly TcpClient _client;
|
||
|
||
private TcpLink(TcpClient client, Stream stream, long tcpUs, long? tlsUs)
|
||
{
|
||
_client = client;
|
||
Socket = client.Client;
|
||
Stream = stream;
|
||
TcpUs = tcpUs;
|
||
TlsUs = tlsUs;
|
||
}
|
||
|
||
public Stream Stream { get; }
|
||
|
||
/// <summary>底层套接字（内核采样用）。</summary>
|
||
public Socket Socket { get; }
|
||
|
||
public long TcpUs { get; }
|
||
|
||
public long? TlsUs { get; }
|
||
|
||
/// <summary>
|
||
/// 连接 → （TLS）→ 认证首帧。<paramref name="sendAuth"/> 为 false 时只测量建立连接的耗时（connect）。
|
||
/// 对端证书按 <see cref="BlockEnv.PeerSpkiSha256"/> 固定；没有指纹（目标任务，目标是任意地址）就只测握手耗时、不校验证书。
|
||
/// 失败时抛出 <see cref="SocketException"/>、<see cref="AuthenticationException"/>、<see cref="IOException"/> 或超时的 <see cref="OperationCanceledException"/>。
|
||
/// </summary>
|
||
public static async Task<TcpLink> OpenAsync(BlockEnv env, int port, bool tls, uint streamId, bool sendAuth, CancellationToken ct)
|
||
{
|
||
var client = new TcpClient(env.PeerAddress!.AddressFamily) { NoDelay = true };
|
||
try
|
||
{
|
||
long start = env.Clock.NowUs();
|
||
await client.ConnectAsync(env.PeerAddress!, port, ct);
|
||
long tcpUs = env.Clock.NowUs() - start;
|
||
Stream stream = client.GetStream();
|
||
long? tlsUs = null;
|
||
if (tls)
|
||
{
|
||
long tlsStart = env.Clock.NowUs();
|
||
byte[]? pin = env.PeerSpkiSha256;
|
||
#pragma warning disable CA5359 // 没有指纹时只测量握手耗时（目标任务的目标是任意地址）
|
||
var ssl = new SslStream(stream, false, (_, certificate, _, _) => pin is null || (certificate is not null && Pinned(certificate, pin)));
|
||
#pragma warning restore CA5359
|
||
await ssl.AuthenticateAsClientAsync(new SslClientAuthenticationOptions { TargetHost = env.PeerAddress!.ToString() }, ct);
|
||
stream = ssl;
|
||
tlsUs = env.Clock.NowUs() - tlsStart;
|
||
}
|
||
|
||
if (sendAuth)
|
||
{
|
||
DataSession session = env.Session() ?? throw new InvalidOperationException("no data session");
|
||
byte[] auth = new byte[TcpWire.AuthBytes];
|
||
Span<byte> nonce = stackalloc byte[TcpWire.NonceBytes];
|
||
RandomNumberGenerator.Fill(nonce);
|
||
TcpWire.WriteAuth(auth, session.SessionId, streamId, nonce, session.Key);
|
||
await stream.WriteAsync(auth, ct);
|
||
}
|
||
|
||
return new TcpLink(client, stream, tcpUs, tlsUs);
|
||
}
|
||
catch
|
||
{
|
||
client.Dispose();
|
||
throw;
|
||
}
|
||
}
|
||
|
||
public void Dispose()
|
||
{
|
||
Stream.Dispose();
|
||
_client.Dispose();
|
||
}
|
||
|
||
private static bool Pinned(X509Certificate certificate, byte[] expected)
|
||
{
|
||
X509Certificate2 full = certificate as X509Certificate2 ?? X509CertificateLoader.LoadCertificate(certificate.GetRawCertData());
|
||
return CryptographicOperations.FixedTimeEquals(SHA256.HashData(full.PublicKey.ExportSubjectPublicKeyInfo()), expected);
|
||
}
|
||
}
|
||
|
||
/// <summary>监听方收到的一条已经认证的数据面 TCP 连接。</summary>
|
||
public sealed class TcpConnection
|
||
{
|
||
public TcpConnection(Stream stream, Socket socket, DataSession session, uint streamId)
|
||
{
|
||
Stream = stream;
|
||
Socket = socket;
|
||
Session = session;
|
||
StreamId = streamId;
|
||
}
|
||
|
||
public Stream Stream { get; }
|
||
|
||
/// <summary>底层套接字（内核采样用）。</summary>
|
||
public Socket Socket { get; }
|
||
|
||
public DataSession Session { get; }
|
||
|
||
public uint StreamId { get; }
|
||
}
|
||
|
||
public delegate Task TcpResponder(TcpConnection connection, CancellationToken ct);
|
||
|
||
/// <summary>
|
||
/// 监听方的 TCP 数据面：每个端口一个监听器。连接的第一个字节是 0x16 就先做 TLS（节点自己的证书），否则直接是明文。
|
||
/// 随后是认证首帧（<c>"MN" | ver | session_id | stream_id | nonce | mac</c>，随机数只能用一次），通过之后交给按 <c>stream_id</c> 登记的处理函数。
|
||
/// 首帧在限定时间内没有到达或者校验失败，直接关闭连接并计数。
|
||
/// </summary>
|
||
public sealed class TcpAcceptor
|
||
{
|
||
private static readonly TimeSpan s_authTimeout = TimeSpan.FromSeconds(5);
|
||
|
||
private readonly IPAddress _bind;
|
||
private readonly IReadOnlyList<int> _ports;
|
||
private readonly X509Certificate2 _certificate;
|
||
private readonly IDataSessions _sessions;
|
||
private readonly DataPlaneCounters _counters;
|
||
private readonly ConcurrentDictionary<uint, TcpResponder> _handlers = new();
|
||
private readonly List<TcpListener> _listeners = [];
|
||
|
||
public TcpAcceptor(IPAddress bind, IReadOnlyList<int> ports, X509Certificate2 certificate, IDataSessions sessions, DataPlaneCounters counters)
|
||
{
|
||
_bind = bind;
|
||
_ports = ports;
|
||
_certificate = certificate;
|
||
_sessions = sessions;
|
||
_counters = counters;
|
||
}
|
||
|
||
public IReadOnlyList<int> BoundPorts { get; private set; } = [];
|
||
|
||
public void Register(uint streamId, TcpResponder handler) => _handlers[streamId] = handler;
|
||
|
||
public void Unregister(uint streamId) => _handlers.TryRemove(streamId, out _);
|
||
|
||
/// <summary>开始监听（端口被占用等错误在这里抛出）。</summary>
|
||
public void Bind()
|
||
{
|
||
var bound = new List<int>();
|
||
foreach (int port in _ports)
|
||
{
|
||
var listener = new TcpListener(_bind, port);
|
||
if (_bind.Equals(IPAddress.IPv6Any))
|
||
{
|
||
listener.Server.DualMode = true;
|
||
}
|
||
|
||
listener.Start(128);
|
||
_listeners.Add(listener);
|
||
bound.Add(((IPEndPoint)listener.LocalEndpoint).Port);
|
||
}
|
||
|
||
BoundPorts = bound;
|
||
}
|
||
|
||
public async Task RunAsync(CancellationToken ct)
|
||
{
|
||
try
|
||
{
|
||
await Task.WhenAll(_listeners.Select(l => AcceptLoopAsync(l, ct)));
|
||
}
|
||
finally
|
||
{
|
||
foreach (TcpListener listener in _listeners)
|
||
{
|
||
listener.Stop();
|
||
}
|
||
}
|
||
}
|
||
|
||
private async Task AcceptLoopAsync(TcpListener listener, CancellationToken ct)
|
||
{
|
||
while (!ct.IsCancellationRequested)
|
||
{
|
||
TcpClient client;
|
||
try
|
||
{
|
||
client = await listener.AcceptTcpClientAsync(ct);
|
||
}
|
||
catch (OperationCanceledException)
|
||
{
|
||
return;
|
||
}
|
||
catch (SocketException)
|
||
{
|
||
continue;
|
||
}
|
||
|
||
_ = Task.Run(() => HandleAsync(client, ct), CancellationToken.None);
|
||
}
|
||
}
|
||
|
||
private async Task HandleAsync(TcpClient client, CancellationToken ct)
|
||
{
|
||
using (client)
|
||
{
|
||
client.NoDelay = true;
|
||
try
|
||
{
|
||
Stream stream = client.GetStream();
|
||
using var authCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
|
||
authCts.CancelAfter(s_authTimeout);
|
||
|
||
byte[] first = new byte[1];
|
||
if (await stream.ReadAsync(first, authCts.Token) != 1)
|
||
{
|
||
return;
|
||
}
|
||
|
||
stream = new PrefixedStream(stream, first[0]);
|
||
if (first[0] == 0x16)
|
||
{
|
||
var ssl = new SslStream(stream, leaveInnerStreamOpen: false);
|
||
await ssl.AuthenticateAsServerAsync(new SslServerAuthenticationOptions { ServerCertificate = _certificate }, authCts.Token);
|
||
stream = ssl;
|
||
}
|
||
|
||
byte[] auth = new byte[TcpWire.AuthBytes];
|
||
await stream.ReadExactlyAsync(auth, authCts.Token);
|
||
if (!TcpWire.TryParseAuth(auth, out long sessionId, out uint streamId, out ReadOnlySpan<byte> nonce)
|
||
|| _sessions.Find(sessionId) is not { } session
|
||
|| !_handlers.TryGetValue(streamId, out TcpResponder? handler)
|
||
|| !TcpWire.VerifyAuth(auth, session.Key)
|
||
|| !session.UseNonce(nonce))
|
||
{
|
||
_counters.DropInvalid();
|
||
return;
|
||
}
|
||
|
||
await handler(new TcpConnection(stream, client.Client, session, streamId), ct);
|
||
}
|
||
catch (Exception ex) when (ex is IOException or SocketException or OperationCanceledException or ObjectDisposedException or AuthenticationException or InvalidDataException or EndOfStreamException)
|
||
{
|
||
// 连接被对端关闭、超时或者格式错误：直接结束这个连接
|
||
}
|
||
}
|
||
}
|
||
|
||
/// <summary>把已经读出的第一个字节放回流的开头（用于探测是 TLS 还是明文）。</summary>
|
||
private sealed class PrefixedStream : Stream
|
||
{
|
||
private readonly Stream _inner;
|
||
private byte? _prefix;
|
||
|
||
public PrefixedStream(Stream inner, byte prefix)
|
||
{
|
||
_inner = inner;
|
||
_prefix = prefix;
|
||
}
|
||
|
||
public override bool CanRead => true;
|
||
|
||
public override bool CanSeek => false;
|
||
|
||
public override bool CanWrite => true;
|
||
|
||
public override long Length => throw new NotSupportedException();
|
||
|
||
public override long Position
|
||
{
|
||
get => throw new NotSupportedException();
|
||
set => throw new NotSupportedException();
|
||
}
|
||
|
||
public override void Flush() => _inner.Flush();
|
||
|
||
public override Task FlushAsync(CancellationToken cancellationToken) => _inner.FlushAsync(cancellationToken);
|
||
|
||
public override int Read(byte[] buffer, int offset, int count) => Read(buffer.AsSpan(offset, count));
|
||
|
||
public override int Read(Span<byte> buffer)
|
||
{
|
||
if (buffer.IsEmpty)
|
||
{
|
||
return 0;
|
||
}
|
||
|
||
if (_prefix is { } prefix)
|
||
{
|
||
buffer[0] = prefix;
|
||
_prefix = null;
|
||
return 1;
|
||
}
|
||
|
||
return _inner.Read(buffer);
|
||
}
|
||
|
||
public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
|
||
{
|
||
if (buffer.IsEmpty)
|
||
{
|
||
return ValueTask.FromResult(0);
|
||
}
|
||
|
||
if (_prefix is { } prefix)
|
||
{
|
||
buffer.Span[0] = prefix;
|
||
_prefix = null;
|
||
return ValueTask.FromResult(1);
|
||
}
|
||
|
||
return _inner.ReadAsync(buffer, cancellationToken);
|
||
}
|
||
|
||
public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => ReadAsync(buffer.AsMemory(offset, count), cancellationToken).AsTask();
|
||
|
||
public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();
|
||
|
||
public override void SetLength(long value) => throw new NotSupportedException();
|
||
|
||
public override void Write(byte[] buffer, int offset, int count) => _inner.Write(buffer, offset, count);
|
||
|
||
public override void Write(ReadOnlySpan<byte> buffer) => _inner.Write(buffer);
|
||
|
||
public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default) => _inner.WriteAsync(buffer, cancellationToken);
|
||
|
||
public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken) => _inner.WriteAsync(buffer, offset, count, cancellationToken);
|
||
|
||
protected override void Dispose(bool disposing)
|
||
{
|
||
if (disposing)
|
||
{
|
||
_inner.Dispose();
|
||
}
|
||
|
||
base.Dispose(disposing);
|
||
}
|
||
}
|
||
}
|