From 65e90efde03c13e170cd0d25f26b2aec3a3ee212 Mon Sep 17 00:00:00 2001
From: Jackson Coxson <jkcoxson@gmail.com>
Date: Thu, 6 Feb 2025 12:23:33 -0700
Subject: [PATCH] Add progress callback for upload image

---
 idevice/Cargo.toml     |  2 +-
 idevice/src/lib.rs     | 21 +++++++++++++--------
 idevice/src/mounter.rs | 17 +++++++++++++++--
 3 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/idevice/Cargo.toml b/idevice/Cargo.toml
index c664bcd..c9e2537 100644
--- a/idevice/Cargo.toml
+++ b/idevice/Cargo.toml
@@ -44,7 +44,7 @@ usbmuxd = []
 tcp = ["tokio/net"]
 tss = ["dep:uuid", "dep:reqwest"]
 xpc = [
-  "tokio/full",
+  "tokio/sync",
   "dep:indexmap",
   "dep:uuid",
   "dep:async-recursion",
diff --git a/idevice/src/lib.rs b/idevice/src/lib.rs
index 644add5..6d37459 100644
--- a/idevice/src/lib.rs
+++ b/idevice/src/lib.rs
@@ -89,20 +89,25 @@ impl Idevice {
 
     /// Sends raw bytes to the socket
     async fn send_raw(&mut self, message: &[u8]) -> Result<(), IdeviceError> {
+        self.send_raw_with_progress(message, |_| async {}).await
+    }
+
+    async fn send_raw_with_progress<Fut>(
+        &mut self,
+        message: &[u8],
+        callback: impl Fn((usize, usize)) -> Fut,
+    ) -> Result<(), IdeviceError>
+    where
+        Fut: std::future::Future<Output = ()>,
+    {
         if let Some(socket) = &mut self.socket {
             let message_parts = message.chunks(1024 * 64);
             let part_len = message_parts.len();
 
-            let mut err = 5;
             for (i, part) in message_parts.enumerate() {
                 debug!("Writing {i}/{part_len}");
-                while let Err(e) = socket.write_all(part).await {
-                    err -= 1;
-                    if err == 0 {
-                        return Err(e.into());
-                    }
-                }
-                err = 5;
+                socket.write_all(part).await?;
+                callback((i, part_len)).await;
             }
             Ok(())
         } else {
diff --git a/idevice/src/mounter.rs b/idevice/src/mounter.rs
index 87bb0f0..4ddc90b 100644
--- a/idevice/src/mounter.rs
+++ b/idevice/src/mounter.rs
@@ -80,6 +80,20 @@ impl ImageMounter {
         image: &[u8],
         signature: Vec<u8>,
     ) -> Result<(), IdeviceError> {
+        self.upload_image_with_progress(image_type, image, signature, |_| async {})
+            .await
+    }
+
+    pub async fn upload_image_with_progress<Fut>(
+        &mut self,
+        image_type: impl Into<String>,
+        image: &[u8],
+        signature: Vec<u8>,
+        callback: impl Fn((usize, usize)) -> Fut,
+    ) -> Result<(), IdeviceError>
+    where
+        Fut: std::future::Future<Output = ()>,
+    {
         let image_type = image_type.into();
         let image_size = match u64::try_from(image.len()) {
             Ok(i) => i,
@@ -110,8 +124,7 @@ impl ImageMounter {
         }
 
         debug!("Sending image bytes");
-        self.idevice.send_raw(image).await?;
-        debug!("Image bytes written");
+        self.idevice.send_raw_with_progress(image, callback).await?;
 
         let res = self.idevice.read_plist().await?;
         match res.get("Status") {