feat(batch27): implement jetstream bootstrap and account wiring
This commit is contained in:
Joseph Doherty
427
dotnet/src/ZB.MOM.NatsNet.Server/NatsServer.JetStreamCore.cs
Normal file
427
dotnet/src/ZB.MOM.NatsNet.Server/NatsServer.JetStreamCore.cs
Normal file
@@ -0,0 +1,427 @@
|
||||
using System.Security.Cryptography;
|
||||
using System.Text;
|
||||
using ZB.MOM.NatsNet.Server.Internal.DataStructures;
|
||||
|
||||
namespace ZB.MOM.NatsNet.Server;
|
||||
|
||||
public sealed partial class NatsServer
|
||||
{
|
||||
private const string JetStreamStoreDir = "jetstream";
|
||||
|
||||
public Exception? EnableJetStream(JetStreamConfig? config)
|
||||
{
|
||||
if (JetStreamEnabled())
|
||||
return new InvalidOperationException("jetstream already enabled");
|
||||
|
||||
Noticef("Starting JetStream");
|
||||
|
||||
if (config == null || config.MaxMemory <= 0 || config.MaxStore <= 0)
|
||||
{
|
||||
config = new JetStreamConfig
|
||||
{
|
||||
StoreDir = string.IsNullOrWhiteSpace(GetOpts().StoreDir)
|
||||
? Path.Combine(Path.GetTempPath(), JetStreamStoreDir)
|
||||
: Path.Combine(GetOpts().StoreDir, JetStreamStoreDir),
|
||||
MaxMemory = GetOpts().JetStreamMaxMemory > 0 ? GetOpts().JetStreamMaxMemory : 1,
|
||||
MaxStore = GetOpts().JetStreamMaxStore > 0 ? GetOpts().JetStreamMaxStore : 1,
|
||||
SyncInterval = GetOpts().SyncInterval,
|
||||
SyncAlways = GetOpts().SyncAlways,
|
||||
Domain = GetOpts().JetStreamDomain,
|
||||
};
|
||||
}
|
||||
else if (!string.IsNullOrWhiteSpace(config.StoreDir))
|
||||
{
|
||||
config.StoreDir = Path.Combine(config.StoreDir, JetStreamStoreDir);
|
||||
}
|
||||
|
||||
if (string.IsNullOrWhiteSpace(config.StoreDir))
|
||||
{
|
||||
config.StoreDir = Path.Combine(Path.GetTempPath(), JetStreamStoreDir);
|
||||
Warnf("Temporary storage directory used, data could be lost on system reboot");
|
||||
}
|
||||
|
||||
var err = CheckStoreDir(config);
|
||||
if (err != null)
|
||||
return err;
|
||||
|
||||
return EnableJetStreamInternal(config);
|
||||
}
|
||||
|
||||
private KeyGen? JsKeyGen(string jsKey, string info)
|
||||
{
|
||||
if (string.IsNullOrEmpty(jsKey))
|
||||
return null;
|
||||
|
||||
return context =>
|
||||
{
|
||||
using var hmac = new HMACSHA256(Encoding.UTF8.GetBytes(jsKey));
|
||||
hmac.TransformBlock(Encoding.UTF8.GetBytes(info), 0, info.Length, null, 0);
|
||||
hmac.TransformFinalBlock(context, 0, context.Length);
|
||||
return hmac.Hash ?? [];
|
||||
};
|
||||
}
|
||||
|
||||
internal (byte[]? Plain, bool UsedFallback, Exception? Error) DecryptMeta(
|
||||
StoreCipher storeCipher,
|
||||
byte[] encryptedKey,
|
||||
byte[] encryptedBuffer,
|
||||
string accountName,
|
||||
string context)
|
||||
{
|
||||
if (encryptedKey.Length == 0)
|
||||
return (null, false, new InvalidOperationException("encryption key missing"));
|
||||
|
||||
var ciphers = storeCipher == StoreCipher.Aes
|
||||
? new[] { StoreCipher.Aes, StoreCipher.ChaCha }
|
||||
: new[] { StoreCipher.ChaCha, StoreCipher.Aes };
|
||||
|
||||
var candidates = new List<(KeyGen Prf, StoreCipher Cipher)>();
|
||||
var opts = GetOpts();
|
||||
|
||||
var prf = JsKeyGen(opts.JetStreamKey, accountName);
|
||||
if (prf == null)
|
||||
return (null, false, new InvalidOperationException("jetstream encryption key is not configured"));
|
||||
|
||||
foreach (var cipher in ciphers)
|
||||
candidates.Add((prf, cipher));
|
||||
|
||||
var oldPrf = JsKeyGen(opts.JetStreamOldKey, accountName);
|
||||
if (oldPrf != null)
|
||||
{
|
||||
foreach (var cipher in ciphers)
|
||||
candidates.Add((oldPrf, cipher));
|
||||
}
|
||||
|
||||
for (var i = 0; i < candidates.Count; i++)
|
||||
{
|
||||
try
|
||||
{
|
||||
var rb = candidates[i].Prf(Encoding.UTF8.GetBytes(context));
|
||||
var kek = JetStreamFileStore.GenEncryptionKey(candidates[i].Cipher, rb);
|
||||
var ns = kek.NonceSize;
|
||||
if (encryptedKey.Length < ns || encryptedBuffer.Length < ns)
|
||||
continue;
|
||||
|
||||
var seed = kek.Open(encryptedKey.AsSpan(0, ns), encryptedKey.AsSpan(ns));
|
||||
var aek = JetStreamFileStore.GenEncryptionKey(candidates[i].Cipher, seed);
|
||||
var plain = aek.Open(encryptedBuffer.AsSpan(0, ns), encryptedBuffer.AsSpan(ns));
|
||||
return (plain, i > 0, null);
|
||||
}
|
||||
catch
|
||||
{
|
||||
// Try the next candidate.
|
||||
}
|
||||
}
|
||||
|
||||
return (null, false, new InvalidOperationException("unable to recover encrypted metadata"));
|
||||
}
|
||||
|
||||
internal Exception? CheckStoreDir(JetStreamConfig cfg)
|
||||
{
|
||||
if (string.IsNullOrWhiteSpace(cfg.StoreDir))
|
||||
return new InvalidOperationException("jetstream store directory is required");
|
||||
|
||||
try
|
||||
{
|
||||
Directory.CreateDirectory(cfg.StoreDir);
|
||||
return null;
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
return ex;
|
||||
}
|
||||
}
|
||||
|
||||
internal Exception? InitJetStreamEncryption()
|
||||
{
|
||||
var opts = GetOpts();
|
||||
|
||||
if (!string.IsNullOrEmpty(opts.JetStreamKey) && !string.IsNullOrEmpty(opts.JetStreamTpm.KeysFile))
|
||||
return new InvalidOperationException("JetStream encryption key may not be used with TPM options");
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
private Exception? EnableJetStreamInternal(JetStreamConfig cfg)
|
||||
{
|
||||
var encryptionErr = InitJetStreamEncryption();
|
||||
if (encryptionErr != null)
|
||||
return encryptionErr;
|
||||
|
||||
try
|
||||
{
|
||||
Directory.CreateDirectory(cfg.StoreDir);
|
||||
}
|
||||
catch (Exception ex)
|
||||
{
|
||||
return ex;
|
||||
}
|
||||
|
||||
var js = new JetStream
|
||||
{
|
||||
Server = this,
|
||||
Config = cfg,
|
||||
Started = DateTime.UtcNow,
|
||||
StandAlone = true,
|
||||
};
|
||||
|
||||
_mu.EnterWriteLock();
|
||||
try
|
||||
{
|
||||
_jetStream = js;
|
||||
_info.JetStream = true;
|
||||
_info.Domain = cfg.Domain;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_mu.ExitWriteLock();
|
||||
}
|
||||
|
||||
var err = EnableJetStreamAccounts();
|
||||
if (err != null)
|
||||
{
|
||||
_mu.EnterWriteLock();
|
||||
try
|
||||
{
|
||||
_jetStream = null;
|
||||
_info.JetStream = false;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_mu.ExitWriteLock();
|
||||
}
|
||||
}
|
||||
|
||||
return err;
|
||||
}
|
||||
|
||||
internal bool CanExtendOtherDomain()
|
||||
{
|
||||
var opts = GetOpts();
|
||||
var sysAcc = SystemAccount()?.GetName();
|
||||
if (string.IsNullOrEmpty(sysAcc))
|
||||
return false;
|
||||
|
||||
foreach (var remote in opts.LeafNode.Remotes)
|
||||
{
|
||||
if (!string.Equals(remote.LocalAccount, sysAcc, StringComparison.Ordinal))
|
||||
continue;
|
||||
|
||||
foreach (var denyImport in remote.DenyImports)
|
||||
{
|
||||
if (SubscriptionIndex.SubjectIsSubsetMatch(denyImport, JsApiSubjects.JsAllApi))
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
internal void UpdateJetStreamInfoStatus(bool enabled)
|
||||
{
|
||||
_mu.EnterWriteLock();
|
||||
try
|
||||
{
|
||||
_info.JetStream = enabled;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_mu.ExitWriteLock();
|
||||
}
|
||||
}
|
||||
|
||||
internal Exception? RestartJetStream()
|
||||
{
|
||||
var opts = GetOpts();
|
||||
var cfg = new JetStreamConfig
|
||||
{
|
||||
StoreDir = opts.StoreDir,
|
||||
SyncInterval = opts.SyncInterval,
|
||||
SyncAlways = opts.SyncAlways,
|
||||
MaxMemory = opts.JetStreamMaxMemory,
|
||||
MaxStore = opts.JetStreamMaxStore,
|
||||
Domain = opts.JetStreamDomain,
|
||||
Strict = !opts.NoJetStreamStrict,
|
||||
};
|
||||
|
||||
Noticef("Restarting JetStream");
|
||||
var err = EnableJetStream(cfg);
|
||||
if (err != null)
|
||||
{
|
||||
Warnf("Can't start JetStream: {0}", err.Message);
|
||||
_ = DisableJetStream();
|
||||
return err;
|
||||
}
|
||||
|
||||
UpdateJetStreamInfoStatus(true);
|
||||
return null;
|
||||
}
|
||||
|
||||
internal void CheckJetStreamExports()
|
||||
{
|
||||
if (SystemAccount() != null)
|
||||
SetupJetStreamExports();
|
||||
}
|
||||
|
||||
internal void SetupJetStreamExports()
|
||||
{
|
||||
var sys = SystemAccount();
|
||||
if (sys == null)
|
||||
return;
|
||||
|
||||
var err = sys.AddServiceExport(JsApiSubjects.JsAllApi, null);
|
||||
if (err != null)
|
||||
Warnf("Error setting up jetstream service exports: {0}", err.Message);
|
||||
}
|
||||
|
||||
internal bool JetStreamOOSPending()
|
||||
{
|
||||
var js = _jetStream;
|
||||
if (js == null)
|
||||
return false;
|
||||
|
||||
js.Lock.EnterWriteLock();
|
||||
try
|
||||
{
|
||||
var wasPending = js.Oos;
|
||||
js.Oos = true;
|
||||
return wasPending;
|
||||
}
|
||||
finally
|
||||
{
|
||||
js.Lock.ExitWriteLock();
|
||||
}
|
||||
}
|
||||
|
||||
internal void SetJetStreamDisabled()
|
||||
{
|
||||
var js = _jetStream;
|
||||
if (js != null)
|
||||
Interlocked.Exchange(ref js.Disabled, 1);
|
||||
}
|
||||
|
||||
internal void HandleOutOfSpace(NatsStream? stream)
|
||||
{
|
||||
if (!JetStreamEnabled() || JetStreamOOSPending())
|
||||
return;
|
||||
|
||||
if (stream != null)
|
||||
Errorf("JetStream out of resources for stream {0}, will be DISABLED", stream.Config.Name);
|
||||
else
|
||||
Errorf("JetStream out of resources, will be DISABLED");
|
||||
|
||||
_ = Task.Run(() => DisableJetStream());
|
||||
}
|
||||
|
||||
public Exception? DisableJetStream()
|
||||
{
|
||||
if (!JetStreamEnabled())
|
||||
return null;
|
||||
|
||||
SetJetStreamDisabled();
|
||||
UpdateJetStreamInfoStatus(false);
|
||||
|
||||
_mu.EnterWriteLock();
|
||||
try
|
||||
{
|
||||
_jetStream = null;
|
||||
}
|
||||
finally
|
||||
{
|
||||
_mu.ExitWriteLock();
|
||||
}
|
||||
|
||||
ShutdownJetStream();
|
||||
ShutdownRaftNodes();
|
||||
return null;
|
||||
}
|
||||
|
||||
private Exception? EnableJetStreamAccounts()
|
||||
{
|
||||
if (GlobalAccountOnly())
|
||||
{
|
||||
var gacc = GlobalAccount();
|
||||
if (gacc == null)
|
||||
return new InvalidOperationException("global account not found");
|
||||
|
||||
gacc.JetStreamLimits ??= new Dictionary<string, object>(StringComparer.Ordinal)
|
||||
{
|
||||
[string.Empty] = new JetStreamAccountLimits
|
||||
{
|
||||
MaxMemory = -1,
|
||||
MaxStore = -1,
|
||||
MaxStreams = -1,
|
||||
MaxConsumers = -1,
|
||||
MaxAckPending = -1,
|
||||
MemoryMaxStreamBytes = -1,
|
||||
StoreMaxStreamBytes = -1,
|
||||
},
|
||||
};
|
||||
|
||||
return ConfigJetStream(gacc);
|
||||
}
|
||||
|
||||
return ConfigAllJetStreamAccounts();
|
||||
}
|
||||
|
||||
internal Exception? ConfigJetStream(Account? acc)
|
||||
{
|
||||
if (acc == null)
|
||||
return null;
|
||||
|
||||
var jsLimits = acc.JetStreamLimits;
|
||||
if (jsLimits != null)
|
||||
return acc.EnableAllJetStreamServiceImportsAndMappings();
|
||||
|
||||
if (!ReferenceEquals(acc, SystemAccount()))
|
||||
{
|
||||
acc.JetStream = null;
|
||||
return acc.EnableJetStreamInfoServiceImportOnly();
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
internal Exception? ConfigAllJetStreamAccounts()
|
||||
{
|
||||
CheckJetStreamExports();
|
||||
if (_jetStream == null)
|
||||
return null;
|
||||
|
||||
foreach (var acc in _accounts.Values)
|
||||
{
|
||||
var err = ConfigJetStream(acc);
|
||||
if (err != null)
|
||||
return err;
|
||||
}
|
||||
|
||||
var storeDir = _jetStream.Config.StoreDir;
|
||||
if (!Directory.Exists(storeDir))
|
||||
return null;
|
||||
|
||||
foreach (var directory in Directory.EnumerateDirectories(storeDir))
|
||||
{
|
||||
var accountName = Path.GetFileName(directory);
|
||||
if (string.IsNullOrWhiteSpace(accountName) || _accounts.ContainsKey(accountName))
|
||||
continue;
|
||||
|
||||
var (resolved, _) = LookupAccount(accountName);
|
||||
if (resolved == null)
|
||||
continue;
|
||||
|
||||
var err = ConfigJetStream(resolved);
|
||||
if (err != null)
|
||||
return err;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
public bool JetStreamEnabled()
|
||||
{
|
||||
var js = _jetStream;
|
||||
return js != null && Interlocked.CompareExchange(ref js.Disabled, 0, 0) == 0;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user