WebDriver+Browser: Implement GET /session/{id}/screenshot

This doesn't follow the spec to a tee. Our OutOfProcessWebView already
has a bitmap that can be used as the window screenshot. Therefore, we
can bypass the steps that assume we need to access the window's frame
buffer in-flight.

We also don't create an HTMLCanvasElement. We would need a Document in
the WebDriver process to do so. Instead, we can still run the encoding
steps exactly as-is using the screenshot bitmap.

Userland/Applications/Browser/WebDriverConnection.cpp

@@ -255,4 +255,16 @@ Messages::WebDriverSessionClient::GetElementTagNameResponse WebDriverConnection:
     return { "" };
 }
 
+Messages::WebDriverSessionClient::TakeScreenshotResponse WebDriverConnection::take_screenshot()
+{
+    dbgln_if(WEBDRIVER_DEBUG, "WebDriverConnection: take_screenshot");
+    if (auto browser_window = m_browser_window.strong_ref()) {
+        auto& tab = browser_window->active_tab();
+        if (tab.on_take_screenshot)
+            return { tab.on_take_screenshot() };
+    }
+
+    return { {} };
+}
+
 }

Userland/Applications/Browser/WebDriverConnection.h

@@ -61,6 +61,7 @@ public:
     virtual Messages::WebDriverSessionClient::GetComputedValueForElementResponse get_computed_value_for_element(i32 element_id, String const& property_name) override;
     virtual Messages::WebDriverSessionClient::GetElementTextResponse get_element_text(i32 element_id) override;
     virtual Messages::WebDriverSessionClient::GetElementTagNameResponse get_element_tag_name(i32 element_id) override;
+    virtual Messages::WebDriverSessionClient::TakeScreenshotResponse take_screenshot() override;
 
 private:
     WebDriverConnection(NonnullOwnPtr<Core::Stream::LocalSocket> socket, NonnullRefPtr<BrowserWindow> browser_window);

Userland/Applications/Browser/WebDriverSessionClient.ipc

@@ -2,6 +2,7 @@
 #include <AK/Vector.h>
 #include <LibGfx/Point.h>
 #include <LibGfx/Rect.h>
+#include <LibGfx/ShareableBitmap.h>
 #include <LibGfx/Size.h>
 #include <LibWeb/Cookie/Cookie.h>
 #include <LibWeb/Cookie/ParsedCookie.h>
@@ -33,5 +34,5 @@ endpoint WebDriverSessionClient {
     get_computed_value_for_element(i32 element_id, String property_name) => (String computed_value)
     get_element_text(i32 element_id) => (String text)
     get_element_tag_name(i32 element_id) => (String tag_name)
-
+    take_screenshot() => (Gfx::ShareableBitmap data)
 }

Userland/Services/WebDriver/Client.cpp

@@ -54,6 +54,7 @@ Vector<Client::Route> Client::s_routes = {
     { HTTP::HttpRequest::Method::POST, { "session", ":session_id", "cookie" }, &Client::handle_add_cookie },
     { HTTP::HttpRequest::Method::DELETE, { "session", ":session_id", "cookie", ":name" }, &Client::handle_delete_cookie },
     { HTTP::HttpRequest::Method::DELETE, { "session", ":session_id", "cookie" }, &Client::handle_delete_all_cookies },
+    { HTTP::HttpRequest::Method::GET, { "session", ":session_id", "screenshot" }, &Client::handle_take_screenshot },
 };
 
 Client::Client(NonnullOwnPtr<Core::Stream::BufferedTCPSocket> socket, Core::Object* parent)
@@ -732,4 +733,14 @@ ErrorOr<JsonValue, WebDriverError> Client::handle_delete_all_cookies(Vector<Stri
     return make_json_value(result);
 }
 
+// 17.1 Take Screenshot, https://w3c.github.io/webdriver/#take-screenshot
+// GET /session/{session id}/screenshot
+ErrorOr<JsonValue, WebDriverError> Client::handle_take_screenshot(Vector<StringView> const& parameters, JsonValue const&)
+{
+    dbgln_if(WEBDRIVER_DEBUG, "Handling GET /session/<session_id>/screenshot");
+    auto* session = TRY(find_session_with_id(parameters[0]));
+    auto result = TRY(session->take_screenshot());
+    return make_json_value(result);
+}
+
 }

Userland/Services/WebDriver/Client.h

@@ -79,6 +79,7 @@ private:
     ErrorOr<JsonValue, WebDriverError> handle_add_cookie(Vector<StringView> const&, JsonValue const& payload);
     ErrorOr<JsonValue, WebDriverError> handle_delete_cookie(Vector<StringView> const&, JsonValue const& payload);
     ErrorOr<JsonValue, WebDriverError> handle_delete_all_cookies(Vector<StringView> const&, JsonValue const& payload);
+    ErrorOr<JsonValue, WebDriverError> handle_take_screenshot(Vector<StringView> const&, JsonValue const& payload);
 
     ErrorOr<Session*, WebDriverError> find_session_with_id(StringView session_id);
     JsonValue make_json_value(JsonValue const&);

Userland/Services/WebDriver/Session.cpp

@@ -10,11 +10,14 @@
 #include "Session.h"
 #include "BrowserConnection.h"
 #include "Client.h"
+#include <AK/Base64.h>
 #include <AK/NumericLimits.h>
 #include <AK/Time.h>
+#include <AK/URL.h>
 #include <LibCore/LocalServer.h>
 #include <LibCore/Stream.h>
 #include <LibCore/System.h>
+#include <LibGfx/PNGWriter.h>
 #include <LibGfx/Point.h>
 #include <LibGfx/Rect.h>
 #include <LibGfx/Size.h>
@@ -1123,4 +1126,52 @@ ErrorOr<JsonValue, WebDriverError> Session::delete_all_cookies()
     return JsonValue();
 }
 
+// https://w3c.github.io/webdriver/#dfn-encoding-a-canvas-as-base64
+static ErrorOr<String, WebDriverError> encode_bitmap_as_canvas_element(Gfx::Bitmap const& bitmap)
+{
+    // 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 (bitmap.width() == 0 || bitmap.height() == 0)
+        return WebDriverError::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.
+    auto file = Gfx::PNGWriter::encode(bitmap);
+
+    // 4. Let data url be a data: URL representing file. [RFC2397]
+    auto data_url = AK::URL::create_with_data("image/png"sv, encode_base64(file), true).to_string();
+
+    // 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 encoded_string;
+}
+
+// 17.1 Take Screenshot, https://w3c.github.io/webdriver/#take-screenshot
+ErrorOr<JsonValue, WebDriverError> Session::take_screenshot()
+{
+    // 1. If the current top-level browsing context is no longer open, return error with error code no such window.
+    TRY(check_for_open_top_level_browsing_context_or_return_error());
+
+    // 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.
+    auto screenshot = m_browser_connection->take_screenshot();
+    if (!screenshot.is_valid())
+        return WebDriverError::from_code(ErrorCode::UnableToCaptureScreen, "Unable to capture screenshot"sv);
+
+    //     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 encoded_string = TRY(encode_bitmap_as_canvas_element(*screenshot.bitmap()));
+
+    // 3. Return success with data encoded string.
+    return encoded_string;
+}
+
 }

Userland/Services/WebDriver/Session.h

@@ -68,6 +68,7 @@ public:
     ErrorOr<JsonValue, WebDriverError> add_cookie(JsonValue const& payload);
     ErrorOr<JsonValue, WebDriverError> delete_cookie(StringView const& name);
     ErrorOr<JsonValue, WebDriverError> delete_all_cookies();
+    ErrorOr<JsonValue, WebDriverError> take_screenshot();
 
 private:
     void delete_cookies(Optional<StringView> const& name = {});

Userland/Services/WebDriver/main.cpp

@@ -35,7 +35,7 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
         return 1;
     }
 
-    TRY(Core::System::pledge("stdio accept rpath inet unix proc exec fattr"));
+    TRY(Core::System::pledge("stdio accept rpath recvfd inet unix proc exec fattr"));
 
     Core::EventLoop loop;
 
@@ -70,6 +70,6 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
     TRY(Core::System::unveil("/tmp", "rwc"));
     TRY(Core::System::unveil(nullptr, nullptr));
 
-    TRY(Core::System::pledge("stdio accept rpath unix proc exec fattr"));
+    TRY(Core::System::pledge("stdio accept rpath recvfd unix proc exec fattr"));
     return loop.exec();
 }