79 lines
2.2 KiB
C#
79 lines
2.2 KiB
C#
namespace NATS.Server.Auth;
|
|
|
|
/// <summary>
|
|
/// Tracks reply subjects that a client is temporarily allowed to publish to.
|
|
/// Reference: Go client.go resp struct, setResponsePermissionIfNeeded.
|
|
/// </summary>
|
|
public sealed class ResponseTracker
|
|
{
|
|
private readonly int _maxMsgs; // 0 = unlimited
|
|
private readonly TimeSpan _expires; // TimeSpan.Zero = no TTL
|
|
private readonly Dictionary<string, (DateTime RegisteredAt, int Count)> _replies = new(StringComparer.Ordinal);
|
|
private readonly object _lock = new();
|
|
|
|
public ResponseTracker(int maxMsgs, TimeSpan expires)
|
|
{
|
|
_maxMsgs = maxMsgs;
|
|
_expires = expires;
|
|
}
|
|
|
|
public int Count
|
|
{
|
|
get { lock (_lock) return _replies.Count; }
|
|
}
|
|
|
|
public void RegisterReply(string replySubject)
|
|
{
|
|
lock (_lock)
|
|
{
|
|
_replies[replySubject] = (DateTime.UtcNow, 0);
|
|
}
|
|
}
|
|
|
|
public bool IsReplyAllowed(string subject)
|
|
{
|
|
lock (_lock)
|
|
{
|
|
if (!_replies.TryGetValue(subject, out var entry))
|
|
return false;
|
|
|
|
if (_expires > TimeSpan.Zero && DateTime.UtcNow - entry.RegisteredAt > _expires)
|
|
{
|
|
_replies.Remove(subject);
|
|
return false;
|
|
}
|
|
|
|
var newCount = entry.Count + 1;
|
|
if (_maxMsgs > 0 && newCount > _maxMsgs)
|
|
{
|
|
_replies.Remove(subject);
|
|
return false;
|
|
}
|
|
|
|
_replies[subject] = (entry.RegisteredAt, newCount);
|
|
return true;
|
|
}
|
|
}
|
|
|
|
public void Prune()
|
|
{
|
|
lock (_lock)
|
|
{
|
|
if (_expires <= TimeSpan.Zero && _maxMsgs <= 0)
|
|
return;
|
|
|
|
var now = DateTime.UtcNow;
|
|
var toRemove = new List<string>();
|
|
foreach (var (key, entry) in _replies)
|
|
{
|
|
if (_expires > TimeSpan.Zero && now - entry.RegisteredAt > _expires)
|
|
toRemove.Add(key);
|
|
else if (_maxMsgs > 0 && entry.Count >= _maxMsgs)
|
|
toRemove.Add(key);
|
|
}
|
|
foreach (var key in toRemove)
|
|
_replies.Remove(key);
|
|
}
|
|
}
|
|
}
|