LibWeb+WebContent+WebDriver: Make the screenshot endpoints asynchronous

These were the last WebDriver endpoints spinning the event loop.

Userland/Libraries/LibWeb/WebDriver/Screenshot.cpp

@@ -1,30 +1,66 @@
 /*
- * Copyright (c) 2022, Tim Flynn <trflynn89@serenityos.org>
+ * Copyright (c) 2022-2024, Tim Flynn <trflynn89@ladybird.org>
  *
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
-#include <AK/Optional.h>
 #include <LibGfx/Bitmap.h>
-#include <LibGfx/Rect.h>
 #include <LibWeb/DOM/Document.h>
 #include <LibWeb/DOM/ElementFactory.h>
-#include <LibWeb/HTML/AnimationFrameCallbackDriver.h>
 #include <LibWeb/HTML/BrowsingContext.h>
 #include <LibWeb/HTML/HTMLCanvasElement.h>
 #include <LibWeb/HTML/TagNames.h>
 #include <LibWeb/HTML/TraversableNavigable.h>
-#include <LibWeb/HTML/Window.h>
 #include <LibWeb/Namespace.h>
 #include <LibWeb/Page/Page.h>
-#include <LibWeb/Platform/EventLoopPlugin.h>
-#include <LibWeb/WebDriver/Error.h>
 #include <LibWeb/WebDriver/Screenshot.h>
 
 namespace Web::WebDriver {
 
+// https://w3c.github.io/webdriver/#dfn-draw-a-bounding-box-from-the-framebuffer
+ErrorOr<JS::NonnullGCPtr<HTML::HTMLCanvasElement>, WebDriver::Error> draw_bounding_box_from_the_framebuffer(HTML::BrowsingContext& browsing_context, DOM::Element& element, Gfx::IntRect rect)
+{
+    // 1. If either the initial viewport's width or height is 0 CSS pixels, return error with error code unable to capture screen.
+    auto viewport_rect = browsing_context.top_level_traversable()->viewport_rect();
+    if (viewport_rect.is_empty())
+        return Error::from_code(ErrorCode::UnableToCaptureScreen, "Viewport is empty"sv);
+
+    auto viewport_device_rect = browsing_context.page().enclosing_device_rect(viewport_rect).to_type<int>();
+
+    // 2. Let paint width be the initial viewport's width – min(rectangle x coordinate, rectangle x coordinate + rectangle width dimension).
+    auto paint_width = viewport_device_rect.width() - min(rect.x(), rect.x() + rect.width());
+
+    // 3. Let paint height be the initial viewport's height – min(rectangle y coordinate, rectangle y coordinate + rectangle height dimension).
+    auto paint_height = viewport_device_rect.height() - min(rect.y(), rect.y() + rect.height());
+
+    // 4. Let canvas be a new canvas element, and set its width and height to paint width and paint height, respectively.
+    auto canvas_element = DOM::create_element(element.document(), HTML::TagNames::canvas, Namespace::HTML).release_value_but_fixme_should_propagate_errors();
+    auto& canvas = verify_cast<HTML::HTMLCanvasElement>(*canvas_element);
+
+    // FIXME: Handle DevicePixelRatio in HiDPI mode.
+    MUST(canvas.set_width(paint_width));
+    MUST(canvas.set_height(paint_height));
+
+    // FIXME: 5. Let context, a canvas context mode, be the result of invoking the 2D context creation algorithm given canvas as the target.
+    if (!canvas.create_bitmap(paint_width, paint_height))
+        return Error::from_code(ErrorCode::UnableToCaptureScreen, "Unable to create a screenshot bitmap"sv);
+
+    // 6. Complete implementation specific steps equivalent to drawing the region of the framebuffer specified by the following coordinates onto context:
+    //    - X coordinate: rectangle x coordinate
+    //    - Y coordinate: rectangle y coordinate
+    //    - Width: paint width
+    //    - Height: paint height
+    Gfx::IntRect paint_rect { rect.x(), rect.y(), paint_width, paint_height };
+
+    auto backing_store = Web::Painting::BitmapBackingStore(*canvas.bitmap());
+    browsing_context.page().client().paint(paint_rect.to_type<Web::DevicePixels>(), backing_store);
+
+    // 7. Return success with canvas.
+    return canvas;
+}
+
 // https://w3c.github.io/webdriver/#dfn-encoding-a-canvas-as-base64
-static Response encode_canvas_element(HTML::HTMLCanvasElement& canvas)
+Response encode_canvas_element(HTML::HTMLCanvasElement& canvas)
 {
     // FIXME: 1. If the canvas element’s bitmap’s origin-clean flag is set to false, return error with error code unable to capture screen.
 
@@ -44,66 +80,7 @@ static Response encode_canvas_element(HTML::HTMLCanvasElement& canvas)
     auto encoded_string = MUST(data_url.substring_from_byte_offset(*index + 1));
 
     // 7. Return success with data encoded string.
-    return JsonValue { move(encoded_string) };
-}
-
-// Common animation callback steps between:
-// https://w3c.github.io/webdriver/#take-screenshot
-// https://w3c.github.io/webdriver/#take-element-screenshot
-Response capture_element_screenshot(Painter const& painter, Page& page, DOM::Element& element, Gfx::IntRect& rect)
-{
-    Optional<Response> encoded_string_or_error;
-
-    // https://w3c.github.io/webdriver/#dfn-draw-a-bounding-box-from-the-framebuffer
-    auto draw_bounding_box_from_the_framebuffer = [&]() -> ErrorOr<JS::NonnullGCPtr<HTML::HTMLCanvasElement>, WebDriver::Error> {
-        // 1. If either the initial viewport's width or height is 0 CSS pixels, return error with error code unable to capture screen.
-        auto viewport_rect = page.top_level_traversable()->viewport_rect();
-        if (viewport_rect.is_empty())
-            return Error::from_code(ErrorCode::UnableToCaptureScreen, "Viewport is empty"sv);
-
-        auto viewport_device_rect = page.enclosing_device_rect(viewport_rect).to_type<int>();
-
-        // 2. Let paint width be the initial viewport's width – min(rectangle x coordinate, rectangle x coordinate + rectangle width dimension).
-        auto paint_width = viewport_device_rect.width() - min(rect.x(), rect.x() + rect.width());
-
-        // 3. Let paint height be the initial viewport's height – min(rectangle y coordinate, rectangle y coordinate + rectangle height dimension).
-        auto paint_height = viewport_device_rect.height() - min(rect.y(), rect.y() + rect.height());
-
-        // 4. Let canvas be a new canvas element, and set its width and height to paint width and paint height, respectively.
-        auto canvas_element = DOM::create_element(element.document(), HTML::TagNames::canvas, Namespace::HTML).release_value_but_fixme_should_propagate_errors();
-        auto& canvas = verify_cast<HTML::HTMLCanvasElement>(*canvas_element);
-
-        // FIXME: Handle DevicePixelRatio in HiDPI mode.
-        MUST(canvas.set_width(paint_width));
-        MUST(canvas.set_height(paint_height));
-
-        // FIXME: 5. Let context, a canvas context mode, be the result of invoking the 2D context creation algorithm given canvas as the target.
-        if (!canvas.create_bitmap(paint_width, paint_height))
-            return Error::from_code(ErrorCode::UnableToCaptureScreen, "Unable to create a screenshot bitmap"sv);
-
-        // 6. Complete implementation specific steps equivalent to drawing the region of the framebuffer specified by the following coordinates onto context:
-        //    - X coordinate: rectangle x coordinate
-        //    - Y coordinate: rectangle y coordinate
-        //    - Width: paint width
-        //    - Height: paint height
-        Gfx::IntRect paint_rect { rect.x(), rect.y(), paint_width, paint_height };
-        painter(paint_rect, *canvas.bitmap());
-
-        // 7. Return success with canvas.
-        return canvas;
-    };
-
-    (void)element.document().window()->animation_frame_callback_driver().add(JS::create_heap_function(element.heap(), [&](double) {
-        auto canvas_or_error = draw_bounding_box_from_the_framebuffer();
-        if (canvas_or_error.is_error()) {
-            encoded_string_or_error = canvas_or_error.release_error();
-            return;
-        }
-        encoded_string_or_error = encode_canvas_element(canvas_or_error.release_value());
-    }));
-
-    Platform::EventLoopPlugin::the().spin_until(JS::create_heap_function(element.document().heap(), [&]() { return encoded_string_or_error.has_value(); }));
-    return encoded_string_or_error.release_value();
+    return JsonValue { encoded_string.to_byte_string() };
 }
 
 }

Userland/Libraries/LibWeb/WebDriver/Screenshot.h

@@ -1,19 +1,19 @@
 /*
- * Copyright (c) 2022, Tim Flynn <trflynn89@serenityos.org>
+ * Copyright (c) 2022-2024, Tim Flynn <trflynn89@ladybird.org>
  *
  * SPDX-License-Identifier: BSD-2-Clause
  */
 
 #pragma once
 
-#include <AK/Function.h>
-#include <LibGfx/Forward.h>
+#include <LibGfx/Rect.h>
+#include <LibJS/Heap/GCPtr.h>
 #include <LibWeb/Forward.h>
 #include <LibWeb/WebDriver/Response.h>
 
 namespace Web::WebDriver {
 
-using Painter = Function<void(Gfx::IntRect const&, Gfx::Bitmap&)>;
-Response capture_element_screenshot(Painter const& painter, Page& page, DOM::Element& element, Gfx::IntRect& rect);
+ErrorOr<JS::NonnullGCPtr<HTML::HTMLCanvasElement>, WebDriver::Error> draw_bounding_box_from_the_framebuffer(HTML::BrowsingContext&, DOM::Element&, Gfx::IntRect);
+Response encode_canvas_element(HTML::HTMLCanvasElement&);
 
 }

Userland/Services/WebContent/WebDriverConnection.cpp

@@ -29,6 +29,7 @@
 #include <LibWeb/DOM/Position.h>
 #include <LibWeb/DOM/ShadowRoot.h>
 #include <LibWeb/Geometry/DOMRect.h>
+#include <LibWeb/HTML/AnimationFrameCallbackDriver.h>
 #include <LibWeb/HTML/AttributeNames.h>
 #include <LibWeb/HTML/BrowsingContext.h>
 #include <LibWeb/HTML/Focus.h>
@@ -2279,65 +2280,80 @@ Messages::WebDriverClient::SendAlertTextResponse WebDriverConnection::send_alert
 // 17.1 Take Screenshot, https://w3c.github.io/webdriver/#take-screenshot
 Messages::WebDriverClient::TakeScreenshotResponse WebDriverConnection::take_screenshot()
 {
-    // 1. If the current top-level browsing context is no longer open, return error with error code no such window.
+    // 1. If session's current top-level browsing context is no longer open, return error with error code no such window.
     TRY(ensure_current_top_level_browsing_context_is_open());
 
-    // 2. When the user agent is next to run the animation frame callbacks:
-    //     a. Let root rect be the current top-level browsing context’s document element’s rectangle.
-    //     b. Let screenshot result be the result of trying to call draw a bounding box from the framebuffer, given root rect as an argument.
-    //     c. Let canvas be a canvas element of screenshot result’s data.
-    //     d. Let encoding result be the result of trying encoding a canvas as Base64 canvas.
-    //     e. Let encoded string be encoding result’s data.
     auto* document = current_top_level_browsing_context()->active_document();
-    auto root_rect = calculate_absolute_rect_of_element(*document->document_element());
+    auto window = document->window();
+
+    // 2. When the user agent is next to run the animation frame callbacks:
+    (void)window->animation_frame_callback_driver().add(JS::create_heap_function(document->heap(), [this, document](double) mutable {
+        // a. Let root rect be session's current top-level browsing context's document element's rectangle.
+        auto root_rect = calculate_absolute_rect_of_element(*document->document_element());
+
+        // b. Let screenshot result be the result of trying to call draw a bounding box from the framebuffer, given root rect as an argument.
+        auto screenshot_result = Web::WebDriver::draw_bounding_box_from_the_framebuffer(*current_top_level_browsing_context(), *document->document_element(), root_rect);
+        if (screenshot_result.is_error()) {
+            async_screenshot_taken(screenshot_result.release_error());
+            return;
+        }
+
+        // c. Let canvas be a canvas element of screenshot result's data.
+        auto canvas = screenshot_result.release_value();
 
-    auto encoded_string = TRY(Web::WebDriver::capture_element_screenshot(
-        [&](auto const& rect, auto& bitmap) {
-            auto backing_store = Web::Painting::BitmapBackingStore(bitmap);
-            current_top_level_browsing_context()->page().client().paint(rect.template to_type<Web::DevicePixels>(), backing_store);
-        },
-        current_top_level_browsing_context()->page(),
-        *document->document_element(),
-        root_rect));
+        // d. Let encoding result be the result of trying encoding a canvas as Base64 canvas.
+        // e. Let encoded string be encoding result's data.
+        auto encoded_string = Web::WebDriver::encode_canvas_element(canvas);
 
-    // 3. Return success with data encoded string.
-    return encoded_string;
+        // 3. Return success with data encoded string.
+        async_screenshot_taken(move(encoded_string));
+    }));
+
+    return JsonValue {};
 }
 
 // 17.2 Take Element Screenshot, https://w3c.github.io/webdriver/#dfn-take-element-screenshot
 Messages::WebDriverClient::TakeElementScreenshotResponse WebDriverConnection::take_element_screenshot(String const& element_id)
 {
-    // 1. If the current top-level browsing context is no longer open, return error with error code no such window.
-    TRY(ensure_current_top_level_browsing_context_is_open());
+    // 1. If session's current browsing context is no longer open, return error with error code no such window.
+    TRY(ensure_current_browsing_context_is_open());
 
-    // 2. Handle any user prompts and return its value if it is an error.
+    auto* document = current_browsing_context().active_document();
+    auto window = document->window();
+
+    // 2. Try to handle any user prompts with session.
     TRY(handle_any_user_prompts());
 
-    // 3. Let element be the result of trying to get a known connected element with url variable element id.
+    // 3. Let element be the result of trying to get a known element with session and URL variables["element id"].
     auto element = TRY(Web::WebDriver::get_known_element(current_browsing_context(), element_id));
 
     // 4. Scroll into view the element.
-    scroll_element_into_view(*element);
+    scroll_element_into_view(element);
 
     // 5. When the user agent is next to run the animation frame callbacks:
-    //     a. Let element rect be element’s rectangle.
-    //     b. Let screenshot result be the result of trying to call draw a bounding box from the framebuffer, given element rect as an argument.
-    //     c. Let canvas be a canvas element of screenshot result’s data.
-    //     d. Let encoding result be the result of trying encoding a canvas as Base64 canvas.
-    //     e. Let encoded string be encoding result’s data.
-    auto element_rect = calculate_absolute_rect_of_element(*element);
-
-    auto encoded_string = TRY(Web::WebDriver::capture_element_screenshot(
-        [&](auto const& rect, auto& bitmap) {
-            auto backing_store = Web::Painting::BitmapBackingStore(bitmap);
-            current_top_level_browsing_context()->page().client().paint(rect.template to_type<Web::DevicePixels>(), backing_store);
-        },
-        current_top_level_browsing_context()->page(),
-        *element,
-        element_rect));
-
-    // 6. Return success with data encoded string.
-    return encoded_string;
+    (void)window->animation_frame_callback_driver().add(JS::create_heap_function(document->heap(), [this, element](double) {
+        // a. Let element rect be element's rectangle.
+        auto element_rect = calculate_absolute_rect_of_element(element);
+
+        // b. Let screenshot result be the result of trying to call draw a bounding box from the framebuffer, given element rect as an argument.
+        auto screenshot_result = Web::WebDriver::draw_bounding_box_from_the_framebuffer(current_browsing_context(), element, element_rect);
+        if (screenshot_result.is_error()) {
+            async_screenshot_taken(screenshot_result.release_error());
+            return;
+        }
+
+        // c. Let canvas be a canvas element of screenshot result's data.
+        auto canvas = screenshot_result.release_value();
+
+        // d. Let encoding result be the result of trying encoding a canvas as Base64 canvas.
+        // e. Let encoded string be encoding result's data.
+        auto encoded_string = Web::WebDriver::encode_canvas_element(canvas);
+
+        // 6. Return success with data encoded string.
+        async_screenshot_taken(move(encoded_string));
+    }));
+
+    return JsonValue {};
 }
 
 // 18.1 Print Page, https://w3c.github.io/webdriver/#dfn-print-page

Userland/Services/WebContent/WebDriverServer.ipc

@@ -7,4 +7,5 @@ endpoint WebDriverServer {
     script_executed(Web::WebDriver::Response response) =|
     actions_performed(Web::WebDriver::Response response) =|
     dialog_closed(Web::WebDriver::Response response) =|
+    screenshot_taken(Web::WebDriver::Response response) =|
 }

Userland/Services/WebDriver/Client.cpp

@@ -757,7 +757,7 @@ Web::WebDriver::Response Client::take_screenshot(Web::WebDriver::Parameters para
 {
     dbgln_if(WEBDRIVER_DEBUG, "Handling GET /session/<session_id>/screenshot");
     auto session = TRY(find_session_with_id(parameters[0]));
-    return session->web_content_connection().take_screenshot();
+    return session->take_screenshot();
 }
 
 // 17.2 Take Element Screenshot, https://w3c.github.io/webdriver/#dfn-take-element-screenshot
@@ -766,7 +766,7 @@ Web::WebDriver::Response Client::take_element_screenshot(Web::WebDriver::Paramet
 {
     dbgln_if(WEBDRIVER_DEBUG, "Handling GET /session/<session_id>/element/<element_id>/screenshot");
     auto session = TRY(find_session_with_id(parameters[0]));
-    return session->web_content_connection().take_element_screenshot(move(parameters[1]));
+    return session->take_element_screenshot(move(parameters[1]));
 }
 
 // 18.1 Print Page, https://w3c.github.io/webdriver/#dfn-print-page

Userland/Services/WebDriver/Session.cpp

@@ -326,4 +326,19 @@ Web::WebDriver::Response Session::accept_alert() const
         return web_content_connection().accept_alert();
     });
 }
+
+Web::WebDriver::Response Session::take_screenshot() const
+{
+    return perform_async_action(web_content_connection().on_screenshot_taken, [&]() {
+        return web_content_connection().take_screenshot();
+    });
+}
+
+Web::WebDriver::Response Session::take_element_screenshot(String element_id) const
+{
+    return perform_async_action(web_content_connection().on_screenshot_taken, [&]() {
+        return web_content_connection().take_element_screenshot(move(element_id));
+    });
+}
+
 }

Userland/Services/WebDriver/Session.h

@@ -84,6 +84,9 @@ public:
     Web::WebDriver::Response dismiss_alert() const;
     Web::WebDriver::Response accept_alert() const;
 
+    Web::WebDriver::Response take_screenshot() const;
+    Web::WebDriver::Response take_element_screenshot(String) const;
+
 private:
     using ServerPromise = Core::Promise<ErrorOr<void>>;
     ErrorOr<NonnullRefPtr<Core::LocalServer>> create_server(NonnullRefPtr<ServerPromise> promise);

Userland/Services/WebDriver/WebContentConnection.cpp

@@ -56,4 +56,10 @@ void WebContentConnection::dialog_closed(Web::WebDriver::Response const& respons
         on_dialog_closed(response);
 }
 
+void WebContentConnection::screenshot_taken(Web::WebDriver::Response const& response)
+{
+    if (on_screenshot_taken)
+        on_screenshot_taken(response);
+}
+
 }

Userland/Services/WebDriver/WebContentConnection.h

@@ -28,6 +28,7 @@ public:
     Function<void(Web::WebDriver::Response)> on_script_executed;
     Function<void(Web::WebDriver::Response)> on_actions_performed;
     Function<void(Web::WebDriver::Response)> on_dialog_closed;
+    Function<void(Web::WebDriver::Response)> on_screenshot_taken;
 
 private:
     virtual void die() override;
@@ -38,6 +39,7 @@ private:
     virtual void script_executed(Web::WebDriver::Response const&) override;
     virtual void actions_performed(Web::WebDriver::Response const&) override;
     virtual void dialog_closed(Web::WebDriver::Response const&) override;
+    virtual void screenshot_taken(Web::WebDriver::Response const&) override;
 };
 
 }