unsquashfs_wrapper/
lib.rs

1use std::{
2    io::{self, BufReader, Error, ErrorKind, Read},
3    path::Path,
4    process::{ChildStdout, Command, Stdio},
5    str,
6    sync::{
7        atomic::{AtomicBool, Ordering},
8        Arc, RwLock,
9    },
10    thread,
11    time::Duration,
12};
13
14use thiserror::Error;
15
16fn handle(stdout: ChildStdout, mut callback: impl FnMut(i32)) -> io::Result<()> {
17    let mut last_progress = 0;
18    let mut reader = BufReader::new(stdout);
19
20    loop {
21        let mut data = [0; 0x1000];
22        let count = reader.read(&mut data)?;
23
24        if count == 0 {
25            return Ok(());
26        }
27
28        if let Ok(string) = str::from_utf8(&data[..count]) {
29            for line in string.split(['\r', '\n']) {
30                let len = line.len();
31                if line.starts_with('[') && line.ends_with('%') && len >= 4 {
32                    if let Ok(progress) = line[len - 4..len - 1].trim().parse::<i32>() {
33                        if last_progress != progress {
34                            callback(progress);
35                            last_progress = progress;
36                        }
37                    }
38                }
39            }
40        }
41    }
42}
43
44#[derive(Clone)]
45pub struct Unsquashfs {
46    cancel: Arc<AtomicBool>,
47    status: Arc<RwLock<Status>>,
48}
49
50pub enum Status {
51    Pending,
52    Working,
53}
54
55impl Default for Unsquashfs {
56    fn default() -> Self {
57        Self {
58            cancel: Arc::new(AtomicBool::new(false)),
59            status: Arc::new(RwLock::new(Status::Pending)),
60        }
61    }
62}
63
64#[derive(Debug, Error)]
65pub enum UnsquashfsError {
66    #[error("`unsquashfs` binary does not exist.")]
67    BinaryDoesNotExist,
68    #[error(transparent)]
69    IO(#[from] io::Error),
70    #[error("`unsquashfs` is not start.")]
71    Pending,
72    #[error("`unsquashfs` failed: {0}, output: {1}")]
73    Failure(io::Error, String),
74}
75
76impl Unsquashfs {
77    pub fn new() -> Self {
78        Unsquashfs::default()
79    }
80
81    pub fn cancel(&self) -> Result<(), UnsquashfsError> {
82        match *self.status.read().unwrap() {
83            Status::Pending => Err(UnsquashfsError::Pending),
84            Status::Working => {
85                self.cancel.store(true, Ordering::SeqCst);
86                Ok(())
87            }
88        }
89    }
90
91    /// Extracts an image using either unsquashfs.
92    pub fn extract(
93        &self,
94        archive: impl AsRef<Path>,
95        directory: impl AsRef<Path>,
96        thread: Option<usize>,
97        callback: impl FnMut(i32),
98    ) -> Result<(), UnsquashfsError> {
99        if which::which("unsquashfs").is_err() {
100            return Err(UnsquashfsError::BinaryDoesNotExist);
101        }
102
103        let archive = archive.as_ref().canonicalize()?;
104        let directory = directory.as_ref().canonicalize()?;
105
106        let directory = directory
107            .to_str()
108            .ok_or_else(|| Error::new(ErrorKind::InvalidData, "Invalid directory path"))?
109            .replace('\'', "'\"'\"'");
110
111        let archive = archive
112            .to_str()
113            .ok_or_else(|| Error::new(ErrorKind::InvalidData, "Invalid archive path"))?
114            .replace('\'', "'\"'\"'");
115
116        let mut command = Command::new("unsquashfs");
117
118        if let Some(limit_thread) = thread {
119            command.arg("-p").arg(limit_thread.to_string());
120        }
121
122        command
123            .arg("-f")
124            .arg("-q")
125            .arg("-d")
126            .arg(directory)
127            .arg(archive);
128
129        let mut child = command
130            .env("COLUMNS", "")
131            .env("LINES", "")
132            .env("TERM", "xterm-256color")
133            .stdout(Stdio::piped())
134            .stderr(Stdio::piped())
135            .spawn()?;
136
137        *self.status.write().unwrap() = Status::Working;
138
139        let cc = self.cancel.clone();
140        let status_clone = self.status.clone();
141
142        let stdout = child
143            .stdout
144            .take()
145            .ok_or_else(|| io::Error::new(ErrorKind::BrokenPipe, "Failed to get stdout"))?;
146
147        let stderr = child
148            .stderr
149            .take()
150            .ok_or_else(|| io::Error::new(ErrorKind::BrokenPipe, "Failed to get stderr"))?;
151
152        let process_control = thread::spawn(move || -> io::Result<()> {
153            loop {
154                let wait = child.try_wait()?;
155
156                if cc.load(Ordering::SeqCst) {
157                    child.kill()?;
158                    cc.store(false, Ordering::SeqCst);
159                    *status_clone.write().unwrap() = Status::Pending;
160                    return Ok(());
161                }
162
163                let Some(wait) = wait else {
164                    thread::sleep(Duration::from_millis(10));
165                    continue;
166                };
167
168                *status_clone.write().unwrap() = Status::Pending;
169
170                if !wait.success() {
171                    return Err(Error::new(
172                        ErrorKind::Other,
173                        format!(
174                            "archive extraction failed with status: {}",
175                            wait.code().unwrap_or(1),
176                        ),
177                    ));
178                } else {
179                    return Ok(());
180                }
181            }
182        });
183
184        handle(stdout, callback)?;
185
186        let mut stderr = BufReader::new(stderr);
187        let mut buf = String::new();
188        stderr.read_to_string(&mut buf).ok();
189
190        process_control
191            .join()
192            .unwrap()
193            .map_err(|e| UnsquashfsError::Failure(e, buf))?;
194
195        Ok(())
196    }
197}
198
199#[cfg(test)]
200pub mod test {
201    use std::{env::temp_dir, fs, thread, time::Duration};
202
203    use crate::Unsquashfs;
204
205    #[test]
206    fn test_extract() {
207        let unsquashfs = Unsquashfs::default();
208        let unsquashfs_clone = unsquashfs.clone();
209
210        let t = thread::spawn(move || {
211            let output = temp_dir().join("unsqfs-wrap-test-extract");
212            fs::create_dir_all(&output).unwrap();
213            unsquashfs
214                .extract(
215                    "testdata/test_extract.squashfs",
216                    &output,
217                    None,
218                    Box::new(move |c| {
219                        dbg!(c);
220                    }),
221                )
222                .unwrap();
223            fs::remove_dir_all(output).unwrap();
224        });
225
226        thread::sleep(Duration::from_millis(10));
227        unsquashfs_clone.cancel().unwrap();
228
229        t.join().unwrap();
230    }
231}
unsquashfs_wrapper/lib.rs

unsquashfs_wrapper/
lib.rs
