feat: add per-account connection/subscription limits with AccountConfig
This commit is contained in:
@@ -10,8 +10,11 @@ public sealed class Account : IDisposable
|
||||
public string Name { get; }
|
||||
public SubList SubList { get; } = new();
|
||||
public Permissions? DefaultPermissions { get; set; }
|
||||
public int MaxConnections { get; set; } // 0 = unlimited
|
||||
public int MaxSubscriptions { get; set; } // 0 = unlimited
|
||||
|
||||
private readonly ConcurrentDictionary<ulong, byte> _clients = new();
|
||||
private int _subscriptionCount;
|
||||
|
||||
public Account(string name)
|
||||
{
|
||||
@@ -19,10 +22,31 @@ public sealed class Account : IDisposable
|
||||
}
|
||||
|
||||
public int ClientCount => _clients.Count;
|
||||
public int SubscriptionCount => Volatile.Read(ref _subscriptionCount);
|
||||
|
||||
public void AddClient(ulong clientId) => _clients[clientId] = 0;
|
||||
/// <summary>Returns false if max connections exceeded.</summary>
|
||||
public bool AddClient(ulong clientId)
|
||||
{
|
||||
if (MaxConnections > 0 && _clients.Count >= MaxConnections)
|
||||
return false;
|
||||
_clients[clientId] = 0;
|
||||
return true;
|
||||
}
|
||||
|
||||
public void RemoveClient(ulong clientId) => _clients.TryRemove(clientId, out _);
|
||||
|
||||
public bool IncrementSubscriptions()
|
||||
{
|
||||
if (MaxSubscriptions > 0 && Volatile.Read(ref _subscriptionCount) >= MaxSubscriptions)
|
||||
return false;
|
||||
Interlocked.Increment(ref _subscriptionCount);
|
||||
return true;
|
||||
}
|
||||
|
||||
public void DecrementSubscriptions()
|
||||
{
|
||||
Interlocked.Decrement(ref _subscriptionCount);
|
||||
}
|
||||
|
||||
public void Dispose() => SubList.Dispose();
|
||||
}
|
||||
|
||||
8
src/NATS.Server/Auth/AccountConfig.cs
Normal file
8
src/NATS.Server/Auth/AccountConfig.cs
Normal file
@@ -0,0 +1,8 @@
|
||||
namespace NATS.Server.Auth;
|
||||
|
||||
public sealed class AccountConfig
|
||||
{
|
||||
public int MaxConnections { get; init; } // 0 = unlimited
|
||||
public int MaxSubscriptions { get; init; } // 0 = unlimited
|
||||
public Permissions? DefaultPermissions { get; init; }
|
||||
}
|
||||
@@ -372,7 +372,13 @@ public sealed class NatsClient : IDisposable
|
||||
{
|
||||
var accountName = result.AccountName ?? Account.GlobalAccountName;
|
||||
Account = server.GetOrCreateAccount(accountName);
|
||||
Account.AddClient(Id);
|
||||
if (!Account.AddClient(Id))
|
||||
{
|
||||
Account = null;
|
||||
await SendErrAndCloseAsync("maximum connections for account exceeded",
|
||||
ClientClosedReason.AuthenticationViolation);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
_logger.LogDebug("Client {ClientId} authenticated as {Identity}", Id, result.Identity);
|
||||
@@ -386,7 +392,13 @@ public sealed class NatsClient : IDisposable
|
||||
if (Account == null && Router is NatsServer server2)
|
||||
{
|
||||
Account = server2.GetOrCreateAccount(Account.GlobalAccountName);
|
||||
Account.AddClient(Id);
|
||||
if (!Account.AddClient(Id))
|
||||
{
|
||||
Account = null;
|
||||
await SendErrAndCloseAsync("maximum connections for account exceeded",
|
||||
ClientClosedReason.AuthenticationViolation);
|
||||
return;
|
||||
}
|
||||
}
|
||||
|
||||
// Validate no_responders requires headers
|
||||
|
||||
@@ -29,6 +29,9 @@ public sealed class NatsOptions
|
||||
// Server tags (exposed via /varz)
|
||||
public Dictionary<string, string>? Tags { get; set; }
|
||||
|
||||
// Account configuration
|
||||
public Dictionary<string, AccountConfig>? Accounts { get; set; }
|
||||
|
||||
// Simple auth (single user)
|
||||
public string? Username { get; set; }
|
||||
public string? Password { get; set; }
|
||||
|
||||
@@ -569,7 +569,17 @@ public sealed class NatsServer : IMessageRouter, ISubListAccess, IDisposable
|
||||
|
||||
public Account GetOrCreateAccount(string name)
|
||||
{
|
||||
return _accounts.GetOrAdd(name, n => new Account(n));
|
||||
return _accounts.GetOrAdd(name, n =>
|
||||
{
|
||||
var acc = new Account(n);
|
||||
if (_options.Accounts != null && _options.Accounts.TryGetValue(n, out var config))
|
||||
{
|
||||
acc.MaxConnections = config.MaxConnections;
|
||||
acc.MaxSubscriptions = config.MaxSubscriptions;
|
||||
acc.DefaultPermissions = config.DefaultPermissions;
|
||||
}
|
||||
return acc;
|
||||
});
|
||||
}
|
||||
|
||||
public void RemoveClient(NatsClient client)
|
||||
|
||||
@@ -32,4 +32,40 @@ public class AccountTests
|
||||
{
|
||||
Account.GlobalAccountName.ShouldBe("$G");
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Account_enforces_max_connections()
|
||||
{
|
||||
var acc = new Account("test") { MaxConnections = 2 };
|
||||
acc.AddClient(1).ShouldBeTrue();
|
||||
acc.AddClient(2).ShouldBeTrue();
|
||||
acc.AddClient(3).ShouldBeFalse(); // exceeds limit
|
||||
acc.ClientCount.ShouldBe(2);
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Account_unlimited_connections_when_zero()
|
||||
{
|
||||
var acc = new Account("test") { MaxConnections = 0 };
|
||||
acc.AddClient(1).ShouldBeTrue();
|
||||
acc.AddClient(2).ShouldBeTrue();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Account_enforces_max_subscriptions()
|
||||
{
|
||||
var acc = new Account("test") { MaxSubscriptions = 2 };
|
||||
acc.IncrementSubscriptions().ShouldBeTrue();
|
||||
acc.IncrementSubscriptions().ShouldBeTrue();
|
||||
acc.IncrementSubscriptions().ShouldBeFalse();
|
||||
}
|
||||
|
||||
[Fact]
|
||||
public void Account_decrement_subscriptions()
|
||||
{
|
||||
var acc = new Account("test") { MaxSubscriptions = 1 };
|
||||
acc.IncrementSubscriptions().ShouldBeTrue();
|
||||
acc.DecrementSubscriptions();
|
||||
acc.IncrementSubscriptions().ShouldBeTrue(); // slot freed
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user