PeerConnectiveChat/PCC.Shared/App/Security/TimestampNonceManager.cs

using System.Buffers.Binary;
using System.Collections.Concurrent;
using System.Security.Cryptography;

namespace PCC.App.Security;

public class TimestampNonceManager : IDisposable
{
    public static int TimestampLength => sizeof(long);

    private readonly ConcurrentDictionary<string, DateTime> _holds = new();
    private readonly TimeSpan _expiration;
    private readonly TimeSpan _maxTimeSkew;
    private readonly Timer _cleanupTimer;

    private readonly int _randomLength;

    public int NonceLength { get; }

    private bool _disposed;

    public TimestampNonceManager(int randomLength, TimeSpan expiration, TimeSpan maxTimeSkew)
    {
        NonceLength = TimestampLength + randomLength;
        _randomLength = randomLength;
        _expiration = expiration;
        _maxTimeSkew = maxTimeSkew;
        _cleanupTimer = new Timer(CleanupTimerCallback, null, expiration, expiration);
    }

    /// <summary> 生成一个新的nonce </summary>
    public (ReadOnlyMemory<byte> nonce, DateTimeOffset timestamp) NewNonce()
    {
        ThrowIfDisposed();

        var nonce = new byte[NonceLength];

        // 把时间戳放在nonce开头
        var timestamp = DateTimeOffset.UtcNow;
        BinaryPrimitives.TryWriteInt64LittleEndian(nonce, timestamp.ToUnixTimeMilliseconds());

        // 剩下的字节作为随机 nonce
        RandomNumberGenerator.Fill(nonce.AsSpan(TimestampLength));

        return (new ReadOnlyMemory<byte>(nonce), timestamp);
    }

    /// <summary> 生成一个新的nonce并追加payload </summary>
    public (ReadOnlyMemory<byte> payloadAndNonce, DateTimeOffset timestamp) NewNonceWithPayload(ReadOnlyMemory<byte> payload)
    {
        ThrowIfDisposed();
        var payloadAndNonce = new byte[NonceLength + payload.Length];

        // 把时间戳放在nonce开头
        var timestamp = DateTimeOffset.UtcNow;
        BinaryPrimitives.TryWriteInt64LittleEndian(payloadAndNonce, timestamp.ToUnixTimeMilliseconds());

        // 剩下的字节作为随机 nonce
        RandomNumberGenerator.Fill(payloadAndNonce.AsSpan(TimestampLength, _randomLength));

        // 追加 payload
        payload.Span.CopyTo(payloadAndNonce.AsSpan(NonceLength));

        return (new ReadOnlyMemory<byte>(payloadAndNonce), timestamp);
    }

    public TimestampNonceResult CheckValid(ReadOnlyMemory<byte> timestampNonce, out DateTimeOffset timestamp)
    {
        ThrowIfDisposed();

        if (timestampNonce.Length != NonceLength) throw new ArgumentException("invalid length", nameof(timestampNonce));

        timestamp = DateTimeOffset.FromUnixTimeMilliseconds(BinaryPrimitives.ReadInt64LittleEndian(timestampNonce.Span[..TimestampLength]));

        // 检查时间戳是否在允许的时间窗口内,超过最大时间差则拒绝
        if (Math.Abs((DateTimeOffset.UtcNow - timestamp).TotalMilliseconds) > _maxTimeSkew.TotalMilliseconds) return TimestampNonceResult.TimestampSkew;

        // 剩下的部分为真正的随机 nonce
        var nonce = Convert.ToHexString(timestampNonce[TimestampLength..].Span);

        // 如果 nonce 已存在，则为重放攻击
        return _holds.TryAdd(nonce, DateTime.UtcNow)
            ? TimestampNonceResult.OK
            : TimestampNonceResult.ReplayAttackDetected;
    }

    /// <summary> 验证nonce并提取payload </summary>
    public (TimestampNonceResult, DateTimeOffset timestamp, ReadOnlyMemory<byte> payload) CheckValidAndExtractPayload(ReadOnlyMemory<byte> payloadAndNonce)
    {
        ThrowIfDisposed();

        if (payloadAndNonce.Length < NonceLength) throw new ArgumentException("invalid length", nameof(payloadAndNonce));

        //提取时间戳、nonce、payload
        var timestamp = DateTimeOffset.FromUnixTimeMilliseconds(BinaryPrimitives.ReadInt64LittleEndian(payloadAndNonce.Span[..TimestampLength]));
        var nonce = Convert.ToHexString(payloadAndNonce.Slice(TimestampLength, _randomLength).Span);
        var payload = payloadAndNonce[NonceLength..];

        // 检查时间戳是否在允许的时间窗口内, 超过最大时间差则拒绝
        if (Math.Abs((DateTimeOffset.UtcNow - timestamp).TotalMilliseconds) > _maxTimeSkew.TotalMilliseconds) return (TimestampNonceResult.TimestampSkew, timestamp, payload);

        // 如果 nonce 已存在，则为重放攻击
        return (_holds.TryAdd(nonce, DateTime.UtcNow) ? TimestampNonceResult.OK : TimestampNonceResult.ReplayAttackDetected, timestamp, payload);
    }

    private void CleanupTimerCallback(object? state)
    {
        var now = DateTime.UtcNow;
        foreach (var item in _holds.Where(p => now - p.Value >= _expiration))
        {
            //TODO: debug trace logging 
            _holds.TryRemove(item.Key, out _);
        }
    }

    private void ThrowIfDisposed() => ObjectDisposedException.ThrowIf(_disposed, this);

    public void Dispose()
    {
        if (_disposed) return;
        _cleanupTimer.Dispose();
        _disposed = true;
        GC.SuppressFinalize(this);
    }
}

public enum TimestampNonceResult
{
    OK = 0,
    TimestampSkew,
    ReplayAttackDetected,
}