Add OnSocketAvailableAsync Hook

src/NATS.Client.Core/INatsConnection.cs

@@ -51,6 +51,16 @@ public interface INatsConnection : INatsClient
     /// </summary>
     NatsHeaderParser HeaderParser { get; }
 
+    /// <summary>
+    /// Hook before TCP connection open.
+    /// </summary>
+    Func<(string Host, int Port), ValueTask<(string Host, int Port)>>? OnConnectingAsync { get; set; }
+
+    /// <summary>
+    /// Hook when socket is available.
+    /// </summary>
+    Func<ISocketConnection, ValueTask<ISocketConnection>>? OnSocketAvailableAsync { get; set; }
+
     /// <summary>
     /// Publishes a serializable message payload to the given subject name, optionally supplying a reply subject.
     /// </summary>

src/NATS.Client.Core/Internal/ISocketConnection.cs → src/NATS.Client.Core/ISocketConnection.cs

@@ -1,6 +1,6 @@
-namespace NATS.Client.Core.Internal;
+namespace NATS.Client.Core;
 
-internal interface ISocketConnection : IAsyncDisposable
+public interface ISocketConnection : IAsyncDisposable
 {
     public Task<Exception> WaitForClosed { get; }
 

src/NATS.Client.Core/NatsConnection.cs

@@ -30,11 +30,6 @@ internal enum NatsEvent
 public partial class NatsConnection : INatsConnection
 {
 #pragma warning disable SA1401
-    /// <summary>
-    /// Hook before TCP connection open.
-    /// </summary>
-    public Func<(string Host, int Port), ValueTask<(string Host, int Port)>>? OnConnectingAsync;
-
     internal readonly ConnectionStatsCounter Counter; // allow to call from external sources
     internal volatile ServerInfo? WritableServerInfo;
 
@@ -135,6 +130,11 @@ private set
 
     public NatsHeaderParser HeaderParser { get; }
 
+    // Hooks
+    public Func<(string Host, int Port), ValueTask<(string Host, int Port)>>? OnConnectingAsync { get; set; }
+
+    public Func<ISocketConnection, ValueTask<ISocketConnection>>? OnSocketAvailableAsync { get; set; }
+
     internal bool IsDisposed
     {
         get => Interlocked.CompareExchange(ref _isDisposed, 0, 0) == 1;
@@ -339,6 +339,12 @@ private async ValueTask InitialConnectAsync()
                     }
                 }
 
+                if (OnSocketAvailableAsync != null)
+                {
+                    _logger.LogInformation(NatsLogEvents.Connection, "Try to invoke OnSocketAvailable");
+                    _socket = await OnSocketAvailableAsync(_socket).ConfigureAwait(false);
+                }
+
                 _currentConnectUri = uri;
                 break;
             }
@@ -629,6 +635,12 @@ private async void ReconnectLoop()
                         }
                     }
 
+                    if (OnSocketAvailableAsync != null)
+                    {
+                        _logger.LogInformation(NatsLogEvents.Connection, "Try to invoke OnSocketAvailable");
+                        _socket = await OnSocketAvailableAsync(_socket).ConfigureAwait(false);
+                    }
+
                     _currentConnectUri = url;
                 }
                 else

tests/NATS.Client.Core.Tests/NatsConnectionTest.cs

@@ -412,6 +412,60 @@ public void NewInboxVeryLongPrefixReturnsPrefixWithNuid()
 
         Assert.Matches("A{512}\\.[A-z0-9]{22}", inbox);
     }
+
+    [Fact]
+    public async Task OnSocketAvailableAsync_ShouldBeInvokedOnInitialConnection()
+    {
+        // Arrange
+        await using var server = NatsServer.Start();
+        var clientOpts = server.ClientOpts(NatsOpts.Default);
+
+        var wasInvoked = false;
+        var nats = new NatsConnection(clientOpts);
+        nats.OnSocketAvailableAsync = async socket =>
+        {
+            wasInvoked = true;
+            await Task.Delay(10);
+            return socket;
+        };
+
+        // Act
+        await nats.ConnectAsync();
+
+        // Assert
+        Assert.True(wasInvoked, "OnSocketAvailableAsync should be invoked on initial connection.");
+    }
+
+    [Fact]
+    public async Task OnSocketAvailableAsync_ShouldBeInvokedOnReconnection()
+    {
+        // Arrange
+        await using var server = NatsServer.Start();
+        var clientOpts = server.ClientOpts(NatsOpts.Default);
+
+        var invocationCount = 0;
+        var nats = new NatsConnection(clientOpts);
+        nats.OnSocketAvailableAsync = async socket =>
+        {
+            invocationCount++;
+            await Task.Delay(10);
+            return socket;
+        };
+
+        // Simulate initial connection
+        await nats.ConnectAsync();
+
+        // Simulate disconnection
+        await server.StopAsync();
+
+        // Act
+        // Simulate reconnection
+        server.StartServerProcess();
+        await nats.ConnectAsync();
+
+        // Assert
+        Assert.Equal(2, invocationCount);
+    }
 }
 
 [JsonSerializable(typeof(SampleClass))]

tests/NATS.Client.JetStream.Tests/NatsJsContextFactoryTest.cs

@@ -106,6 +106,10 @@ public class MockConnection : INatsConnection
 
         public NatsHeaderParser HeaderParser { get; } = new NatsHeaderParser(Encoding.UTF8);
 
+        public Func<(string Host, int Port), ValueTask<(string Host, int Port)>>? OnConnectingAsync { get; set; }
+
+        public Func<ISocketConnection, ValueTask<ISocketConnection>>? OnSocketAvailableAsync { get; set; }
+
         public ValueTask<TimeSpan> PingAsync(CancellationToken cancellationToken = default) => throw new NotImplementedException();
 
         public ValueTask PublishAsync<T>(string subject, T data, NatsHeaders? headers = default, string? replyTo = default, INatsSerialize<T>? serializer = default, NatsPubOpts? opts = default, CancellationToken cancellationToken = default) => throw new NotImplementedException();