From 58864dc1663a6e774a9b8e8839b2def3e5b2ff9e Mon Sep 17 00:00:00 2001
From: Garrett Sutton <garrettsutton97@gmail.com>
Date: Tue, 14 Jan 2025 17:00:11 -0600
Subject: [PATCH] NKeyHandler, JwtHandler, and SeedHandler

---
 .../Internal/UserCredentials.cs               | 24 ++++++++++++++-----
 src/NATS.Client.Core/NatsAuthOpts.cs          | 22 +++++++++++++++++
 2 files changed, 40 insertions(+), 6 deletions(-)

diff --git a/src/NATS.Client.Core/Internal/UserCredentials.cs b/src/NATS.Client.Core/Internal/UserCredentials.cs
index 7ce5f3c2..817e6e87 100644
--- a/src/NATS.Client.Core/Internal/UserCredentials.cs
+++ b/src/NATS.Client.Core/Internal/UserCredentials.cs
@@ -1,4 +1,5 @@
 using System.Text;
+using Microsoft.Extensions.Logging;
 using NATS.Client.Core.NaCl;
 
 namespace NATS.Client.Core.Internal;
@@ -8,8 +9,11 @@ internal class UserCredentials
     public UserCredentials(NatsAuthOpts authOpts)
     {
         Jwt = authOpts.Jwt;
+        JwtHandler = authOpts.JwtHandler;
         Seed = authOpts.Seed;
+        SeedHandler = authOpts.SeedHandler;
         NKey = authOpts.NKey;
+        NKeyHandler = authOpts.NKeyHandler;
         Token = authOpts.Token;
         TokenHandler = authOpts.TokenHandler;
 
@@ -26,20 +30,28 @@ public UserCredentials(NatsAuthOpts authOpts)
 
     public string? Jwt { get; }
 
+    public Func<ValueTask<string>>? JwtHandler { get; }
+
     public string? Seed { get; }
 
+    public Func<ValueTask<string>>? SeedHandler { get; }
+
     public string? NKey { get; }
 
+    public Func<ValueTask<string>>? NKeyHandler { get; }
+
     public string? Token { get; }
 
     public Func<ValueTask<string>>? TokenHandler { get; }
 
-    public string? Sign(string? nonce)
+    public async Task<string?> SignAsync(string? nonce)
     {
-        if (Seed == null || nonce == null)
+        var seed = SeedHandler != null ? await SeedHandler().ConfigureAwait(false) : Seed;
+
+        if (seed == null || nonce == null)
             return null;
 
-        using var kp = NKeys.FromSeed(Seed);
+        using var kp = NKeys.FromSeed(seed);
         var bytes = kp.Sign(Encoding.ASCII.GetBytes(nonce));
         var sig = CryptoBytes.ToBase64String(bytes);
 
@@ -48,10 +60,10 @@ public UserCredentials(NatsAuthOpts authOpts)
 
     internal async Task AuthenticateAsync(ClientOpts opts, ServerInfo? info)
     {
-        opts.JWT = Jwt;
-        opts.NKey = NKey;
+        opts.JWT = JwtHandler != null ? await JwtHandler().ConfigureAwait(false) : Jwt;
+        opts.NKey = NKeyHandler != null ? await NKeyHandler().ConfigureAwait(false) : NKey;
         opts.AuthToken = TokenHandler != null ? await TokenHandler().ConfigureAwait(false) : Token;
-        opts.Sig = info is { AuthRequired: true, Nonce: { } } ? Sign(info.Nonce) : null;
+        opts.Sig = info is { AuthRequired: true, Nonce: { } } ? await SignAsync(info.Nonce).ConfigureAwait(false) : null;
     }
 
     private (string, string) LoadCredsFile(string path)
diff --git a/src/NATS.Client.Core/NatsAuthOpts.cs b/src/NATS.Client.Core/NatsAuthOpts.cs
index 53c57dd7..9d5e2085 100644
--- a/src/NATS.Client.Core/NatsAuthOpts.cs
+++ b/src/NATS.Client.Core/NatsAuthOpts.cs
@@ -18,10 +18,28 @@ public record NatsAuthOpts
 
     public string? Jwt { get; init; }
 
+    /// <summary>
+    /// A callback that returns a JWT string.
+    /// If this is set, it takes precedence over <see cref="Jwt"/>
+    /// </summary>
+    public Func<ValueTask<string>>? JwtHandler { get; init; }
+
     public string? NKey { get; init; }
 
+    /// <summary>
+    /// A callback that returns an NKey string.
+    /// If this is set, it takes precedence over <see cref="NKey"/>.
+    /// </summary>
+    public Func<ValueTask<string>>? NKeyHandler { get; init; }
+
     public string? Seed { get; init; }
 
+    /// <summary>
+    /// A callback that returns a seed string.
+    /// If this is set, it takes precedence over <see cref="Seed"/>.
+    /// </summary>
+    public Func<ValueTask<string>>? SeedHandler { get; init; }
+
     public string? CredsFile { get; init; }
 
     public string? NKeyFile { get; init; }
@@ -31,7 +49,11 @@ public record NatsAuthOpts
                                && string.IsNullOrEmpty(Token)
                                && TokenHandler == null
                                && string.IsNullOrEmpty(Jwt)
+                               && JwtHandler == null
+                               && string.IsNullOrEmpty(NKey)
+                               && NKeyHandler == null
                                && string.IsNullOrEmpty(Seed)
+                               && SeedHandler == null
                                && string.IsNullOrEmpty(CredsFile)
                                && string.IsNullOrEmpty(NKeyFile);
 }