Add screenshot features for linux and win32

Cargo.toml

@@ -50,6 +50,7 @@ chrono = "0.4.19"
 tempfile = "3.2.0"
 
 [target."cfg(target_os = \"linux\")".dependencies]
+cairo-rs = { version = "0.9", features = ["png"] }
 webkit2gtk = { version = "0.11", features = [ "v2_10" ] }
 webkit2gtk-sys = "0.13.0"
 gio = "0.9"
@@ -65,6 +66,7 @@ windows-webview2 = { version = "0.1", optional = true }
 windows = { version = "0.7", optional = true }
 
 [target."cfg(any(target_os = \"ios\", target_os = \"macos\"))".dependencies]
+block = "0.1"
 cocoa = "0.24"
 core-graphics = "0.22"
 objc = "0.2"

examples/screenshot.rs

@@ -0,0 +1,106 @@
+// Copyright 2019-2021 Tauri Programme within The Commons Conservancy
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-License-Identifier: MIT
+
+fn main() -> wry::Result<()> {
+  use std::{fs::File, io::Write};
+  use wry::{
+    application::{
+      event::{Event, StartCause, WindowEvent},
+      event_loop::{ControlFlow, EventLoop},
+      menu::{Menu, MenuItem, MenuType},
+      window::WindowBuilder,
+    },
+    webview::{ScreenshotRegion, WebViewBuilder},
+  };
+
+  let custom_screenshot_visible = MenuItem::new("Visible");
+  let custom_screenshot_fulldocument = MenuItem::new("Full document");
+  let custom_screenshot_visible_id = custom_screenshot_visible.id();
+  let custom_screenshot_fulldocument_id = custom_screenshot_fulldocument.id();
+
+  // macOS require to have at least Copy, Paste, Select all etc..
+  // to works fine. You should always add them.
+  #[cfg(any(target_os = "linux", target_os = "macos"))]
+  let menu = vec![
+    Menu::new("File", vec![MenuItem::CloseWindow]),
+    Menu::new(
+      "Edit",
+      vec![
+        MenuItem::Undo,
+        MenuItem::Redo,
+        MenuItem::Separator,
+        MenuItem::Cut,
+        MenuItem::Copy,
+        MenuItem::Paste,
+        MenuItem::Separator,
+        MenuItem::SelectAll,
+      ],
+    ),
+    Menu::new(
+      // on macOS first menu is always app name
+      "Screenshot",
+      vec![custom_screenshot_visible, custom_screenshot_fulldocument],
+    ),
+  ];
+
+  // Attention, Windows only support custom menu for now.
+  // If we add any `MenuItem::*` they'll not render
+  // We need to use custom menu with `Menu::new()` and catch
+  // the events in the EventLoop.
+  #[cfg(target_os = "windows")]
+  let menu = vec![Menu::new(
+    "Screenshot",
+    vec![custom_screenshot_visible, custom_screenshot_fulldocument],
+  )];
+
+  // Build our event loop
+  let event_loop = EventLoop::new();
+  // Build the window
+  let window = WindowBuilder::new()
+    .with_title("Hello World")
+    .with_menu(menu)
+    .build(&event_loop)?;
+  // Build the webview
+  let webview = WebViewBuilder::new(window)?
+    .with_url("https://html5test.com")?
+    .build()?;
+
+  // launch WRY process
+  event_loop.run(move |event, _, control_flow| {
+    *control_flow = ControlFlow::Wait;
+
+    match event {
+      Event::NewEvents(StartCause::Init) => println!("Wry has started!"),
+      Event::WindowEvent {
+        event: WindowEvent::CloseRequested,
+        ..
+      } => *control_flow = ControlFlow::Exit,
+      // Catch menu events
+      Event::MenuEvent {
+        menu_id,
+        origin: MenuType::Menubar,
+      } => {
+        let on_screenshot = |image: wry::Result<Vec<u8>>| {
+          let image = image.expect("No image?");
+          let mut file = File::create("baaaaar.png").expect("Couldn't create the dang file");
+          file
+            .write(image.as_slice())
+            .expect("Couldn't write the dang file");
+        };
+        if menu_id == custom_screenshot_visible_id {
+          webview
+            .screenshot(ScreenshotRegion::Visible, on_screenshot)
+            .expect("Unable to screenshot");
+        }
+        if menu_id == custom_screenshot_fulldocument_id {
+          webview
+            .screenshot(ScreenshotRegion::FullDocument, on_screenshot)
+            .expect("Unable to screenshot");
+        }
+        println!("Clicked on {:?}", menu_id);
+      }
+      _ => (),
+    }
+  });
+}

src/lib.rs

@@ -108,6 +108,12 @@ pub type Result<T> = std::result::Result<T, Error>;
 /// Errors returned by wry.
 #[derive(Error, Debug)]
 pub enum Error {
+  #[cfg(target_os = "linux")]
+  #[error(transparent)]
+  CairoError(#[from] cairo::Error),
+  #[cfg(target_os = "linux")]
+  #[error(transparent)]
+  CairoIoError(#[from] cairo::IoError),
   #[cfg(target_os = "linux")]
   #[error(transparent)]
   GlibError(#[from] glib::Error),

src/webview/mod.rs

@@ -346,6 +346,13 @@ impl WebView {
     Ok(())
   }
 
+  pub fn screenshot<F>(&self, region: ScreenshotRegion, handler: F) -> Result<()>
+  where
+    F: Fn(Result<Vec<u8>>) -> () + 'static + Send,
+  {
+    self.webview.screenshot(region, handler)
+  }
+
   /// Resize the WebView manually. This is required on Windows because its WebView API doesn't
   /// provide a way to resize automatically.
   pub fn resize(&self) -> Result<()> {
@@ -440,6 +447,13 @@ impl RpcResponse {
   }
 }
 
+#[derive(Debug, Clone)]
+#[non_exhaustive]
+pub enum ScreenshotRegion {
+  Visible,
+  FullDocument,
+}
+
 /// An event enumeration sent to [`FileDropHandler`].
 #[derive(Debug, Serialize, Clone)]
 pub enum FileDropEvent {

src/webview/webkitgtk/mod.rs

@@ -2,17 +2,20 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-License-Identifier: MIT
 
+use std::convert::TryFrom;
+
 use std::{path::PathBuf, rc::Rc};
 
+use cairo::ImageSurface;
 use gdk::{WindowEdge, WindowExt, RGBA};
 use gio::Cancellable;
 use glib::{signal::Inhibit, Bytes, Cast, FileError};
 use gtk::{BoxExt, ContainerExt, WidgetExt};
 use url::Url;
 use webkit2gtk::{
-  SecurityManagerExt, SettingsExt, URISchemeRequestExt, UserContentInjectedFrames,
-  UserContentManager, UserContentManagerExt, UserScript, UserScriptInjectionTime,
-  WebContextBuilder, WebContextExt, WebView, WebViewExt, WebViewExtManual,
+  SecurityManagerExt, SettingsExt, SnapshotOptions, SnapshotRegion, URISchemeRequestExt,
+  UserContentInjectedFrames, UserContentManager, UserContentManagerExt, UserScript,
+  UserScriptInjectionTime, WebContextBuilder, WebContextExt, WebView, WebViewExt, WebViewExtManual,
   WebsiteDataManagerBuilder,
 };
 use webkit2gtk_sys::{
@@ -21,7 +24,7 @@ use webkit2gtk_sys::{
 
 use crate::{
   application::{platform::unix::*, window::Window},
-  webview::{mimetype::MimeType, FileDropEvent, RpcRequest, RpcResponse},
+  webview::{mimetype::MimeType, FileDropEvent, RpcRequest, RpcResponse, ScreenshotRegion},
   Error, Result,
 };
 
@@ -213,6 +216,47 @@ impl InnerWebView {
     let _ = self.eval("window.print()");
   }
 
+  pub fn screenshot<F>(&self, region: ScreenshotRegion, handler: F) -> Result<()>
+  where
+    F: Fn(Result<Vec<u8>>) -> () + 'static + Send,
+  {
+    let cancellable: Option<&Cancellable> = None;
+    let cb = move |result: std::result::Result<cairo::Surface, glib::Error>| match result {
+      Ok(surface) => match ImageSurface::try_from(surface) {
+        Ok(image) => {
+          let mut bytes = Vec::new();
+          match image.write_to_png(&mut bytes) {
+            Ok(_) => handler(Ok(bytes)),
+            Err(err) => handler(Err(Error::CairoIoError(err))),
+          }
+        }
+        Err(_) => handler(Err(Error::CairoError(cairo::Error::SurfaceTypeMismatch))),
+      },
+      Err(err) => handler(Err(Error::GlibError(err))),
+    };
+
+    match region {
+      ScreenshotRegion::FullDocument => {
+        self.webview.get_snapshot(
+          SnapshotRegion::FullDocument,
+          SnapshotOptions::NONE,
+          cancellable,
+          cb,
+        );
+      }
+      ScreenshotRegion::Visible => {
+        self.webview.get_snapshot(
+          SnapshotRegion::Visible,
+          SnapshotOptions::NONE,
+          cancellable,
+          cb,
+        );
+      }
+    };
+
+    Ok(())
+  }
+
   pub fn eval(&self, js: &str) -> Result<()> {
     let cancellable: Option<&Cancellable> = None;
     self.webview.run_javascript(js, cancellable, |_| ());

src/webview/webview2/win32/mod.rs

@@ -5,25 +5,30 @@
 mod file_drop;
 
 use crate::{
-  webview::{mimetype::MimeType, FileDropEvent, RpcRequest, RpcResponse},
-  Result,
+  application::{
+    event_loop::{ControlFlow, EventLoop},
+    platform::{run_return::EventLoopExtRunReturn, windows::WindowExtWindows},
+    window::Window,
+  },
+  webview::{mimetype::MimeType, FileDropEvent, RpcRequest, RpcResponse, ScreenshotRegion},
+  Error, Result,
 };
 
 use file_drop::FileDropController;
 
-use std::{collections::HashSet, os::raw::c_void, path::PathBuf, rc::Rc};
+use std::{
+  collections::HashSet,
+  io::{Read, Seek, SeekFrom},
+  os::raw::c_void,
+  path::PathBuf,
+  rc::Rc,
+};
 
 use once_cell::unsync::OnceCell;
 use url::Url;
 use webview2::{Controller, PermissionKind, PermissionState, WebView};
 use winapi::{shared::windef::HWND, um::winuser::GetClientRect};
 
-use crate::application::{
-  event_loop::{ControlFlow, EventLoop},
-  platform::{run_return::EventLoopExtRunReturn, windows::WindowExtWindows},
-  window::Window,
-};
-
 pub struct InnerWebView {
   controller: Rc<OnceCell<Controller>>,
   webview: Rc<OnceCell<WebView>>,
@@ -241,6 +246,35 @@ impl InnerWebView {
     let _ = self.eval("window.print()");
   }
 
+  pub fn screenshot<F>(&self, region: ScreenshotRegion, handler: F) -> Result<()>
+  where
+    F: Fn(Result<Vec<u8>>) -> () + 'static + Send,
+  {
+    if let Some(w) = self.webview.get() {
+      match region {
+        ScreenshotRegion::Visible => {
+          let mut stream = webview2::Stream::from_bytes(&[]);
+          w.capture_preview(
+            webview2::CapturePreviewImageFormat::PNG,
+            stream.clone(),
+            move |res| {
+              res?;
+              let mut bytes = Vec::new();
+              stream.seek(SeekFrom::Start(0)).unwrap();
+              match stream.read_to_end(&mut bytes) {
+                Ok(_) => handler(Ok(bytes)),
+                Err(err) => handler(Err(Error::Io(err))),
+              }
+              Ok(())
+            },
+          )?;
+        }
+        _ => todo!("{:?} screenshots for WebView2", region),
+      }
+    }
+    Ok(())
+  }
+
   pub fn eval(&self, js: &str) -> Result<()> {
     if let Some(w) = self.webview.get() {
       w.execute_script(js, |_| (Ok(())))?;

src/webview/webview2/winrt/mod.rs

@@ -243,6 +243,13 @@ impl InnerWebView {
     let _ = self.eval("window.print()");
   }
 
+  pub fn screenshot<F>(&self, region: ScreenshotRegion, handler: F) -> Result<()>
+  where
+    F: Fn(Result<Vec<u8>>) -> () + 'static + Send,
+  {
+    todo!();
+  }
+
   pub fn eval(&self, js: &str) -> Result<()> {
     if let Some(w) = self.webview.get() {
       let _ = w.ExecuteScriptAsync(js)?;

src/webview/wkwebview/mod.rs

@@ -11,11 +11,12 @@ use std::{
   slice, str,
 };
 
+use block::{Block, ConcreteBlock};
 use cocoa::base::id;
 #[cfg(target_os = "macos")]
 use cocoa::{
   appkit::{NSView, NSViewHeightSizable, NSViewWidthSizable},
-  base::YES,
+  base::{nil, YES},
 };
 
 use core_graphics::geometry::{CGPoint, CGRect, CGSize};
@@ -36,7 +37,7 @@ use crate::application::platform::ios::WindowExtIOS;
 
 use crate::{
   application::window::Window,
-  webview::{mimetype::MimeType, FileDropEvent, RpcRequest, RpcResponse},
+  webview::{mimetype::MimeType, FileDropEvent, RpcRequest, RpcResponse, ScreenshotRegion},
   Result,
 };
 
@@ -399,6 +400,33 @@ impl InnerWebView {
       let () = msg_send![print_operation, runOperationModalForWindow: self.ns_window delegate: null::<*const c_void>() didRunSelector: null::<*const c_void>() contextInfo: null::<*const c_void>()];
     }
   }
+
+  pub fn screenshot<F>(&self, region: ScreenshotRegion, handler: F) -> Result<()>
+  where
+    F: Fn(Result<Vec<u8>>) -> () + 'static + Send,
+  {
+    unsafe {
+      let config: id = msg_send![class!(WKSnapshotConfiguration), new];
+      let handler = ConcreteBlock::new(move |image: id, _error: id| {
+        let cgref: id =
+          msg_send![image, CGImageForProposedRect:null::<*const c_void>() context:nil hints:nil];
+        let bitmap_image_ref: id = msg_send![class!(NSBitmapImageRep), alloc];
+        let newrep: id = msg_send![bitmap_image_ref, initWithCGImage: cgref];
+        let size: id = msg_send![image, size];
+        let () = msg_send![newrep, setSize: size];
+        let nsdata: id = msg_send![newrep, representationUsingType:4 properties:nil];
+        let bytes: *const u8 = msg_send![nsdata, bytes];
+        let len: usize = msg_send![nsdata, length];
+        let vector = slice::from_raw_parts(bytes, len).to_vec();
+        handler(Ok(vector));
+      });
+      let handler = handler.copy();
+      let handler: &Block<(id, id), ()> = &handler;
+      let () =
+        msg_send![self.webview, takeSnapshotWithConfiguration:config completionHandler:handler];
+    }
+    Ok(())
+  }
 }
 
 pub fn platform_webview_version() -> Result<String> {