Refactor global bingins (#2343)

* Refactor global bingins

* fix typo

* Improve comment

* Make injected source static

* arrows

* Queue utils registration

* Add close reason

Author: kblok

lib/PuppeteerSharp.Tests/CDPSessionTests/CreateCDPSessionTests.cs

@@ -83,7 +83,7 @@ public async Task ShouldBeAbleToDetachSession()
             Assert.AreEqual(3, evalResponse["result"]["value"].ToObject<int>());
             await client.DetachAsync();
 
-            var exception = Assert.ThrowsAsync<PuppeteerException>(()
+            var exception = Assert.ThrowsAsync<TargetClosedException>(()
                 => client.SendAsync("Runtime.evaluate", new RuntimeEvaluateRequest
                 {
                     Expression = "3 + 1",

lib/PuppeteerSharp.Tests/WontImplementTests.cs

@@ -34,6 +34,7 @@ public WontImplementTests(): base()
         [PuppeteerTest("EventEmitter.spec.ts", "removeAllListeners", "removes every listener from all events by default")]
         [PuppeteerTest("EventEmitter.spec.ts", "removeAllListeners", "returns the emitter for chaining")]
         [PuppeteerTest("EventEmitter.spec.ts", "removeAllListeners", "can filter to remove only listeners for a given event name")]
+        [PuppeteerTest("elementhandle.spec.ts", "Custom queries", "should work with function shorthands")]
         [PuppeteerTest("emulation.spec.ts", "Page.emulateMediaType", "should throw in case of bad argument")]
         [PuppeteerTest("emulation.spec.ts", "Page.emulateMediaFeatures", "should throw in case of bad argument")]
         [PuppeteerTest("emulation.spec.ts", "Page.emulateVisionDeficiency", "should throw for invalid vision deficiencies")]

lib/PuppeteerSharp/Browser.cs

@@ -136,8 +136,6 @@ public bool IsClosed
 
         internal LauncherBase Launcher { get; set; }
 
-        internal CustomQueriesManager CustomQueriesManager { get; } = new();
-
         internal ITargetManager TargetManager { get; }
 
         internal Func<Target, bool> IsPageTargetFunc { get; set; }
@@ -245,16 +243,16 @@ public void RegisterCustomQueryHandler(string name, CustomQueryHandler queryHand
                 throw new ArgumentNullException(nameof(queryHandler));
             }
 
-            CustomQueriesManager.RegisterCustomQueryHandler(name, queryHandler);
+            Connection.CustomQuerySelectorRegistry.RegisterCustomQueryHandler(name, queryHandler);
         }
 
         /// <inheritdoc/>
         public void UnregisterCustomQueryHandler(string name)
-            => CustomQueriesManager.UnregisterCustomQueryHandler(name);
+            => Connection.CustomQuerySelectorRegistry.UnregisterCustomQueryHandler(name);
 
         /// <inheritdoc/>
         public void ClearCustomQueryHandlers()
-            => CustomQueriesManager.ClearCustomQueryHandlers();
+            => Connection.CustomQuerySelectorRegistry.ClearCustomQueryHandlers();
 
         /// <inheritdoc />
         public void Dispose()
@@ -311,7 +309,7 @@ internal static async Task<Browser> CreateAsync(
         }
 
         internal IEnumerable<string> GetCustomQueryHandlerNames()
-            => CustomQueriesManager.GetCustomQueryHandlerNames();
+            => Connection.CustomQuerySelectorRegistry.GetCustomQueryHandlerNames();
 
         internal async Task<IPage> CreatePageInContextAsync(string contextId)
         {

lib/PuppeteerSharp/CDPSession.cs

@@ -64,10 +64,11 @@ public async Task<JObject> SendAsync(string method, object args = null, bool wai
         {
             if (Connection == null)
             {
-                throw new PuppeteerException(
+                throw new TargetClosedException(
                     $"Protocol error ({method}): Session closed. " +
                     $"Most likely the {TargetType} has been closed." +
-                    $"Close reason: {CloseReason}");
+                    $"Close reason: {CloseReason}",
+                    CloseReason);
             }
 
             var id = Connection.GetMessageID();

lib/PuppeteerSharp/Connection.cs

@@ -10,6 +10,7 @@
 using PuppeteerSharp.Helpers;
 using PuppeteerSharp.Helpers.Json;
 using PuppeteerSharp.Messaging;
+using PuppeteerSharp.QueryHandlers;
 using PuppeteerSharp.Transport;
 
 namespace PuppeteerSharp
@@ -93,6 +94,13 @@ internal Connection(string url, int delay, bool enqueueAsyncMessages, IConnectio
 
         internal AsyncMessageQueue MessageQueue { get; }
 
+        // The connection is a good place to keep the state of custom queries and injectors.
+        // Although I consider that the Browser class would be a better place for this,
+        // The connection is being shared between all the components involved in one browser instance
+        internal CustomQuerySelectorRegistry CustomQuerySelectorRegistry { get; } = new();
+
+        internal ScriptInjector ScriptInjector { get; } = new();
+
         /// <inheritdoc />
         public void Dispose()
         {

lib/PuppeteerSharp/QueryHandlers/CustomQueryHandler.cs renamed to lib/PuppeteerSharp/CustomQueryHandler.cs

@@ -1,4 +1,4 @@
-namespace PuppeteerSharp.QueryHandlers
+namespace PuppeteerSharp
 {
     // This class represents both InternalQueryHandler and CustomQueryHandler upstream.
 

lib/PuppeteerSharp/ElementHandle.cs

@@ -43,7 +43,7 @@ internal ElementHandle(
 
         internal Frame Frame { get; }
 
-        internal CustomQueriesManager CustomQueriesManager => Page.Browser.CustomQueriesManager;
+        internal CustomQuerySelectorRegistry CustomQuerySelectorRegistry => Client.Connection.CustomQuerySelectorRegistry;
 
         private string DebuggerDisplay =>
             string.IsNullOrEmpty(RemoteObject.ClassName) ? ToString() : $"{RemoteObject.ClassName}@{RemoteObject.Description}";
@@ -90,8 +90,7 @@ public async Task<IElementHandle> WaitForSelectorAsync(string selector, WaitForS
                 throw new ArgumentNullException(nameof(selector));
             }
 
-            var customQueriesManager = Frame.FrameManager.Page.Browser.CustomQueriesManager;
-            var (updatedSelector, queryHandler) = customQueriesManager.GetQueryHandlerAndSelector(selector);
+            var (updatedSelector, queryHandler) = CustomQuerySelectorRegistry.GetQueryHandlerAndSelector(selector);
             return await queryHandler.WaitForAsync(null, this, updatedSelector, options).ConfigureAwait(false);
         }
 
@@ -261,7 +260,7 @@ public Task<IElementHandle> QuerySelectorAsync(string selector)
                 throw new ArgumentNullException(nameof(selector));
             }
 
-            var (updatedSelector, queryHandler) = CustomQueriesManager.GetQueryHandlerAndSelector(selector);
+            var (updatedSelector, queryHandler) = CustomQuerySelectorRegistry.GetQueryHandlerAndSelector(selector);
             return queryHandler.QueryOneAsync(this, updatedSelector);
         }
 
@@ -273,7 +272,7 @@ public async Task<IElementHandle[]> QuerySelectorAllAsync(string selector)
                 throw new ArgumentNullException(nameof(selector));
             }
 
-            var (updatedSelector, queryHandler) = CustomQueriesManager.GetQueryHandlerAndSelector(selector);
+            var (updatedSelector, queryHandler) = CustomQuerySelectorRegistry.GetQueryHandlerAndSelector(selector);
             var result = new List<IElementHandle>();
             await foreach (var item in queryHandler.QueryAllAsync(this, updatedSelector))
             {

lib/PuppeteerSharp/ExecutionContext.cs

@@ -9,6 +9,7 @@
 using Newtonsoft.Json.Linq;
 using PuppeteerSharp.Helpers;
 using PuppeteerSharp.Messaging;
+using PuppeteerSharp.QueryHandlers;
 
 namespace PuppeteerSharp
 {
@@ -18,9 +19,9 @@ public class ExecutionContext : IExecutionContext
         internal const string EvaluationScriptUrl = "__puppeteer_evaluation_script__";
 
         private static readonly Regex _sourceUrlRegex = new(@"^[\040\t]*\/\/[@#] sourceURL=\s*\S*?\s*$", RegexOptions.Multiline);
-        private static string _injectedSource;
 
         private readonly string _evaluationScriptSuffix = $"//# sourceURL={EvaluationScriptUrl}";
+        private readonly TaskQueue _puppeteerUtilQueue = new();
         private IJSHandle _puppeteerUtil;
 
         internal ExecutionContext(
@@ -98,12 +99,27 @@ public async Task<IJSHandle> QueryObjectsAsync(IJSHandle prototypeHandle)
 
         internal async Task<IJSHandle> GetPuppeteerUtilAsync()
         {
-            if (_puppeteerUtil == null)
+            await _puppeteerUtilQueue.Enqueue(async () =>
             {
-                var injectedSource = GetInjectedSource();
-                _puppeteerUtil = await EvaluateExpressionHandleAsync(injectedSource).ConfigureAwait(false);
-            }
-
+                if (_puppeteerUtil == null)
+                {
+                    await Client.Connection.ScriptInjector.InjectAsync(
+                        async (string script) =>
+                        {
+                            if (_puppeteerUtil != null)
+                            {
+                                await _puppeteerUtil.DisposeAsync().ConfigureAwait(false);
+                            }
+
+                            await InstallGlobalBindingAsync(new Binding(
+                                "__ariaQuerySelector",
+                                (Func<IElementHandle, string, Task<IElementHandle>>)Client.Connection.CustomQuerySelectorRegistry.InternalQueryHandlers["aria"].QueryOneAsync))
+                                .ConfigureAwait(false);
+                            _puppeteerUtil = await EvaluateExpressionHandleAsync(script).ConfigureAwait(false);
+                        },
+                        _puppeteerUtil == null).ConfigureAwait(false);
+                }
+            }).ConfigureAwait(false);
             return _puppeteerUtil;
         }
 
@@ -138,22 +154,6 @@ internal async Task<IElementHandle> AdoptElementHandleAsync(IElementHandle eleme
             return CreateJSHandle(obj.Object) as ElementHandle;
         }
 
-        private static string GetInjectedSource()
-        {
-            if (string.IsNullOrEmpty(_injectedSource))
-            {
-                var assembly = Assembly.GetExecutingAssembly();
-                var resourceName = "PuppeteerSharp.Injected.injected.js";
-
-                using var stream = assembly.GetManifestResourceStream(resourceName);
-                using var reader = new StreamReader(stream);
-                var fileContent = reader.ReadToEnd();
-                _injectedSource = fileContent;
-            }
-
-            return _injectedSource;
-        }
-
         private static string GetExceptionMessage(EvaluateExceptionResponseDetails exceptionDetails)
         {
             if (exceptionDetails.Exception != null)
@@ -175,6 +175,24 @@ private static string GetExceptionMessage(EvaluateExceptionResponseDetails excep
             return message;
         }
 
+        private async Task InstallGlobalBindingAsync(Binding binding)
+        {
+            try
+            {
+                if (World != null)
+                {
+                    World.Bindings.TryAdd(binding.Name, binding);
+                    await World.AddBindingToContextAsync(this, binding.Name).ConfigureAwait(false);
+                }
+            }
+            catch
+            {
+                // If the binding cannot be added, then either the browser doesn't support
+                // bindings (e.g. Firefox) or the context is broken. Either breakage is
+                // okay, so we ignore the error.
+            }
+        }
+
         private async Task<T> RemoteObjectTaskToObject<T>(Task<RemoteObject> remote)
         {
             var response = await remote.ConfigureAwait(false);

lib/PuppeteerSharp/Frame.cs

@@ -102,8 +102,7 @@ public async Task<IElementHandle> WaitForSelectorAsync(string selector, WaitForS
                 throw new ArgumentNullException(nameof(selector));
             }
 
-            var customQueriesManager = ((Browser)FrameManager.Page.Browser).CustomQueriesManager;
-            var (updatedSelector, queryHandler) = customQueriesManager.GetQueryHandlerAndSelector(selector);
+            var (updatedSelector, queryHandler) = Client.Connection.CustomQuerySelectorRegistry.GetQueryHandlerAndSelector(selector);
             return await queryHandler.WaitForAsync(this, null, updatedSelector, options).ConfigureAwait(false);
         }