Implement TakeElementScreenshot WebDriver command

12 files changed, 134 insertions, 28 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 89f14314016..69a47ff87a0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -5621,6 +5621,7 @@ dependencies = [
  "serde_json 1.0.13 (registry+https://github.com/rust-lang/crates.io-index)",
  "servo_config 0.0.1",
  "servo_url 0.0.1",
+ "style_traits 0.0.1",
  "url 2.1.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "uuid 0.7.4 (registry+https://github.com/rust-lang/crates.io-index)",
  "webdriver 0.40.0 (registry+https://github.com/rust-lang/crates.io-index)",
diff --git a/components/compositing/compositor.rs b/components/compositing/compositor.rs
index 7ac80ee4832..f0474363fa4 100644
--- a/components/compositing/compositor.rs
+++ b/components/compositing/compositor.rs
@@ -14,7 +14,7 @@ use crate::CompositionPipeline;
 use crate::SendableFrameTree;
 use crossbeam_channel::Sender;
 use embedder_traits::Cursor;
-use euclid::{Point2D, Scale, Vector2D};
+use euclid::{Point2D, Rect, Scale, Vector2D};
 use gfx_traits::Epoch;
 #[cfg(feature = "gl")]
 use image::{DynamicImage, ImageFormat};
@@ -442,8 +442,8 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
                 self.touch_handler.on_event_processed(result);
             },
 
-            (Msg::CreatePng(reply), ShutdownState::NotShuttingDown) => {
-                let res = self.composite_specific_target(CompositeTarget::WindowAndPng);
+            (Msg::CreatePng(rect, reply), ShutdownState::NotShuttingDown) => {
+                let res = self.composite_specific_target(CompositeTarget::WindowAndPng, rect);
                 if let Err(ref e) = res {
                     info!("Error retrieving PNG: {:?}", e);
                 }
@@ -1229,7 +1229,7 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
 
     pub fn composite(&mut self) {
         let target = self.composite_target;
-        match self.composite_specific_target(target) {
+        match self.composite_specific_target(target, None) {
             Ok(_) => {
                 if self.output_file.is_some() || self.exit_after_load {
                     println!("Shutting down the Constellation after generating an output file or exit flag specified");
@@ -1256,6 +1256,7 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
     fn composite_specific_target(
         &mut self,
         target: CompositeTarget,
+        rect: Option<Rect<f32, CSSPixel>>,
     ) -> Result<Option<Image>, UnableToComposite> {
         let size = self.embedder_coordinates.framebuffer.to_u32();
 
@@ -1347,6 +1348,22 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
             }
         }
 
+        let (x, y, width, height) = match rect {
+            Some(rect) => {
+                let rect = self.device_pixels_per_page_px().transform_rect(&rect);
+
+                let x = rect.origin.x as i32;
+                // We need to convert to the bottom-left origin coordinate
+                // system used by OpenGL
+                let y = (size.height as f32 - rect.origin.y - rect.size.height) as i32;
+                let w = rect.size.width as u32;
+                let h = rect.size.height as u32;
+
+                (x, y, w, h)
+            },
+            None => (0, 0, size.width, size.height),
+        };
+
         let rv = match target {
             CompositeTarget::Window => None,
             #[cfg(feature = "gl")]
@@ -1354,8 +1371,10 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
                 let img = gl::draw_img(
                     &*self.window.gl(),
                     rt_info,
-                    FramebufferUintLength::new(size.width),
-                    FramebufferUintLength::new(size.height),
+                    x,
+                    y,
+                    FramebufferUintLength::new(width),
+                    FramebufferUintLength::new(height),
                 );
                 Some(Image {
                     width: img.width(),
@@ -1378,8 +1397,10 @@ impl<Window: WindowMethods + ?Sized> IOCompositor<Window> {
                                 let img = gl::draw_img(
                                     gl,
                                     rt_info,
-                                    FramebufferUintLength::new(size.width),
-                                    FramebufferUintLength::new(size.height),
+                                    x,
+                                    y,
+                                    FramebufferUintLength::new(width),
+                                    FramebufferUintLength::new(height),
                                 );
                                 let dynamic_image = DynamicImage::ImageRgb8(img);
                                 if let Err(e) = dynamic_image.write_to(&mut file, ImageFormat::PNG)
diff --git a/components/compositing/compositor_thread.rs b/components/compositing/compositor_thread.rs
index f095ab65b1c..e10aff78844 100644
--- a/components/compositing/compositor_thread.rs
+++ b/components/compositing/compositor_thread.rs
@@ -8,6 +8,7 @@ use crate::compositor::CompositingReason;
 use crate::SendableFrameTree;
 use crossbeam_channel::{Receiver, Sender};
 use embedder_traits::EventLoopWaker;
+use euclid::Rect;
 use gfx_traits::Epoch;
 use ipc_channel::ipc::IpcSender;
 use msg::constellation_msg::{PipelineId, TopLevelBrowsingContextId};
@@ -17,6 +18,7 @@ use profile_traits::time;
 use script_traits::{AnimationState, ConstellationMsg, EventResult, MouseButton, MouseEventType};
 use std::fmt::{Debug, Error, Formatter};
 use style_traits::viewport::ViewportConstraints;
+use style_traits::CSSPixel;
 use webrender_api;
 use webrender_api::units::{DeviceIntPoint, DeviceIntSize};
 use webvr_traits::WebVRMainThreadHeartbeat;
@@ -80,7 +82,7 @@ pub enum Msg {
     /// Script has handled a touch event, and either prevented or allowed default actions.
     TouchEventProcessed(EventResult),
     /// Composite to a PNG file and return the Image over a passed channel.
-    CreatePng(IpcSender<Option<Image>>),
+    CreatePng(Option<Rect<f32, CSSPixel>>, IpcSender<Option<Image>>),
     /// Alerts the compositor that the viewport has been constrained in some manner
     ViewportConstrained(PipelineId, ViewportConstraints),
     /// A reply to the compositor asking if the output image is stable.
diff --git a/components/compositing/gl.rs b/components/compositing/gl.rs
index 44d69254692..45d33f19a7d 100644
--- a/components/compositing/gl.rs
+++ b/components/compositing/gl.rs
@@ -82,6 +82,8 @@ pub fn initialize_png(
 pub fn draw_img(
     gl: &dyn gl::Gl,
     render_target_info: RenderTargetInfo,
+    x: i32,
+    y: i32,
     width: FramebufferUintLength,
     height: FramebufferUintLength,
 ) -> RgbImage {
@@ -96,8 +98,8 @@ pub fn draw_img(
     gl.bind_vertex_array(0);
 
     let mut pixels = gl.read_pixels(
-        0,
-        0,
+        x,
+        y,
         width as gl::GLsizei,
         height as gl::GLsizei,
         gl::RGB,
diff --git a/components/constellation/constellation.rs b/components/constellation/constellation.rs
index b443163f8d5..f43c383d3de 100644
--- a/components/constellation/constellation.rs
+++ b/components/constellation/constellation.rs
@@ -3499,9 +3499,9 @@ where
                         y,
                     ));
             },
-            WebDriverCommandMsg::TakeScreenshot(_, reply) => {
+            WebDriverCommandMsg::TakeScreenshot(_, rect, reply) => {
                 self.compositor_proxy
-                    .send(ToCompositorMsg::CreatePng(reply));
+                    .send(ToCompositorMsg::CreatePng(rect, reply));
             },
         }
     }
diff --git a/components/script/script_thread.rs b/components/script/script_thread.rs
index e5bfd5ebe8f..3f547d7cfd6 100644
--- a/components/script/script_thread.rs
+++ b/components/script/script_thread.rs
@@ -2198,6 +2198,14 @@ impl ScriptThread {
             WebDriverScriptCommand::GetElementRect(node_id, reply) => {
                 webdriver_handlers::handle_get_rect(&*documents, pipeline_id, node_id, reply)
             },
+            WebDriverScriptCommand::GetBoundingClientRect(node_id, reply) => {
+                webdriver_handlers::handle_get_bounding_client_rect(
+                    &*documents,
+                    pipeline_id,
+                    node_id,
+                    reply,
+                )
+            },
             WebDriverScriptCommand::GetElementText(node_id, reply) => {
                 webdriver_handlers::handle_get_text(&*documents, pipeline_id, node_id, reply)
             },
diff --git a/components/script/webdriver_handlers.rs b/components/script/webdriver_handlers.rs
index 22d3e56fcaf..564dc7ceb4d 100644
--- a/components/script/webdriver_handlers.rs
+++ b/components/script/webdriver_handlers.rs
@@ -3,6 +3,7 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/. */
 
 use crate::dom::bindings::codegen::Bindings::CSSStyleDeclarationBinding::CSSStyleDeclarationMethods;
+use crate::dom::bindings::codegen::Bindings::DOMRectBinding::DOMRectMethods;
 use crate::dom::bindings::codegen::Bindings::DocumentBinding::DocumentMethods;
 use crate::dom::bindings::codegen::Bindings::ElementBinding::ElementMethods;
 use crate::dom::bindings::codegen::Bindings::HTMLElementBinding::HTMLElementMethods;
@@ -776,6 +777,30 @@ pub fn handle_get_rect(
         .unwrap();
 }
 
+pub fn handle_get_bounding_client_rect(
+    documents: &Documents,
+    pipeline: PipelineId,
+    element_id: String,
+    reply: IpcSender<Result<Rect<f32>, ErrorStatus>>,
+) {
+    reply
+        .send(
+            find_node_by_unique_id(documents, pipeline, element_id).and_then(|node| match node
+                .downcast::<Element>(
+            ) {
+                Some(element) => {
+                    let rect = element.GetBoundingClientRect();
+                    Ok(Rect::new(
+                        Point2D::new(rect.X() as f32, rect.Y() as f32),
+                        Size2D::new(rect.Width() as f32, rect.Height() as f32),
+                    ))
+                },
+                None => Err(ErrorStatus::UnknownError),
+            }),
+        )
+        .unwrap();
+}
+
 pub fn handle_get_text(
     documents: &Documents,
     pipeline: PipelineId,
diff --git a/components/script_traits/lib.rs b/components/script_traits/lib.rs
index 790c49eff94..f0aa1516ca2 100644
--- a/components/script_traits/lib.rs
+++ b/components/script_traits/lib.rs
@@ -25,10 +25,7 @@ use canvas_traits::webgl::WebGLPipeline;
 use crossbeam_channel::{Receiver, RecvTimeoutError, Sender};
 use devtools_traits::{DevtoolScriptControlMsg, ScriptToDevtoolsControlMsg, WorkerId};
 use embedder_traits::{Cursor, EventLoopWaker};
-use euclid::{
-    default::{Point2D, Rect},
-    Length, Scale, Size2D, Vector2D,
-};
+use euclid::{default::Point2D, Length, Rect, Scale, Size2D, UnknownUnit, Vector2D};
 use gfx_traits::Epoch;
 use http::HeaderMap;
 use hyper::Method;
@@ -294,7 +291,7 @@ pub enum ConstellationControlMsg {
     /// Sends a DOM event.
     SendEvent(PipelineId, CompositorEvent),
     /// Notifies script of the viewport.
-    Viewport(PipelineId, Rect<f32>),
+    Viewport(PipelineId, Rect<f32, UnknownUnit>),
     /// Notifies script of a new set of scroll offsets.
     SetScrollState(
         PipelineId,
@@ -807,7 +804,11 @@ pub enum WebDriverCommandMsg {
         IpcSender<WindowSizeData>,
     ),
     /// Take a screenshot of the window.
-    TakeScreenshot(TopLevelBrowsingContextId, IpcSender<Option<Image>>),
+    TakeScreenshot(
+        TopLevelBrowsingContextId,
+        Option<Rect<f32, CSSPixel>>,
+        IpcSender<Option<Image>>,
+    ),
 }
 
 /// Messages to the constellation.
diff --git a/components/script_traits/webdriver_msg.rs b/components/script_traits/webdriver_msg.rs
index e7087d8825d..9716bd220d6 100644
--- a/components/script_traits/webdriver_msg.rs
+++ b/components/script_traits/webdriver_msg.rs
@@ -74,6 +74,7 @@ pub enum WebDriverScriptCommand {
     GetElementRect(String, IpcSender<Result<Rect<f64>, ErrorStatus>>),
     GetElementTagName(String, IpcSender<Result<String, ErrorStatus>>),
     GetElementText(String, IpcSender<Result<String, ErrorStatus>>),
+    GetBoundingClientRect(String, IpcSender<Result<Rect<f32>, ErrorStatus>>),
     GetBrowsingContextId(
         WebDriverFrameId,
         IpcSender<Result<BrowsingContextId, ErrorStatus>>,
diff --git a/components/webdriver_server/Cargo.toml b/components/webdriver_server/Cargo.toml
index 7fd9e7dae69..31c980b9b38 100644
--- a/components/webdriver_server/Cargo.toml
+++ b/components/webdriver_server/Cargo.toml
@@ -29,6 +29,7 @@ serde_json = "1"
 script_traits = {path = "../script_traits"}
 servo_config = {path = "../config"}
 servo_url = {path = "../url"}
+style_traits = {path = "../style_traits"}
 url = "2.0"
 uuid = {version = "0.7", features = ["v4"]}
 webdriver = "0.40"
diff --git a/components/webdriver_server/lib.rs b/components/webdriver_server/lib.rs
index a40b766a66a..0bd7d140c32 100644
--- a/components/webdriver_server/lib.rs
+++ b/components/webdriver_server/lib.rs
@@ -20,7 +20,7 @@ use crate::actions::InputSourceState;
 use base64;
 use capabilities::ServoCapabilities;
 use crossbeam_channel::Sender;
-use euclid::Size2D;
+use euclid::{Rect, Size2D};
 use hyper::Method;
 use image::{DynamicImage, ImageFormat, RgbImage};
 use ipc_channel::ipc::{self, IpcReceiver, IpcSender};
@@ -44,6 +44,7 @@ use std::mem;
 use std::net::{SocketAddr, SocketAddrV4};
 use std::thread;
 use std::time::Duration;
+use style_traits::CSSPixel;
 use uuid::Uuid;
 use webdriver::actions::ActionSequence;
 use webdriver::capabilities::{Capabilities, CapabilitiesMatching};
@@ -1445,16 +1446,20 @@ impl Handler {
         Ok(WebDriverResponse::Void)
     }
 
-    fn handle_take_screenshot(&self) -> WebDriverResult<WebDriverResponse> {
+    fn take_screenshot(&self, rect: Option<Rect<f32, CSSPixel>>) -> WebDriverResult<String> {
         let mut img = None;
-        let top_level_id = self.session()?.top_level_browsing_context_id;
 
         let interval = 1000;
-        let iterations = 30_000 / interval;
+        let iterations = 30000 / interval;
 
         for _ in 0..iterations {
             let (sender, receiver) = ipc::channel().unwrap();
-            let cmd_msg = WebDriverCommandMsg::TakeScreenshot(top_level_id, sender);
+
+            let cmd_msg = WebDriverCommandMsg::TakeScreenshot(
+                self.session()?.top_level_browsing_context_id,
+                rect,
+                sender,
+            );
             self.constellation_chan
                 .send(ConstellationMsg::WebDriverCommand(cmd_msg))
                 .unwrap();
@@ -1464,7 +1469,7 @@ impl Handler {
                 break;
             };
 
-            thread::sleep(Duration::from_millis(interval))
+            thread::sleep(Duration::from_millis(interval));
         }
 
         let img = match img {
@@ -1483,19 +1488,50 @@ impl Handler {
             PixelFormat::RGB8,
             "Unexpected screenshot pixel format"
         );
-        let rgb = RgbImage::from_raw(img.width, img.height, img.bytes.to_vec()).unwrap();
 
+        let rgb = RgbImage::from_raw(img.width, img.height, img.bytes.to_vec()).unwrap();
         let mut png_data = Vec::new();
         DynamicImage::ImageRgb8(rgb)
             .write_to(&mut png_data, ImageFormat::PNG)
             .unwrap();
 
-        let encoded = base64::encode(&png_data);
+        Ok(base64::encode(&png_data))
+    }
+
+    fn handle_take_screenshot(&self) -> WebDriverResult<WebDriverResponse> {
+        let encoded = self.take_screenshot(None)?;
+
         Ok(WebDriverResponse::Generic(ValueResponse(
             serde_json::to_value(encoded)?,
         )))
     }
 
+    fn handle_take_element_screenshot(
+        &self,
+        element: &WebElement,
+    ) -> WebDriverResult<WebDriverResponse> {
+        let (sender, receiver) = ipc::channel().unwrap();
+
+        let command = WebDriverScriptCommand::GetBoundingClientRect(element.to_string(), sender);
+        self.browsing_context_script_command(command)?;
+
+        match receiver.recv().unwrap() {
+            Ok(rect) => {
+                let encoded = self.take_screenshot(Some(Rect::from_untyped(&rect)))?;
+
+                Ok(WebDriverResponse::Generic(ValueResponse(
+                    serde_json::to_value(encoded)?,
+                )))
+            },
+            Err(_) => {
+                return Err(WebDriverError::new(
+                    ErrorStatus::StaleElementReference,
+                    "Element not found",
+                ));
+            },
+        }
+    }
+
     fn handle_get_prefs(
         &self,
         parameters: &GetPrefsParameters,
@@ -1635,6 +1671,9 @@ impl WebDriverHandler<ServoExtensionRoute> for Handler {
             WebDriverCommand::GetTimeouts => self.handle_get_timeouts(),
             WebDriverCommand::SetTimeouts(ref x) => self.handle_set_timeouts(x),
             WebDriverCommand::TakeScreenshot => self.handle_take_screenshot(),
+            WebDriverCommand::TakeElementScreenshot(ref x) => {
+                self.handle_take_element_screenshot(x)
+            },
             WebDriverCommand::Extension(ref extension) => match *extension {
                 ServoExtensionCommand::GetPrefs(ref x) => self.handle_get_prefs(x),
                 ServoExtensionCommand::SetPrefs(ref x) => self.handle_set_prefs(x),
diff --git a/tests/wpt/metadata/webdriver/tests/take_element_screenshot/screenshot.py.ini b/tests/wpt/metadata/webdriver/tests/take_element_screenshot/screenshot.py.ini
index f959c3368af..73c9205fac2 100644
--- a/tests/wpt/metadata/webdriver/tests/take_element_screenshot/screenshot.py.ini
+++ b/tests/wpt/metadata/webdriver/tests/take_element_screenshot/screenshot.py.ini
@@ -1,2 +1,7 @@
 [screenshot.py]
-  disabled: Unimplemented WebDriver command
+  [test_no_browsing_context]
+    expected: ERROR
+
+  [test_format_and_dimensions]
+    expected: FAIL
+