feat: add per-account connection/subscription limits with AccountConfig

This commit is contained in:
Joseph Doherty
2026-02-23 00:46:16 -05:00
parent cc0fe04f3c
commit 6afe11ad4d
6 changed files with 97 additions and 4 deletions

View File

@@ -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();
}

View 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; }
}

View File

@@ -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

View File

@@ -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; }

View File

@@ -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)

View File

@@ -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
}
}