[dotnet] [bidi] Provide extension points for custom BiDi modules

Author: RenderMichael

dotnet/src/webdriver/BiDi/BiDi.cs

@@ -27,7 +27,7 @@ namespace OpenQA.Selenium.BiDi;
 public class BiDi : IAsyncDisposable
 {
     private readonly ITransport _transport;
-    private readonly Broker _broker;
+    protected Broker Broker { get; }
 
     private readonly Lazy<Modules.Session.SessionModule> _sessionModule;
     private readonly Lazy<Modules.BrowsingContext.BrowsingContextModule> _browsingContextModule;
@@ -38,21 +38,21 @@ public class BiDi : IAsyncDisposable
     private readonly Lazy<Modules.Log.LogModule> _logModule;
     private readonly Lazy<Modules.Storage.StorageModule> _storageModule;
 
-    internal BiDi(string url)
+    protected internal BiDi(string url)
     {
         var uri = new Uri(url);
 
         _transport = new WebSocketTransport(new Uri(url));
-        _broker = new Broker(this, _transport);
+        Broker = new Broker(this, _transport);
 
-        _sessionModule = new Lazy<Modules.Session.SessionModule>(() => new Modules.Session.SessionModule(_broker));
-        _browsingContextModule = new Lazy<Modules.BrowsingContext.BrowsingContextModule>(() => new Modules.BrowsingContext.BrowsingContextModule(_broker));
-        _browserModule = new Lazy<Modules.Browser.BrowserModule>(() => new Modules.Browser.BrowserModule(_broker));
-        _networkModule = new Lazy<Modules.Network.NetworkModule>(() => new Modules.Network.NetworkModule(_broker));
-        _inputModule = new Lazy<Modules.Input.InputModule>(() => new Modules.Input.InputModule(_broker));
-        _scriptModule = new Lazy<Modules.Script.ScriptModule>(() => new Modules.Script.ScriptModule(_broker));
-        _logModule = new Lazy<Modules.Log.LogModule>(() => new Modules.Log.LogModule(_broker));
-        _storageModule = new Lazy<Modules.Storage.StorageModule>(() => new Modules.Storage.StorageModule(_broker));
+        _sessionModule = new Lazy<Modules.Session.SessionModule>(() => new Modules.Session.SessionModule(Broker));
+        _browsingContextModule = new Lazy<Modules.BrowsingContext.BrowsingContextModule>(() => new Modules.BrowsingContext.BrowsingContextModule(Broker));
+        _browserModule = new Lazy<Modules.Browser.BrowserModule>(() => new Modules.Browser.BrowserModule(Broker));
+        _networkModule = new Lazy<Modules.Network.NetworkModule>(() => new Modules.Network.NetworkModule(Broker));
+        _inputModule = new Lazy<Modules.Input.InputModule>(() => new Modules.Input.InputModule(Broker));
+        _scriptModule = new Lazy<Modules.Script.ScriptModule>(() => new Modules.Script.ScriptModule(Broker));
+        _logModule = new Lazy<Modules.Log.LogModule>(() => new Modules.Log.LogModule(Broker));
+        _storageModule = new Lazy<Modules.Storage.StorageModule>(() => new Modules.Storage.StorageModule(Broker));
     }
 
     internal Modules.Session.SessionModule SessionModule => _sessionModule.Value;
@@ -73,7 +73,7 @@ public static async Task<BiDi> ConnectAsync(string url)
     {
         var bidi = new BiDi(url);
 
-        await bidi._broker.ConnectAsync(default).ConfigureAwait(false);
+        await bidi.Broker.ConnectAsync().ConfigureAwait(false);
 
         return bidi;
     }
@@ -83,10 +83,15 @@ public Task EndAsync(Modules.Session.EndOptions? options = null)
         return SessionModule.EndAsync(options);
     }
 
-    public async ValueTask DisposeAsync()
+    protected virtual async ValueTask DisposeAsyncCore()
     {
-        await _broker.DisposeAsync().ConfigureAwait(false);
-
+        await Broker.DisposeAsync().ConfigureAwait(false);
         _transport?.Dispose();
     }
+
+    public async ValueTask DisposeAsync()
+    {
+        await DisposeAsyncCore();
+        GC.SuppressFinalize(this);
+    }
 }

dotnet/src/webdriver/BiDi/Communication/Broker.cs

@@ -27,6 +27,7 @@
 using System.Linq;
 using System.Text.Json;
 using System.Text.Json.Serialization;
+using System.Text.Json.Serialization.Metadata;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -52,7 +53,7 @@ public class Broker : IAsyncDisposable
     private Task? _eventEmitterTask;
     private CancellationTokenSource? _receiveMessagesCancellationTokenSource;
 
-    private readonly BiDiJsonSerializerContext _jsonSerializerContext;
+    private readonly JsonSerializerOptions _jsonSerializerContext;
 
     internal Broker(BiDi bidi, ITransport transport)
     {
@@ -99,13 +100,22 @@ internal Broker(BiDi bidi, ITransport transport)
                 new Json.Converters.Enumerable.GetUserContextsResultConverter(),
                 new Json.Converters.Enumerable.GetClientWindowsResultConverter(),
                 new Json.Converters.Enumerable.GetRealmsResultConverter(),
+            },
+            TypeInfoResolverChain =
+            {
+                BiDiJsonSerializerContext.Default
             }
         };
 
-        _jsonSerializerContext = new BiDiJsonSerializerContext(jsonSerializerOptions);
+        _jsonSerializerContext = jsonSerializerOptions;
+    }
+
+    public void ProvideCustomSerializationContext(JsonSerializerContext extensionContext)
+    {
+        _jsonSerializerContext.TypeInfoResolverChain.Add(extensionContext);
     }
 
-    public async Task ConnectAsync(CancellationToken cancellationToken)
+    public async Task ConnectAsync(CancellationToken cancellationToken = default)
     {
         await _transport.ConnectAsync(cancellationToken).ConfigureAwait(false);
 
@@ -120,7 +130,8 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
         {
             var data = await _transport.ReceiveAsync(cancellationToken).ConfigureAwait(false);
 
-            var message = JsonSerializer.Deserialize(new ReadOnlySpan<byte>(data), _jsonSerializerContext.Message);
+            var messageTypeInfo = (JsonTypeInfo<Message>)_jsonSerializerContext.GetTypeInfo(typeof(Message));
+            var message = JsonSerializer.Deserialize(new ReadOnlySpan<byte>(data), messageTypeInfo);
 
             switch (message)
             {
@@ -131,9 +142,9 @@ private async Task ReceiveMessagesAsync(CancellationToken cancellationToken)
                 case MessageEvent messageEvent:
                     _pendingEvents.Add(messageEvent);
                     break;
-                case MessageError mesageError:
-                    _pendingCommands[mesageError.Id].SetException(new BiDiException($"{mesageError.Error}: {mesageError.Message}"));
-                    _pendingCommands.TryRemove(mesageError.Id, out _);
+                case MessageError messageError:
+                    _pendingCommands[messageError.Id].SetException(new BiDiException($"{messageError.Error}: {messageError.Message}"));
+                    _pendingCommands.TryRemove(messageError.Id, out _);
                     break;
             }
         }