From 40b9d248be9744e595b94595acebee36b5effd38 Mon Sep 17 00:00:00 2001
From: Timothy Flynn <trflynn89@pm.me>
Date: Thu, 10 Nov 2022 14:57:16 -0500
Subject: [PATCH] LibWeb: Implement screenshot painting inside Web::WebDriver

This will allow for implementing the screenshot action closer to the
spec, as we can now use HTMLCanvasElement to encode the bitmap, and
capture the screenshot on the animation frame loop.
---
 Userland/Libraries/LibWeb/CMakeLists.txt      |  1 +
 .../Libraries/LibWeb/WebDriver/Screenshot.cpp | 76 +++++++++++++++++++
 .../Libraries/LibWeb/WebDriver/Screenshot.h   | 19 +++++
 3 files changed, 96 insertions(+)
 create mode 100644 Userland/Libraries/LibWeb/WebDriver/Screenshot.cpp
 create mode 100644 Userland/Libraries/LibWeb/WebDriver/Screenshot.h

diff --git a/Userland/Libraries/LibWeb/CMakeLists.txt b/Userland/Libraries/LibWeb/CMakeLists.txt
index 681ea3fb92..8ba66c0bca 100644
--- a/Userland/Libraries/LibWeb/CMakeLists.txt
+++ b/Userland/Libraries/LibWeb/CMakeLists.txt
@@ -442,6 +442,7 @@ set(SOURCES
     WebDriver/Error.cpp
     WebDriver/ExecuteScript.cpp
     WebDriver/Response.cpp
+    WebDriver/Screenshot.cpp
     WebGL/WebGLContextAttributes.cpp
     WebGL/WebGLContextEvent.cpp
     WebGL/WebGLRenderingContext.cpp
diff --git a/Userland/Libraries/LibWeb/WebDriver/Screenshot.cpp b/Userland/Libraries/LibWeb/WebDriver/Screenshot.cpp
new file mode 100644
index 0000000000..40cfada4bd
--- /dev/null
+++ b/Userland/Libraries/LibWeb/WebDriver/Screenshot.cpp
@@ -0,0 +1,76 @@
+/*
+ * Copyright (c) 2022, Tim Flynn <trflynn89@serenityos.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/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-encoding-a-canvas-as-base64
+static Response encode_canvas_element(HTML::HTMLCanvasElement const& 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.
+
+    // 2. If the canvas element’s bitmap has no pixels (i.e. either its horizontal dimension or vertical dimension is zero) then return error with error code unable to capture screen.
+    if (canvas.bitmap()->width() == 0 || canvas.bitmap()->height() == 0)
+        return Error::from_code(ErrorCode::UnableToCaptureScreen, "Captured screenshot is empty"sv);
+
+    // 3. Let file be a serialization of the canvas element’s bitmap as a file, using "image/png" as an argument.
+    // 4. Let data url be a data: URL representing file. [RFC2397]
+    auto data_url = canvas.to_data_url("image/png"sv, {});
+
+    // 5. Let index be the index of "," in data url.
+    auto index = data_url.find(',');
+    VERIFY(index.has_value());
+
+    // 6. Let encoded string be a substring of data url using (index + 1) as the start argument.
+    auto encoded_string = data_url.substring(*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;
+
+    element.document().window().animation_frame_callback_driver().add([&](auto) {
+        auto viewport_rect = page.top_level_browsing_context().viewport_rect();
+        rect.intersect(viewport_rect);
+
+        auto canvas_element = DOM::create_element(element.document(), HTML::TagNames::canvas, Namespace::HTML);
+        auto& canvas = verify_cast<HTML::HTMLCanvasElement>(*canvas_element);
+
+        if (!canvas.create_bitmap(rect.width(), rect.height())) {
+            encoded_string_or_error = Error::from_code(ErrorCode::UnableToCaptureScreen, "Unable to create a screenshot bitmap"sv);
+            return;
+        }
+
+        painter(rect, *canvas.bitmap());
+        encoded_string_or_error = encode_canvas_element(canvas);
+    });
+
+    Platform::EventLoopPlugin::the().spin_until([&]() { return encoded_string_or_error.has_value(); });
+    return encoded_string_or_error.release_value();
+}
+
+}
diff --git a/Userland/Libraries/LibWeb/WebDriver/Screenshot.h b/Userland/Libraries/LibWeb/WebDriver/Screenshot.h
new file mode 100644
index 0000000000..5cf1856aea
--- /dev/null
+++ b/Userland/Libraries/LibWeb/WebDriver/Screenshot.h
@@ -0,0 +1,19 @@
+/*
+ * Copyright (c) 2022, Tim Flynn <trflynn89@serenityos.org>
+ *
+ * SPDX-License-Identifier: BSD-2-Clause
+ */
+
+#pragma once
+
+#include <AK/Function.h>
+#include <LibGfx/Forward.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);
+
+}