297 lines
10 KiB
Rust
297 lines
10 KiB
Rust
use std::fs;
|
|
use std::io;
|
|
use std::os::unix::fs::{OpenOptionsExt, PermissionsExt};
|
|
use std::path::PathBuf;
|
|
use std::sync::mpsc::{Receiver, SyncSender, TrySendError, sync_channel};
|
|
use std::thread;
|
|
use std::time::{Duration, SystemTime};
|
|
|
|
use smithay::reexports::calloop::channel;
|
|
|
|
use super::screenshot::{Canvas, OutputFrame, compose};
|
|
|
|
pub(crate) const QUEUE_CAPACITY: usize = 4;
|
|
/// A client normally opens and unlinks its screenshot immediately. If it dies
|
|
/// first, the file is left behind, so staged files are swept well after any
|
|
/// realistic request could still be reading one.
|
|
const STAGE_TTL: Duration = Duration::from_secs(60 * 60);
|
|
pub(crate) const SWEEP_INTERVAL: Duration = Duration::from_secs(5 * 60);
|
|
const DIR_MODE: u32 = 0o700;
|
|
const FILE_MODE: u32 = 0o600;
|
|
|
|
pub(crate) struct Job {
|
|
pub(crate) request: u64,
|
|
pub(crate) frames: Vec<OutputFrame>,
|
|
}
|
|
|
|
pub(crate) struct Encoded {
|
|
pub(crate) request: u64,
|
|
pub(crate) result: Result<PathBuf, String>,
|
|
}
|
|
|
|
pub(crate) struct Worker {
|
|
queue: SyncSender<Job>,
|
|
}
|
|
|
|
impl Worker {
|
|
pub(crate) fn spawn(results: channel::SyncSender<Encoded>) -> Option<Self> {
|
|
let (queue, receiver) = sync_channel::<Job>(QUEUE_CAPACITY);
|
|
thread::Builder::new()
|
|
.name("ferese-screenshot".to_owned())
|
|
.spawn(move || run(receiver, results))
|
|
.ok()
|
|
.map(|_| Self { queue })
|
|
}
|
|
|
|
// The compositor thread must never block, so a full queue is reported back
|
|
// to the caller instead of waiting.
|
|
pub(crate) fn submit(&self, job: Job) -> Result<(), Job> {
|
|
match self.queue.try_send(job) {
|
|
Ok(()) => Ok(()),
|
|
Err(TrySendError::Full(job)) | Err(TrySendError::Disconnected(job)) => Err(job),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn run(receiver: Receiver<Job>, results: channel::SyncSender<Encoded>) {
|
|
while let Ok(job) = receiver.recv() {
|
|
let result = encode(&job);
|
|
// This send cannot block. Each request produces exactly one Encode, so at
|
|
// most QUEUE_CAPACITY jobs can be queued while one more is being encoded,
|
|
// and RESULT_QUEUE_CAPACITY is asserted to exceed that. Blocking here
|
|
// would stall the worker with no way for the loop to recover it.
|
|
if results
|
|
.send(Encoded {
|
|
request: job.request,
|
|
result,
|
|
})
|
|
.is_err()
|
|
{
|
|
return;
|
|
}
|
|
}
|
|
}
|
|
|
|
// Removes staged screenshots whose client never collected them. Only files
|
|
// older than the TTL are touched, so an in-flight request is never disturbed.
|
|
pub(crate) fn sweep_stale_files() {
|
|
let Ok(directory) = private_directory() else {
|
|
return;
|
|
};
|
|
let Ok(entries) = fs::read_dir(&directory) else {
|
|
return;
|
|
};
|
|
for entry in entries.flatten() {
|
|
let path = entry.path();
|
|
if path.extension().is_none_or(|extension| extension != "png") {
|
|
continue;
|
|
}
|
|
let Ok(metadata) = entry.metadata() else {
|
|
continue;
|
|
};
|
|
if !metadata.is_file() {
|
|
continue;
|
|
}
|
|
// A file dated in the future is treated as fresh rather than swept.
|
|
let stale = metadata
|
|
.modified()
|
|
.ok()
|
|
.and_then(|modified| SystemTime::now().duration_since(modified).ok())
|
|
.is_some_and(|age| age > STAGE_TTL);
|
|
if stale && fs::remove_file(&path).is_ok() {
|
|
tracing::debug!(path = %path.display(), "removed an uncollected screenshot");
|
|
}
|
|
}
|
|
}
|
|
|
|
fn encode(job: &Job) -> Result<PathBuf, String> {
|
|
let canvas = compose(&job.frames)?;
|
|
let png = encode_png(&canvas)?;
|
|
write_private_file(&png)
|
|
}
|
|
|
|
fn encode_png(canvas: &Canvas) -> Result<Vec<u8>, String> {
|
|
let (width, height) = (canvas.width, canvas.height);
|
|
// A final sanity check on the composed size before handing it to the
|
|
// encoder, which would otherwise panic on a zero dimension.
|
|
if width == 0 || height == 0 {
|
|
return Err("Screenshot is empty".into());
|
|
}
|
|
let image = image::RgbaImage::from_raw(width, height, canvas.pixels.clone())
|
|
.ok_or("Screenshot buffer does not match its dimensions")?;
|
|
let mut out = Vec::new();
|
|
let encoder = image::codecs::png::PngEncoder::new_with_quality(
|
|
std::io::Cursor::new(&mut out),
|
|
image::codecs::png::CompressionType::Fast,
|
|
image::codecs::png::FilterType::Adaptive,
|
|
);
|
|
image::ImageEncoder::write_image(
|
|
encoder,
|
|
image.as_raw(),
|
|
width,
|
|
height,
|
|
image::ExtendedColorType::Rgba8,
|
|
)
|
|
.map_err(|error| format!("Could not encode screenshot: {error}"))?;
|
|
tracing::debug!(encoded = out.len(), "encoded screenshot");
|
|
Ok(out)
|
|
}
|
|
|
|
// The client is told to open this path and unlink it immediately, so the
|
|
// directory is ours, private, and files are created with the final mode.
|
|
fn write_private_file(png: &[u8]) -> Result<PathBuf, String> {
|
|
let directory = private_directory()?;
|
|
for attempt in 0..64u32 {
|
|
let path = directory.join(format!("shot-{}-{attempt}.png", std::process::id()));
|
|
let mut options = fs::OpenOptions::new();
|
|
options.write(true).create_new(true).mode(FILE_MODE);
|
|
match options.open(&path) {
|
|
Ok(mut file) => {
|
|
use io::Write;
|
|
file.write_all(png)
|
|
.and_then(|()| file.sync_all())
|
|
.map_err(|error| format!("Could not write screenshot: {error}"))?;
|
|
return Ok(path);
|
|
}
|
|
Err(error) if error.kind() == io::ErrorKind::AlreadyExists => continue,
|
|
Err(error) => return Err(format!("Could not write screenshot: {error}")),
|
|
}
|
|
}
|
|
Err("Could not find an unused screenshot filename".into())
|
|
}
|
|
|
|
fn private_directory() -> Result<PathBuf, String> {
|
|
let runtime = std::env::var_os("XDG_RUNTIME_DIR")
|
|
.ok_or("XDG_RUNTIME_DIR is not set, so screenshots cannot be staged")?;
|
|
let root = PathBuf::from(runtime);
|
|
let directory = root.join("ferese-screenshots");
|
|
match fs::symlink_metadata(&directory) {
|
|
Ok(metadata) if metadata.file_type().is_symlink() => {
|
|
return Err("Refusing to stage screenshots in a symlinked directory".into());
|
|
}
|
|
Ok(metadata) => {
|
|
use std::os::unix::fs::MetadataExt;
|
|
if !metadata.is_dir() || metadata.uid() != unsafe { libc::geteuid() } {
|
|
return Err("Refusing to stage screenshots in an untrusted directory".into());
|
|
}
|
|
}
|
|
Err(error) if error.kind() == io::ErrorKind::NotFound => {
|
|
fs::create_dir(&directory)
|
|
.and_then(|()| {
|
|
fs::set_permissions(&directory, fs::Permissions::from_mode(DIR_MODE))
|
|
})
|
|
.map_err(|error| format!("Could not create screenshot directory: {error}"))?;
|
|
}
|
|
Err(error) => return Err(format!("Could not inspect screenshot directory: {error}")),
|
|
}
|
|
Ok(directory)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use smithay::utils::Transform;
|
|
|
|
fn frame(width: i32, height: i32) -> OutputFrame {
|
|
let stride = width as usize * 4;
|
|
let mut pixels = vec![0u8; stride * height as usize];
|
|
let (quads, _) = pixels.as_chunks_mut::<4>();
|
|
for pixel in quads {
|
|
pixel[3] = 255;
|
|
}
|
|
OutputFrame {
|
|
preserve_alpha: false,
|
|
width,
|
|
height,
|
|
stride,
|
|
pixels,
|
|
transform: Transform::Normal,
|
|
location: (0, 0),
|
|
scale: 1.0,
|
|
logical_width: width,
|
|
logical_height: height,
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn encode_produces_a_decodable_png() {
|
|
let job = Job {
|
|
request: 1,
|
|
frames: vec![frame(4, 3)],
|
|
};
|
|
let canvas = compose(&job.frames).unwrap();
|
|
let png = encode_png(&canvas).expect("png");
|
|
assert_eq!(&png[..8], b"\x89PNG\r\n\x1a\n", "png signature");
|
|
|
|
let decoded = image::load_from_memory_with_format(&png, image::ImageFormat::Png)
|
|
.expect("decodes back")
|
|
.to_rgba8();
|
|
assert_eq!(decoded.dimensions(), (4, 3));
|
|
// The opaque alpha we force is preserved through the round trip.
|
|
assert_eq!(decoded.get_pixel(0, 0)[3], 255);
|
|
}
|
|
|
|
#[test]
|
|
fn staged_files_are_owner_only() {
|
|
let png = b"\x89PNG\r\n\x1a\nnot really a png".to_vec();
|
|
let path = write_private_file(&png).expect("staged");
|
|
use std::os::unix::fs::PermissionsExt;
|
|
let mode = fs::metadata(&path).unwrap().permissions().mode() & 0o777;
|
|
assert_eq!(
|
|
mode, FILE_MODE,
|
|
"staged screenshots are not group/world readable"
|
|
);
|
|
fs::remove_file(&path).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn the_sweep_keeps_fresh_files_and_ignores_other_types() {
|
|
let png = b"\x89PNG\r\n\x1a\n".to_vec();
|
|
let fresh = write_private_file(&png).expect("staged");
|
|
|
|
let directory = private_directory().unwrap();
|
|
let other = directory.join("keep-me.txt");
|
|
fs::write(&other, b"not a screenshot").unwrap();
|
|
|
|
sweep_stale_files();
|
|
assert!(
|
|
fresh.exists(),
|
|
"a freshly staged screenshot is still in use"
|
|
);
|
|
assert!(other.exists(), "unrelated files are left alone");
|
|
|
|
fs::remove_file(&fresh).unwrap();
|
|
fs::remove_file(&other).unwrap();
|
|
}
|
|
|
|
#[test]
|
|
fn the_sweep_removes_files_past_the_ttl() {
|
|
let png = b"\x89PNG\r\n\x1a\n".to_vec();
|
|
let path = write_private_file(&png).expect("staged");
|
|
// Backdate it past the TTL without waiting an hour.
|
|
let aged = SystemTime::now() - STAGE_TTL - Duration::from_secs(60);
|
|
fs::File::options()
|
|
.write(true)
|
|
.open(&path)
|
|
.unwrap()
|
|
.set_modified(aged)
|
|
.expect("backdate");
|
|
|
|
sweep_stale_files();
|
|
assert!(!path.exists(), "an uncollected screenshot is reclaimed");
|
|
}
|
|
|
|
#[test]
|
|
fn a_missing_runtime_dir_is_reported_rather_than_ignored() {
|
|
// private_directory must never fall back to /tmp or the cwd.
|
|
let result = private_directory();
|
|
match result {
|
|
Ok(path) => assert!(
|
|
path.ends_with("ferese-screenshots"),
|
|
"screenshots are staged in the runtime dir, not a fallback location"
|
|
),
|
|
Err(error) => assert!(error.contains("XDG_RUNTIME_DIR") || error.contains("untrusted")),
|
|
}
|
|
}
|
|
}
|