Skip to content

Commit

Permalink
Add initial impl of org.gnome.Shell.Screenshot
Browse files Browse the repository at this point in the history
Enough to make the portal all-outputs screenshot work. With this,
Flameshot kinda-works.
  • Loading branch information
YaLTeR committed Sep 20, 2023
1 parent 8b4f2cf commit 1963aaa
Show file tree
Hide file tree
Showing 6 changed files with 198 additions and 30 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ tracy-client = { version = "0.16.2", default-features = false }
xcursor = "0.3.4"
zbus = { version = "3.14.1" }
async-io = "1.13.0"
async-channel = "1.9.0"

[dependencies.smithay]
git = "https://github.com/Smithay/smithay.git"
Expand Down
4 changes: 2 additions & 2 deletions src/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ pub struct DebugConfig {
#[knuffel(child, unwrap(argument), default = 1.)]
pub animation_slowdown: f64,
#[knuffel(child)]
pub screen_cast_in_non_session_instances: bool,
pub dbus_interfaces_in_non_session_instances: bool,
#[knuffel(child)]
pub wait_for_frame_completion_before_queueing: bool,
#[knuffel(child)]
Expand All @@ -149,7 +149,7 @@ impl Default for DebugConfig {
fn default() -> Self {
Self {
animation_slowdown: 1.,
screen_cast_in_non_session_instances: false,
dbus_interfaces_in_non_session_instances: false,
wait_for_frame_completion_before_queueing: false,
enable_color_transformations_capability: false,
enable_overlay_planes: false,
Expand Down
57 changes: 57 additions & 0 deletions src/dbus/gnome_shell_screenshot.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
use std::path::PathBuf;

use smithay::reexports::calloop;
use zbus::{dbus_interface, fdo};

pub struct Screenshot {
to_niri: calloop::channel::Sender<ScreenshotToNiri>,
from_niri: async_channel::Receiver<NiriToScreenshot>,
}

pub enum ScreenshotToNiri {
TakeScreenshot { include_cursor: bool },
}

pub enum NiriToScreenshot {
ScreenshotResult(Option<PathBuf>),
}

#[dbus_interface(name = "org.gnome.Shell.Screenshot")]
impl Screenshot {
async fn screenshot(
&self,
include_cursor: bool,
_flash: bool,
_filename: PathBuf,
) -> fdo::Result<(bool, PathBuf)> {
if let Err(err) = self
.to_niri
.send(ScreenshotToNiri::TakeScreenshot { include_cursor })
{
warn!("error sending message to niri: {err:?}");
return Err(fdo::Error::Failed("internal error".to_owned()));
}

let filename = match self.from_niri.recv().await {
Ok(NiriToScreenshot::ScreenshotResult(Some(filename))) => filename,
Ok(NiriToScreenshot::ScreenshotResult(None)) => {
return Err(fdo::Error::Failed("internal error".to_owned()));
}
Err(err) => {
warn!("error receiving message from niri: {err:?}");
return Err(fdo::Error::Failed("internal error".to_owned()));
}
};

Ok((true, filename))
}
}

impl Screenshot {
pub fn new(
to_niri: calloop::channel::Sender<ScreenshotToNiri>,
from_niri: async_channel::Receiver<NiriToScreenshot>,
) -> Self {
Self { to_niri, from_niri }
}
}
1 change: 1 addition & 0 deletions src/dbus/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
pub mod gnome_shell_screenshot;
pub mod mutter_display_config;
pub mod mutter_screen_cast;
pub mod mutter_service_channel;
164 changes: 136 additions & 28 deletions src/niri.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
use std::cell::RefCell;
use std::cmp::max;
use std::collections::HashMap;
use std::os::unix::io::AsRawFd;
use std::path::PathBuf;
use std::process::Command;
use std::rc::Rc;
use std::sync::{Arc, Mutex};
Expand All @@ -15,6 +17,7 @@ use smithay::backend::renderer::element::surface::{
render_elements_from_surface_tree, WaylandSurfaceRenderElement,
};
use smithay::backend::renderer::element::texture::{TextureBuffer, TextureRenderElement};
use smithay::backend::renderer::element::utils::{Relocate, RelocateRenderElement};
use smithay::backend::renderer::element::{
render_elements, AsRenderElements, Element, Kind, RenderElement, RenderElementStates,
};
Expand Down Expand Up @@ -59,6 +62,7 @@ use zbus::fdo::RequestNameFlags;

use crate::backend::{Backend, Tty, Winit};
use crate::config::Config;
use crate::dbus::gnome_shell_screenshot::{self, NiriToScreenshot, ScreenshotToNiri};
use crate::dbus::mutter_display_config::DisplayConfig;
use crate::dbus::mutter_screen_cast::{self, ScreenCast, ToNiriMsg};
use crate::dbus::mutter_service_channel::ServiceChannel;
Expand Down Expand Up @@ -351,6 +355,43 @@ impl Niri {
.unwrap();
let screen_cast = ScreenCast::new(backend.connectors(), to_niri);

let (to_niri, from_screenshot) = calloop::channel::channel();
let (to_screenshot, from_niri) = async_channel::unbounded();
event_loop
.insert_source(from_screenshot, move |event, _, data| match event {
calloop::channel::Event::Msg(ScreenshotToNiri::TakeScreenshot {
include_cursor,
}) => {
let renderer = data.state.backend.renderer();
let on_done = {
let to_screenshot = to_screenshot.clone();
move |path| {
let msg = NiriToScreenshot::ScreenshotResult(Some(path));
if let Err(err) = to_screenshot.send_blocking(msg) {
warn!("error sending path to screenshot: {err:?}");
}
}
};

let res =
data.state
.niri
.screenshot_all_outputs(renderer, include_cursor, on_done);

if let Err(err) = res {
warn!("error taking a screenshot: {err:?}");

let msg = NiriToScreenshot::ScreenshotResult(None);
if let Err(err) = to_screenshot.send_blocking(msg) {
warn!("error sending None to screenshot: {err:?}");
}
}
}
calloop::channel::Event::Closed => (),
})
.unwrap();
let screenshot = gnome_shell_screenshot::Screenshot::new(to_niri, from_niri);

let mut zbus_conn = None;
let mut inhibit_power_key_fd = None;
if std::env::var_os("NOTIFY_SOCKET").is_some() {
Expand Down Expand Up @@ -395,25 +436,33 @@ impl Niri {
.build()
.unwrap();

if pipewire.is_some() {
{
let server = conn.object_server();
server
.at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
.unwrap();
server
.at(
"/org/gnome/Mutter/DisplayConfig",
DisplayConfig::new(backend.connectors()),
)
.unwrap();

let flags = RequestNameFlags::AllowReplacement
| RequestNameFlags::ReplaceExisting
| RequestNameFlags::DoNotQueue;
conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)

server
.at("/org/gnome/Shell/Screenshot", screenshot)
.unwrap();
conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
conn.request_name_with_flags("org.gnome.Shell.Screenshot", flags)
.unwrap();

if pipewire.is_some() {
server
.at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
.unwrap();
server
.at(
"/org/gnome/Mutter/DisplayConfig",
DisplayConfig::new(backend.connectors()),
)
.unwrap();
conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
.unwrap();
conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
.unwrap();
}
}

zbus_conn = Some(conn);
Expand Down Expand Up @@ -448,28 +497,37 @@ impl Niri {
warn!("error inhibiting power key: {err:?}");
}
}
} else if pipewire.is_some() && config_.debug.screen_cast_in_non_session_instances {
} else if config_.debug.dbus_interfaces_in_non_session_instances {
let conn = zbus::blocking::Connection::session().unwrap();
let flags = RequestNameFlags::AllowReplacement
| RequestNameFlags::ReplaceExisting
| RequestNameFlags::DoNotQueue;

{
let server = conn.object_server();

server
.at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
.at("/org/gnome/Shell/Screenshot", screenshot)
.unwrap();
server
.at(
"/org/gnome/Mutter/DisplayConfig",
DisplayConfig::new(backend.connectors()),
)
conn.request_name_with_flags("org.gnome.Shell.Screenshot", flags)
.unwrap();
}

let flags = RequestNameFlags::AllowReplacement
| RequestNameFlags::ReplaceExisting
| RequestNameFlags::DoNotQueue;
conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
.unwrap();
conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
.unwrap();
if pipewire.is_some() {
server
.at("/org/gnome/Mutter/ScreenCast", screen_cast.clone())
.unwrap();
server
.at(
"/org/gnome/Mutter/DisplayConfig",
DisplayConfig::new(backend.connectors()),
)
.unwrap();
conn.request_name_with_flags("org.gnome.Mutter.ScreenCast", flags)
.unwrap();
conn.request_name_with_flags("org.gnome.Mutter.DisplayConfig", flags)
.unwrap();
}
}

zbus_conn = Some(conn);
}
Expand Down Expand Up @@ -1083,6 +1141,56 @@ impl Niri {

Ok(())
}

pub fn screenshot_all_outputs(
&mut self,
renderer: &mut GlesRenderer,
include_pointer: bool,
on_done: impl FnOnce(PathBuf) + Send + 'static,
) -> anyhow::Result<()> {
let _span = tracy_client::span!("Niri::screenshot_all_outputs");

let mut elements = vec![];
let mut size = Size::from((0, 0));

let outputs: Vec<_> = self.global_space.outputs().cloned().collect();
for output in outputs {
let geom = self.global_space.output_geometry(&output).unwrap();
let geom = geom.to_physical(1);

size.w = max(size.w, geom.loc.x + geom.size.w);
size.h = max(size.h, geom.loc.y + geom.size.h);

let output_elements = self.render(renderer, &output, include_pointer);
elements.extend(output_elements.into_iter().map(|elem| {
RelocateRenderElement::from_element(elem, geom.loc, Relocate::Relative)
}));
}

let pixels = render_to_vec(renderer, size, &elements)?;

let path = make_screenshot_path().context("error making screenshot path")?;
debug!("saving screenshot to {path:?}");

thread::spawn(move || {
let res = image::save_buffer(
&path,
&pixels,
size.w as u32,
size.h as u32,
image::ColorType::Rgba8,
);

if let Err(err) = res {
warn!("error saving screenshot image: {err:?}");
return;
}

on_done(path);
});

Ok(())
}
}

render_elements! {
Expand Down

0 comments on commit 1963aaa

Please sign in to comment.