Nats web socket opts improvements (nats-io#623)

* NatsWebSocketOpts improvements (nats-io#610) * Separate model NatsWebSocketOpts * Сalled NatsTlsOpts.AuthenticateAsClientAsync parameters before passing it to the ConfigureWebSocketOpts * RequestHeaders overwrite the header specified in ConfigureWebSocketOpts * NatsWebSocketOpts improvements (nats-io#610) * Added test * Build fixes * Build fixes * Build fixes * dotnet format * nats-io#623 reccomendations Signed-off-by: Caleb Lloyd <[email protected]> * NatsWebSocketOpts improvements (nats-io#610) * Fence added, an occurs exceptions when add headers to blazor * Revert "NatsWebSocketOpts improvements (nats-io#610)" This reverts commit e6b9d1c. * WebSocketSecure tests Signed-off-by: Caleb Lloyd <[email protected]> * fix format Signed-off-by: Caleb Lloyd <[email protected]> * simplify WebSocketOptionsTest Signed-off-by: Caleb Lloyd <[email protected]> --------- Signed-off-by: Caleb Lloyd <[email protected]> Co-authored-by: Ziya Suzen <[email protected]> Co-authored-by: Caleb Lloyd <[email protected]>
divyeshio · Sep 13, 2024 · dd0aca3 · dd0aca3
1 parent 7124125
commit dd0aca3
Show file tree

Hide file tree

Showing 9 changed files with 188 additions and 358 deletions.
diff --git a/src/NATS.Client.Core/Internal/WebSocketConnection.cs b/src/NATS.Client.Core/Internal/WebSocketConnection.cs
@@ -1,3 +1,4 @@
+using System.Net.Security;
 using System.Net.Sockets;
 using System.Net.WebSockets;
 using System.Runtime.CompilerServices;
@@ -39,13 +40,13 @@ public Task ConnectAsync(Uri uri, CancellationToken cancellationToken)
     /// <summary>
     /// Connect with Timeout. When failed, Dispose this connection.
     /// </summary>
-    public async ValueTask ConnectAsync(Uri uri, NatsOpts opts)
+    public async ValueTask ConnectAsync(NatsUri uri, NatsOpts opts)
     {
         using var cts = new CancellationTokenSource(opts.ConnectTimeout);
         try
         {
-            await InvokeCallbackForClientWebSocketOptionsAsync(opts, uri, _socket.Options, cts.Token).ConfigureAwait(false);
-            await _socket.ConnectAsync(uri, cts.Token).ConfigureAwait(false);
+            await opts.WebSocketOpts.ApplyClientWebSocketOptionsAsync(_socket.Options, uri, opts.TlsOpts, cts.Token).ConfigureAwait(false);
+            await _socket.ConnectAsync(uri.Uri, cts.Token).ConfigureAwait(false);
         }
         catch (Exception ex)
         {
@@ -131,13 +132,4 @@ public void SignalDisconnected(Exception exception)
     {
         _waitForClosedSource.TrySetResult(exception);
     }
-
-    [MethodImpl(MethodImplOptions.AggressiveInlining)]
-    private async Task InvokeCallbackForClientWebSocketOptionsAsync(NatsOpts opts, Uri uri, ClientWebSocketOptions options, CancellationToken token)
-    {
-        if (opts.ConfigureWebSocketOpts != null)
-        {
-            await opts.ConfigureWebSocketOpts(uri, options, token).ConfigureAwait(false);
-        }
-    }
 }
diff --git a/src/NATS.Client.Core/NatsConnection.cs b/src/NATS.Client.Core/NatsConnection.cs
@@ -317,7 +317,7 @@ private async ValueTask InitialConnectAsync()
                 if (uri.IsWebSocket)
                 {
                     var conn = new WebSocketConnection();
-                    await conn.ConnectAsync(uri.Uri, Opts).ConfigureAwait(false);
+                    await conn.ConnectAsync(uri, Opts).ConfigureAwait(false);
                     _socket = conn;
                 }
                 else
@@ -605,7 +605,7 @@ private async void ReconnectLoop()
                     {
                         _logger.LogDebug(NatsLogEvents.Connection, "Trying to reconnect using WebSocket {Url} [{ReconnectCount}]", url, reconnectCount);
                         var conn = new WebSocketConnection();
-                        await conn.ConnectAsync(url.Uri, Opts).ConfigureAwait(false);
+                        await conn.ConnectAsync(url, Opts).ConfigureAwait(false);
                         _socket = conn;
                     }
                     else

diff --git a/src/NATS.Client.Core/NatsOpts.cs b/src/NATS.Client.Core/NatsOpts.cs
@@ -28,6 +28,8 @@ public sealed record NatsOpts
 
     public NatsTlsOpts TlsOpts { get; init; } = NatsTlsOpts.Default;
 
+    public NatsWebSocketOpts WebSocketOpts { get; init; } = NatsWebSocketOpts.Default;
+
     public INatsSerializerRegistry SerializerRegistry { get; init; } = NatsDefaultSerializerRegistry.Default;
 
     public ILoggerFactory LoggerFactory { get; init; } = NullLoggerFactory.Instance;
@@ -115,28 +117,6 @@ public sealed record NatsOpts
     /// </remarks>
     public BoundedChannelFullMode SubPendingChannelFullMode { get; init; } = BoundedChannelFullMode.DropNewest;
 
-    /// <summary>
-    /// An optional async callback handler for manipulation of ClientWebSocketOptions used for WebSocket connections.
-    /// </summary>
-    /// <remarks>
-    /// This can be used to set authorization header and other HTTP header values.
-    /// Note: Setting HTTP header values is not supported by Blazor WebAssembly as the underlying browser implementation does not support adding headers to a WebSocket.
-    /// The callback's execution time contributes to the connection establishment subject to the <see cref="ConnectTimeout"/>.
-    /// Implementors should use the passed CancellationToken for async operations called by this handler.
-    /// </remarks>
-    /// <example>
-    /// await using var nats = new NatsConnection(new NatsOpts
-    /// {
-    ///     Url = "ws://localhost:8080",
-    ///     ConfigureWebSocketOpts = (serverUri, clientWsOpts, ct) =>
-    ///     {
-    ///         clientWsOpts.SetRequestHeader("authorization", $"Bearer MY_TOKEN");
-    ///         return ValueTask.CompletedTask;
-    ///     },
-    /// });
-    /// </example>
-    public Func<Uri, ClientWebSocketOptions, CancellationToken, ValueTask>? ConfigureWebSocketOpts { get; init; } = null;
-
     internal NatsUri[] GetSeedUris()
     {
         var urls = Url.Split(',');

diff --git a/src/NATS.Client.Core/NatsWebSocketOpts.cs b/src/NATS.Client.Core/NatsWebSocketOpts.cs
@@ -0,0 +1,79 @@
+using System.Net.WebSockets;
+using System.Security.Cryptography.X509Certificates;
+using Microsoft.Extensions.Primitives;
+using NATS.Client.Core.Internal;
+
+namespace NATS.Client.Core;
+
+/// <summary>
+/// Options for ClientWebSocketOptions
+/// </summary>
+public sealed record NatsWebSocketOpts
+{
+    public static readonly NatsWebSocketOpts Default = new();
+
+    /// <summary>
+    /// An optional dictionary of HTTP request headers to be sent with the WebSocket request.
+    /// </summary>
+    /// <remarks>
+    /// Not supported when running in the Browser, such as when using Blazor WebAssembly,
+    /// as the underlying Browser implementation does not support adding headers to a WebSocket.
+    /// </remarks>
+    public IDictionary<string, StringValues>? RequestHeaders { get; init; }
+
+    /// <summary>
+    /// An optional async callback handler for manipulation of ClientWebSocketOptions used for WebSocket connections.
+    /// Implementors should use the passed CancellationToken for async operations called by this handler.
+    /// </summary>
+    public Func<Uri, ClientWebSocketOptions, CancellationToken, ValueTask>? ConfigureClientWebSocketOptions { get; init; } = null;
+
+    internal async ValueTask ApplyClientWebSocketOptionsAsync(
+        ClientWebSocketOptions clientWebSocketOptions,
+        NatsUri uri,
+        NatsTlsOpts tlsOpts,
+        CancellationToken cancellationToken)
+    {
+        if (RequestHeaders != null)
+        {
+            foreach (var entry in RequestHeaders)
+            {
+                // SetRequestHeader overwrites if called multiple times;
+                // RFC7230 Section 3.2.2 allows for combining them with a comma
+                // https://www.rfc-editor.org/rfc/rfc7230#section-3.2.2
+                clientWebSocketOptions.SetRequestHeader(entry.Key, string.Join(",", entry.Value));
+            }
+        }
+
+        if (tlsOpts.HasTlsCerts)
+        {
+            var authenticateAsClientOptions = await tlsOpts.AuthenticateAsClientOptionsAsync(uri).ConfigureAwait(false);
+            var collection = new X509CertificateCollection();
+
+            // must match LoadClientCertFromX509 method in SslClientAuthenticationOptions.cs
+#if NET8_0_OR_GREATER
+            if (authenticateAsClientOptions.ClientCertificateContext != null)
+            {
+                collection.Add(authenticateAsClientOptions.ClientCertificateContext.TargetCertificate);
+            }
+#else
+            if (authenticateAsClientOptions.ClientCertificates != null)
+            {
+                collection.AddRange(authenticateAsClientOptions.ClientCertificates);
+            }
+#endif
+            if (collection.Count > 0)
+            {
+                clientWebSocketOptions.ClientCertificates = collection;
+            }
+
+#if !NETSTANDARD2_0
+            clientWebSocketOptions.RemoteCertificateValidationCallback = authenticateAsClientOptions.RemoteCertificateValidationCallback;
+#endif
+        }
+
+        if (ConfigureClientWebSocketOptions != null)
+        {
+            await ConfigureClientWebSocketOptions(uri.Uri, clientWebSocketOptions, cancellationToken).ConfigureAwait(false);
+        }
+    }
+}
diff --git a/tests/NATS.Client.Core.Tests/NatsConnectionTest.Transports.cs b/tests/NATS.Client.Core.Tests/NatsConnectionTest.Transports.cs
@@ -23,3 +23,11 @@ public NatsConnectionTestWs(ITestOutputHelper output)
     {
     }
 }
+
+public class NatsConnectionTestWss : NatsConnectionTest
+{
+    public NatsConnectionTestWss(ITestOutputHelper output)
+        : base(output, TransportType.WebSocketSecure)
+    {
+    }
+}
diff --git a/tests/NATS.Client.Core.Tests/TlsClientTest.cs b/tests/NATS.Client.Core.Tests/TlsClientTest.cs
@@ -11,13 +11,15 @@ public class TlsClientTest
 
     public TlsClientTest(ITestOutputHelper output) => _output = output;
 
-    [Fact]
-    public async Task Client_connect_using_certificate()
+    [Theory]
+    [InlineData(TransportType.Tls)]
+    [InlineData(TransportType.WebSocketSecure)]
+    public async Task Client_connect_using_certificate(TransportType transportType)
     {
         await using var server = NatsServer.Start(
             new NullOutputHelper(),
             new NatsServerOptsBuilder()
-                .UseTransport(TransportType.Tls, tlsVerify: true)
+                .UseTransport(transportType, tlsVerify: true)
                 .Build());
 
         var clientOpts = server.ClientOpts(NatsOpts.Default with { Name = "tls-test-client" });
@@ -56,13 +58,15 @@ public async Task Client_connect_using_certificate_and_revocation_check()
         Assert.Contains("remote certificate was rejected", exception.InnerException!.InnerException!.Message);
     }
 
-    [Fact]
-    public async Task Client_cannot_connect_without_certificate()
+    [Theory]
+    [InlineData(TransportType.Tls)]
+    [InlineData(TransportType.WebSocketSecure)]
+    public async Task Client_cannot_connect_without_certificate(TransportType transportType)
     {
         await using var server = NatsServer.Start(
             new NullOutputHelper(),
             new NatsServerOptsBuilder()
-                .UseTransport(TransportType.Tls, tlsVerify: true)
+                .UseTransport(transportType, tlsVerify: true)
                 .Build());
 
         var clientOpts = server.ClientOpts(NatsOpts.Default);