#[cfg(not(any(target_os = "android", target_os = "ios")))] use crate::clipboard::{update_clipboard, ClipboardSide}; use base::message_proto::*; use hbb_common::ResultType; use std::sync::Mutex; lazy_static::lazy_static! { static ref SCREENSHOT: Mutex = Default::default(); } pub enum ScreenshotAction { SaveAs(String), CopyToClipboard, Discard, } impl Default for ScreenshotAction { fn default() -> Self { Self::Discard } } impl From<&str> for ScreenshotAction { fn from(value: &str) -> Self { match value.chars().next() { Some('0') => { if let Some((pos, _)) = value.char_indices().nth(2) { let substring = &value[pos..]; Self::SaveAs(substring.to_string()) } else { Self::default() } } Some('1') => Self::CopyToClipboard, Some('2') => Self::default(), _ => Self::default(), } } } impl Into for ScreenshotAction { fn into(self) -> String { match self { Self::SaveAs(p) => format!("0:{p}"), Self::CopyToClipboard => "1".to_owned(), Self::Discard => "2".to_owned(), } } } #[derive(Default)] pub struct Screenshot { data: Option, } impl Screenshot { fn set_screenshot(&mut self, data: bytes::Bytes) { self.data.replace(data); } fn handle_screenshot(&mut self, action: String) -> String { let Some(data) = self.data.as_ref().cloned() else { return "No cached screenshot".to_owned(); }; match Self::handle_screenshot_(data, action) { Ok(()) => { self.data = None; "".to_owned() } Err(e) => e.to_string(), } } fn handle_screenshot_(data: bytes::Bytes, action: String) -> ResultType<()> { match ScreenshotAction::from(&action as &str) { ScreenshotAction::SaveAs(p) => { std::fs::write(p, data)?; } ScreenshotAction::CopyToClipboard => { #[cfg(not(any(target_os = "android", target_os = "ios")))] { let clips = vec![Clipboard { compress: false, content: data, format: ClipboardFormat::ImagePng.into(), ..Default::default() }]; update_clipboard(clips, ClipboardSide::Client); } } ScreenshotAction::Discard => {} } Ok(()) } } pub fn set_screenshot(data: bytes::Bytes) { SCREENSHOT.lock().unwrap().set_screenshot(data); } pub fn handle_screenshot(action: String) -> String { SCREENSHOT.lock().unwrap().handle_screenshot(action) } #[cfg(test)] mod tests { use super::Screenshot; use std::time::{SystemTime, UNIX_EPOCH}; #[test] fn preserves_cached_screenshot_when_save_fails() { let data = bytes::Bytes::from_static(b"screenshot data"); let mut screenshot = Screenshot { data: Some(data.clone()), }; let unique = SystemTime::now() .duration_since(UNIX_EPOCH) .unwrap() .as_nanos(); let missing_parent = std::env::temp_dir() .join(format!("rustdesk-screenshot-missing-parent-{unique}")) .join("screenshot.png"); let valid_path = std::env::temp_dir().join(format!("rustdesk-screenshot-{unique}.png")); let error = screenshot.handle_screenshot(format!("0:{}", missing_parent.display())); assert!(!error.is_empty()); assert_eq!(screenshot.data.as_deref(), Some(data.as_ref())); assert_eq!( screenshot.handle_screenshot(format!("0:{}", valid_path.display())), "" ); assert!(screenshot.data.is_none()); assert_eq!(std::fs::read(&valid_path).unwrap(), data.as_ref()); std::fs::remove_file(valid_path).unwrap(); } }