uboot_shell/
lib.rs

1#[macro_use]
2extern crate log;
3
4use std::{
5    fs::File,
6    io::*,
7    path::PathBuf,
8    sync::{
9        Arc,
10        atomic::{AtomicBool, Ordering},
11    },
12    thread,
13    time::{Duration, Instant},
14};
15
16mod crc;
17mod ymodem;
18
19macro_rules! dbg {
20    ($($arg:tt)*) => {{
21        debug!("$ {}", &std::fmt::format(format_args!($($arg)*)));
22    }};
23}
24
25const CTRL_C: u8 = 0x03;
26const INT_STR: &str = "<INTERRUPT>";
27const INT: &[u8] = INT_STR.as_bytes();
28
29pub struct UbootShell {
30    pub tx: Option<Box<dyn Write + Send>>,
31    pub rx: Option<Box<dyn Read + Send>>,
32    perfix: String,
33}
34
35impl UbootShell {
36    /// Create a new UbootShell instance, block wait for uboot shell.
37    pub fn new(tx: impl Write + Send + 'static, rx: impl Read + Send + 'static) -> Result<Self> {
38        let mut s = Self {
39            tx: Some(Box::new(tx)),
40            rx: Some(Box::new(rx)),
41            perfix: "".to_string(),
42        };
43        s.wait_for_shell()?;
44        debug!("shell ready, perfix: `{}`", s.perfix);
45        Ok(s)
46    }
47
48    fn rx(&mut self) -> &mut Box<dyn Read + Send> {
49        self.rx.as_mut().unwrap()
50    }
51
52    fn tx(&mut self) -> &mut Box<dyn Write + Send> {
53        self.tx.as_mut().unwrap()
54    }
55
56    fn wait_for_interrupt(&mut self) -> Result<Vec<u8>> {
57        let mut tx = self.tx.take().unwrap();
58
59        let ok = Arc::new(AtomicBool::new(false));
60
61        let tx_handle = thread::spawn({
62            let ok = ok.clone();
63            move || {
64                while !ok.load(Ordering::Acquire) {
65                    let _ = tx.write_all(&[CTRL_C]);
66                    thread::sleep(Duration::from_millis(20));
67                }
68                tx
69            }
70        });
71        let mut history: Vec<u8> = Vec::new();
72        let mut interrupt_line: Vec<u8> = Vec::new();
73        debug!("wait for interrupt");
74        loop {
75            match self.read_byte() {
76                Ok(ch) => {
77                    history.push(ch);
78
79                    if history.last() == Some(&b'\n') {
80                        let line = history.trim_ascii_end();
81                        dbg!("{}", String::from_utf8_lossy(line));
82                        let it = line.ends_with(INT);
83                        if it {
84                            interrupt_line.extend_from_slice(line);
85                        }
86                        history.clear();
87                        if it {
88                            ok.store(true, Ordering::Release);
89                            break;
90                        }
91                    }
92                }
93
94                Err(ref e) if e.kind() == ErrorKind::TimedOut => {
95                    continue;
96                }
97                Err(e) => {
98                    return Err(e);
99                }
100            }
101        }
102
103        self.tx = Some(tx_handle.join().unwrap());
104
105        Ok(interrupt_line)
106    }
107
108    fn clear_shell(&mut self) -> Result<()> {
109        let _ = self.read_to_end(&mut vec![]);
110        Ok(())
111    }
112
113    fn wait_for_shell(&mut self) -> Result<()> {
114        let mut line = self.wait_for_interrupt()?;
115        debug!("got {}", String::from_utf8_lossy(&line));
116        line.resize(line.len() - INT.len(), 0);
117        self.perfix = String::from_utf8_lossy(&line).to_string();
118        self.clear_shell()?;
119        Ok(())
120    }
121
122    fn read_byte(&mut self) -> Result<u8> {
123        let mut buff = [0u8; 1];
124        let time_out = Duration::from_secs(5);
125        let start = Instant::now();
126
127        loop {
128            match self.rx().read_exact(&mut buff) {
129                Ok(_) => return Ok(buff[0]),
130                Err(e) => {
131                    if e.kind() == ErrorKind::TimedOut {
132                        if start.elapsed() > time_out {
133                            return Err(std::io::Error::new(
134                                std::io::ErrorKind::TimedOut,
135                                "Timeout",
136                            ));
137                        }
138                    } else {
139                        return Err(e);
140                    }
141                }
142            }
143        }
144    }
145
146    pub fn wait_for_reply(&mut self, val: &str) -> Result<String> {
147        let mut reply = Vec::new();
148        let mut display = Vec::new();
149        debug!("wait for `{}`", val);
150        loop {
151            let byte = self.read_byte()?;
152            reply.push(byte);
153            display.push(byte);
154            if byte == b'\n' {
155                dbg!("{}", String::from_utf8_lossy(&display).trim_end());
156                display.clear();
157            }
158
159            if reply.ends_with(val.as_bytes()) {
160                dbg!("{}", String::from_utf8_lossy(&display).trim_end());
161                break;
162            }
163        }
164        Ok(String::from_utf8_lossy(&reply)
165            .trim()
166            .trim_end_matches(&self.perfix)
167            .to_string())
168    }
169
170    pub fn cmd_without_reply(&mut self, cmd: &str) -> Result<()> {
171        self.tx().write_all(cmd.as_bytes())?;
172        self.tx().write_all("\n".as_bytes())?;
173        // self.tx().flush()?;
174        // self.wait_for_reply(cmd)?;
175        // debug!("cmd ok");
176        Ok(())
177    }
178
179    fn _cmd(&mut self, cmd: &str) -> Result<String> {
180        let _ = self.read_to_end(&mut vec![]);
181        let ok_str = "cmd-ok";
182        let cmd_with_id = format!("{cmd}&& echo {ok_str}");
183        self.cmd_without_reply(&cmd_with_id)?;
184        let perfix = self.perfix.clone();
185        let res = self
186            .wait_for_reply(&perfix)?
187            .trim_end()
188            .trim_end_matches(self.perfix.as_str().trim())
189            .trim_end()
190            .to_string();
191        if res.ends_with(ok_str) {
192            let res = res
193                .trim()
194                .trim_end_matches(ok_str)
195                .trim_end()
196                .trim_start_matches(&cmd_with_id)
197                .trim()
198                .to_string();
199            Ok(res)
200        } else {
201            Err(Error::other(format!(
202                "command `{cmd}` failed, response: {res}",
203            )))
204        }
205    }
206
207    pub fn cmd(&mut self, cmd: &str) -> Result<String> {
208        info!("cmd: {cmd}");
209        let mut retry = 3;
210        while retry > 0 {
211            match self._cmd(cmd) {
212                Ok(res) => return Ok(res),
213                Err(e) => {
214                    warn!("cmd `{}` failed: {}, retrying...", cmd, e);
215                    retry -= 1;
216                    thread::sleep(Duration::from_millis(100));
217                }
218            }
219        }
220        Err(Error::other(format!(
221            "command `{cmd}` failed after retries",
222        )))
223    }
224
225    pub fn set_env(&mut self, name: impl Into<String>, value: impl Into<String>) -> Result<()> {
226        self.cmd(&format!("setenv {} {}", name.into(), value.into()))?;
227        Ok(())
228    }
229
230    pub fn env(&mut self, name: impl Into<String>) -> Result<String> {
231        let name = name.into();
232        let s = self.cmd(&format!("echo ${}", name))?;
233        let sp = s
234            .split("\n")
235            .filter(|s| !s.trim().is_empty())
236            .collect::<Vec<_>>();
237        let s = sp
238            .last()
239            .ok_or(Error::new(
240                ErrorKind::NotFound,
241                format!("env {} not found", name),
242            ))?
243            .to_string();
244        Ok(s)
245    }
246
247    pub fn env_int(&mut self, name: impl Into<String>) -> Result<usize> {
248        let name = name.into();
249        let line = self.env(&name)?;
250        debug!("env {name} = {line}");
251
252        parse_int(&line).ok_or(Error::new(
253            ErrorKind::InvalidData,
254            format!("env {name} is not a number"),
255        ))
256    }
257
258    pub fn loady(
259        &mut self,
260        addr: usize,
261        file: impl Into<PathBuf>,
262        on_progress: impl Fn(usize, usize),
263    ) -> Result<String> {
264        self.cmd_without_reply(&format!("loady {:#x}", addr,))?;
265        let crc = self.wait_for_load_crc()?;
266        let mut p = ymodem::Ymodem::new(crc);
267
268        let file = file.into();
269        let name = file.file_name().unwrap().to_str().unwrap();
270
271        let mut file = File::open(&file).unwrap();
272
273        let size = file.metadata().unwrap().len() as usize;
274
275        p.send(self, &mut file, name, size, |p| {
276            on_progress(p, size);
277        })?;
278        let perfix = self.perfix.clone();
279        self.wait_for_reply(&perfix)
280    }
281
282    fn wait_for_load_crc(&mut self) -> Result<bool> {
283        let mut reply = Vec::new();
284        loop {
285            let byte = self.read_byte()?;
286            reply.push(byte);
287            print_raw(&[byte]);
288
289            if reply.ends_with(b"C") {
290                return Ok(true);
291            }
292            let res = String::from_utf8_lossy(&reply);
293            if res.contains("try 'help'") {
294                panic!("{}", res);
295            }
296        }
297    }
298}
299
300impl Read for UbootShell {
301    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
302        self.rx().read(buf)
303    }
304}
305
306impl Write for UbootShell {
307    fn write(&mut self, buf: &[u8]) -> Result<usize> {
308        self.tx().write(buf)
309    }
310
311    fn flush(&mut self) -> Result<()> {
312        self.tx().flush()
313    }
314}
315
316fn parse_int(line: &str) -> Option<usize> {
317    let mut line = line.trim();
318    let mut radix = 10;
319    if line.starts_with("0x") {
320        line = &line[2..];
321        radix = 16;
322    }
323    u64::from_str_radix(line, radix).ok().map(|o| o as _)
324}
325
326fn print_raw(buff: &[u8]) {
327    #[cfg(target_os = "windows")]
328    print_raw_win(buff);
329    #[cfg(not(target_os = "windows"))]
330    stdout().write_all(buff).unwrap();
331}
332
333#[cfg(target_os = "windows")]
334fn print_raw_win(buff: &[u8]) {
335    use std::sync::Mutex;
336    static PRINT_BUFF: Mutex<Vec<u8>> = Mutex::new(Vec::new());
337
338    let mut g = PRINT_BUFF.lock().unwrap();
339
340    g.extend_from_slice(buff);
341
342    if g.ends_with(b"\n") {
343        let s = String::from_utf8_lossy(&g[..]);
344        println!("{}", s.trim());
345        g.clear();
346    }
347}
uboot_shell/lib.rs

uboot_shell/
lib.rs
