fix(screenshot): keep cached image when saving fails (#16120)

Signed-off-by: cuishuang <imcusg@gmail.com>

src/client/screenshot.rs

@@ -59,11 +59,14 @@ impl Screenshot {
     }
 
     fn handle_screenshot(&mut self, action: String) -> String {
-        let Some(data) = self.data.take() else {
+        let Some(data) = self.data.as_ref().cloned() else {
             return "No cached screenshot".to_owned();
         };
         match Self::handle_screenshot_(data, action) {
-            Ok(()) => "".to_owned(),
+            Ok(()) => {
+                self.data = None;
+                "".to_owned()
+            }
             Err(e) => e.to_string(),
         }
     }
@@ -98,3 +101,37 @@ pub fn set_screenshot(data: bytes::Bytes) {
 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();
+    }
+}