Skip to main content

uboot_shell/
lib.rs

1//! # uboot-shell
2//!
3//! A Rust library for communicating with U-Boot bootloader over serial connection.
4//!
5//! This crate provides functionality to interact with U-Boot shell, execute commands,
6//! transfer files via YMODEM protocol, and manage environment variables.
7//!
8//! ## Features
9//!
10//! - Automatic U-Boot shell detection and synchronization
11//! - Command execution with retry support
12//! - YMODEM file transfer protocol implementation
13//! - Environment variable management
14//! - CRC16-CCITT checksum support
15//!
16//! ## Quick Start
17//!
18//! ```rust,no_run
19//! use uboot_shell::UbootShell;
20//! use std::io::{Read, Write};
21//!
22//! // Open serial port (using serialport crate)
23//! let port = serialport::new("/dev/ttyUSB0", 115200)
24//!     .open()
25//!     .unwrap();
26//! let rx = port.try_clone().unwrap();
27//! let tx = port;
28//!
29//! // Create U-Boot shell instance (blocks until shell is ready)
30//! let mut uboot = UbootShell::new(tx, rx).unwrap();
31//!
32//! // Execute commands
33//! let output = uboot.cmd("help").unwrap();
34//! println!("{}", output);
35//!
36//! // Get/set environment variables
37//! let bootargs = uboot.env("bootargs").unwrap();
38//! uboot.set_env("myvar", "myvalue").unwrap();
39//!
40//! // Transfer file via YMODEM
41//! uboot.loady(0x80000000, "kernel.bin", |sent, total| {
42//!     println!("Progress: {}/{}", sent, total);
43//! }).unwrap();
44//! ```
45//!
46//! ## Modules
47//!
48//! - [`crc`] - CRC16-CCITT checksum implementation
49//! - [`ymodem`] - YMODEM file transfer protocol
50
51#[macro_use]
52extern crate log;
53
54use std::{
55    fs::File,
56    io::*,
57    path::PathBuf,
58    sync::{
59        Arc,
60        atomic::{AtomicBool, Ordering},
61    },
62    thread,
63    time::{Duration, Instant},
64};
65
66/// CRC16-CCITT checksum implementation.
67pub mod crc;
68
69/// YMODEM file transfer protocol implementation.
70pub mod ymodem;
71
72macro_rules! dbg {
73    ($($arg:tt)*) => {{
74        debug!("$ {}", &std::fmt::format(format_args!($($arg)*)));
75    }};
76}
77
78const CTRL_C: u8 = 0x03;
79const INT_STR: &str = "<INTERRUPT>";
80const INT: &[u8] = INT_STR.as_bytes();
81
82/// U-Boot shell communication interface.
83///
84/// `UbootShell` provides methods to interact with U-Boot bootloader
85/// over a serial connection. It handles shell synchronization,
86/// command execution, and file transfers.
87///
88/// # Example
89///
90/// ```rust,no_run
91/// use uboot_shell::UbootShell;
92///
93/// // Assuming tx and rx are Read/Write implementations
94/// # fn example(tx: impl std::io::Write + Send + 'static, rx: impl std::io::Read + Send + 'static) {
95/// let mut shell = UbootShell::new(tx, rx).unwrap();
96/// let result = shell.cmd("printenv").unwrap();
97/// # }
98/// ```
99pub struct UbootShell {
100    /// Transmit channel for sending data to U-Boot.
101    pub tx: Option<Box<dyn Write + Send>>,
102    /// Receive channel for reading data from U-Boot.
103    pub rx: Option<Box<dyn Read + Send>>,
104    /// Shell prompt prefix detected during initialization.
105    perfix: String,
106}
107
108impl UbootShell {
109    /// Creates a new UbootShell instance and waits for U-Boot shell to be ready.
110    ///
111    /// This function will block until it successfully detects the U-Boot shell prompt.
112    /// It sends interrupt signals (Ctrl+C) to ensure the shell is in a clean state.
113    ///
114    /// # Arguments
115    ///
116    /// * `tx` - A writable stream for sending data to U-Boot
117    /// * `rx` - A readable stream for receiving data from U-Boot
118    ///
119    /// # Returns
120    ///
121    /// Returns `Ok(UbootShell)` if the shell is successfully initialized,
122    /// or an `Err` if communication fails.
123    ///
124    /// # Errors
125    ///
126    /// Returns an error if the serial I/O fails or the prompt cannot be detected
127    /// within the internal retry loop.
128    ///
129    /// # Example
130    ///
131    /// ```rust,no_run
132    /// use uboot_shell::UbootShell;
133    ///
134    /// let port = serialport::new("/dev/ttyUSB0", 115200).open().unwrap();
135    /// let rx = port.try_clone().unwrap();
136    /// let mut uboot = UbootShell::new(port, rx).unwrap();
137    /// ```
138    pub fn new(tx: impl Write + Send + 'static, rx: impl Read + Send + 'static) -> Result<Self> {
139        let mut s = Self {
140            tx: Some(Box::new(tx)),
141            rx: Some(Box::new(rx)),
142            perfix: "".to_string(),
143        };
144        s.wait_for_shell()?;
145        debug!("shell ready, perfix: `{}`", s.perfix);
146        Ok(s)
147    }
148
149    fn rx(&mut self) -> &mut Box<dyn Read + Send> {
150        self.rx.as_mut().unwrap()
151    }
152
153    fn tx(&mut self) -> &mut Box<dyn Write + Send> {
154        self.tx.as_mut().unwrap()
155    }
156
157    fn wait_for_interrupt(&mut self) -> Result<Vec<u8>> {
158        let mut tx = self.tx.take().unwrap();
159
160        let ok = Arc::new(AtomicBool::new(false));
161
162        let tx_handle = thread::spawn({
163            let ok = ok.clone();
164            move || {
165                while !ok.load(Ordering::Acquire) {
166                    let _ = tx.write_all(&[CTRL_C]);
167                    thread::sleep(Duration::from_millis(20));
168                }
169                tx
170            }
171        });
172        let mut history: Vec<u8> = Vec::new();
173        let mut interrupt_line: Vec<u8> = Vec::new();
174        debug!("wait for interrupt");
175        loop {
176            match self.read_byte() {
177                Ok(ch) => {
178                    history.push(ch);
179
180                    if history.last() == Some(&b'\n') {
181                        let line = history.trim_ascii_end();
182                        dbg!("{}", String::from_utf8_lossy(line));
183                        let it = line.ends_with(INT);
184                        if it {
185                            interrupt_line.extend_from_slice(line);
186                        }
187                        history.clear();
188                        if it {
189                            ok.store(true, Ordering::Release);
190                            break;
191                        }
192                    }
193                }
194
195                Err(ref e) if e.kind() == ErrorKind::TimedOut => {
196                    continue;
197                }
198                Err(e) => {
199                    return Err(e);
200                }
201            }
202        }
203
204        self.tx = Some(tx_handle.join().unwrap());
205
206        Ok(interrupt_line)
207    }
208
209    fn clear_shell(&mut self) -> Result<()> {
210        let _ = self.read_to_end(&mut vec![]);
211        Ok(())
212    }
213
214    fn wait_for_shell(&mut self) -> Result<()> {
215        let mut line = self.wait_for_interrupt()?;
216        debug!("got {}", String::from_utf8_lossy(&line));
217        line.resize(line.len() - INT.len(), 0);
218        self.perfix = String::from_utf8_lossy(&line).to_string();
219        self.clear_shell()?;
220        Ok(())
221    }
222
223    fn read_byte(&mut self) -> Result<u8> {
224        let mut buff = [0u8; 1];
225        let time_out = Duration::from_secs(5);
226        let start = Instant::now();
227
228        loop {
229            match self.rx().read_exact(&mut buff) {
230                Ok(_) => return Ok(buff[0]),
231                Err(e) => {
232                    if e.kind() == ErrorKind::TimedOut {
233                        if start.elapsed() > time_out {
234                            return Err(std::io::Error::new(
235                                std::io::ErrorKind::TimedOut,
236                                "Timeout",
237                            ));
238                        }
239                    } else {
240                        return Err(e);
241                    }
242                }
243            }
244        }
245    }
246
247    /// Waits for a specific string to appear in the U-Boot output.
248    ///
249    /// Reads from the serial connection until the specified string is found.
250    ///
251    /// # Arguments
252    ///
253    /// * `val` - The string to wait for
254    ///
255    /// # Returns
256    ///
257    /// Returns the accumulated output up to and including the matched string.
258    ///
259    /// # Errors
260    ///
261    /// Returns an error when the underlying read operation times out or fails.
262    pub fn wait_for_reply(&mut self, val: &str) -> Result<String> {
263        let mut reply = Vec::new();
264        let mut display = Vec::new();
265        debug!("wait for `{}`", val);
266        loop {
267            let byte = self.read_byte()?;
268            reply.push(byte);
269            display.push(byte);
270            if byte == b'\n' {
271                dbg!("{}", String::from_utf8_lossy(&display).trim_end());
272                display.clear();
273            }
274
275            if reply.ends_with(val.as_bytes()) {
276                dbg!("{}", String::from_utf8_lossy(&display).trim_end());
277                break;
278            }
279        }
280        Ok(String::from_utf8_lossy(&reply)
281            .trim()
282            .trim_end_matches(&self.perfix)
283            .to_string())
284    }
285
286    /// Sends a command to U-Boot without waiting for the response.
287    ///
288    /// This is useful for commands that don't produce output or when
289    /// you want to handle the response separately.
290    ///
291    /// # Arguments
292    ///
293    /// * `cmd` - The command string to send
294    ///
295    /// # Errors
296    ///
297    /// Returns any I/O error that occurs while writing to the serial stream.
298    pub fn cmd_without_reply(&mut self, cmd: &str) -> Result<()> {
299        self.tx().write_all(cmd.as_bytes())?;
300        self.tx().write_all("\n".as_bytes())?;
301        // self.tx().flush()?;
302        // self.wait_for_reply(cmd)?;
303        // debug!("cmd ok");
304        Ok(())
305    }
306
307    fn _cmd(&mut self, cmd: &str) -> Result<String> {
308        let _ = self.read_to_end(&mut vec![]);
309        let ok_str = "cmd-ok";
310        let cmd_with_id = format!("{cmd}&& echo {ok_str}");
311        self.cmd_without_reply(&cmd_with_id)?;
312        let perfix = self.perfix.clone();
313        let res = self
314            .wait_for_reply(&perfix)?
315            .trim_end()
316            .trim_end_matches(self.perfix.as_str().trim())
317            .trim_end()
318            .to_string();
319        if res.ends_with(ok_str) {
320            let res = res
321                .trim()
322                .trim_end_matches(ok_str)
323                .trim_end()
324                .trim_start_matches(&cmd_with_id)
325                .trim()
326                .to_string();
327            Ok(res)
328        } else {
329            Err(Error::other(format!(
330                "command `{cmd}` failed, response: {res}",
331            )))
332        }
333    }
334
335    /// Executes a command in U-Boot shell and returns the output.
336    ///
337    /// This method sends the command, waits for completion, and verifies
338    /// that the command executed successfully. It includes automatic retry
339    /// logic (up to 3 attempts) for improved reliability.
340    ///
341    /// # Arguments
342    ///
343    /// * `cmd` - The command string to execute
344    ///
345    /// # Returns
346    ///
347    /// Returns `Ok(String)` with the command output on success,
348    /// or an `Err` if the command fails after all retries.
349    ///
350    /// # Errors
351    ///
352    /// Returns an error if the command fails after retries or if serial I/O fails.
353    ///
354    /// # Example
355    ///
356    /// ```rust,no_run
357    /// # use uboot_shell::UbootShell;
358    /// # fn example(uboot: &mut UbootShell) {
359    /// let output = uboot.cmd("printenv bootargs").unwrap();
360    /// println!("bootargs: {}", output);
361    /// # }
362    /// ```
363    pub fn cmd(&mut self, cmd: &str) -> Result<String> {
364        info!("cmd: {cmd}");
365        let mut retry = 3;
366        while retry > 0 {
367            match self._cmd(cmd) {
368                Ok(res) => return Ok(res),
369                Err(e) => {
370                    warn!("cmd `{}` failed: {}, retrying...", cmd, e);
371                    retry -= 1;
372                    thread::sleep(Duration::from_millis(100));
373                }
374            }
375        }
376        Err(Error::other(format!(
377            "command `{cmd}` failed after retries",
378        )))
379    }
380
381    /// Sets a U-Boot environment variable.
382    ///
383    /// # Arguments
384    ///
385    /// * `name` - The name of the environment variable
386    /// * `value` - The value to set
387    ///
388    /// # Example
389    ///
390    /// ```rust,no_run
391    /// # use uboot_shell::UbootShell;
392    /// # fn example(uboot: &mut UbootShell) {
393    /// uboot.set_env("bootargs", "console=ttyS0,115200").unwrap();
394    /// # }
395    /// ```
396    ///
397    /// # Errors
398    ///
399    /// Returns any error from the underlying command execution.
400    pub fn set_env(&mut self, name: impl Into<String>, value: impl Into<String>) -> Result<()> {
401        self.cmd(&format!("setenv {} {}", name.into(), value.into()))?;
402        Ok(())
403    }
404
405    /// Gets the value of a U-Boot environment variable.
406    ///
407    /// # Arguments
408    ///
409    /// * `name` - The name of the environment variable
410    ///
411    /// # Returns
412    ///
413    /// Returns `Ok(String)` with the variable value, or an `Err` if not found.
414    ///
415    /// # Errors
416    ///
417    /// Returns `ErrorKind::NotFound` if the variable is not set or cannot be read.
418    ///
419    /// # Example
420    ///
421    /// ```rust,no_run
422    /// # use uboot_shell::UbootShell;
423    /// # fn example(uboot: &mut UbootShell) {
424    /// let bootargs = uboot.env("bootargs").unwrap();
425    /// # }
426    /// ```
427    ///
428    /// # Errors
429    ///
430    /// Returns `ErrorKind::NotFound` if the variable is not set or cannot be read.
431    pub fn env(&mut self, name: impl Into<String>) -> Result<String> {
432        let name = name.into();
433        let s = self.cmd(&format!("echo ${}", name))?;
434        let sp = s
435            .split("\n")
436            .filter(|s| !s.trim().is_empty())
437            .collect::<Vec<_>>();
438        let s = sp
439            .last()
440            .ok_or(Error::new(
441                ErrorKind::NotFound,
442                format!("env {} not found", name),
443            ))?
444            .to_string();
445        Ok(s)
446    }
447
448    /// Gets a U-Boot environment variable as an integer.
449    ///
450    /// Supports both decimal and hexadecimal (0x prefix) formats.
451    ///
452    /// # Arguments
453    ///
454    /// * `name` - The name of the environment variable
455    ///
456    /// # Returns
457    ///
458    /// Returns `Ok(usize)` with the parsed integer value,
459    /// or an `Err` if not found or not a valid number.
460    ///
461    /// # Errors
462    ///
463    /// Returns `ErrorKind::InvalidData` if the value is not a valid integer.
464    pub fn env_int(&mut self, name: impl Into<String>) -> Result<usize> {
465        let name = name.into();
466        let line = self.env(&name)?;
467        debug!("env {name} = {line}");
468
469        parse_int(&line).ok_or(Error::new(
470            ErrorKind::InvalidData,
471            format!("env {name} is not a number"),
472        ))
473    }
474
475    /// Transfers a file to U-Boot memory using YMODEM protocol.
476    ///
477    /// Uses the U-Boot `loady` command to receive files via YMODEM protocol.
478    /// The file will be loaded to the specified memory address.
479    ///
480    /// # Arguments
481    ///
482    /// * `addr` - The memory address where the file will be loaded
483    /// * `file` - Path to the file to transfer
484    /// * `on_progress` - Callback function called with (bytes_sent, total_bytes)
485    ///
486    /// # Returns
487    ///
488    /// Returns `Ok(String)` with the U-Boot response on success.
489    ///
490    /// # Errors
491    ///
492    /// Returns an error if the file cannot be opened, the path has a non-UTF-8
493    /// file name, or if the serial transfer fails.
494    ///
495    /// # Example
496    ///
497    /// ```rust,no_run
498    /// # use uboot_shell::UbootShell;
499    /// # fn example(uboot: &mut UbootShell) {
500    /// uboot.loady(0x80000000, "kernel.bin", |sent, total| {
501    ///     println!("Progress: {}/{} bytes", sent, total);
502    /// }).unwrap();
503    /// # }
504    /// ```
505    pub fn loady(
506        &mut self,
507        addr: usize,
508        file: impl Into<PathBuf>,
509        on_progress: impl Fn(usize, usize),
510    ) -> Result<String> {
511        self.cmd_without_reply(&format!("loady {:#x}", addr,))?;
512        let crc = self.wait_for_load_crc()?;
513        let mut p = ymodem::Ymodem::new(crc);
514
515        let file = file.into();
516        let name = file
517            .file_name()
518            .and_then(|name| name.to_str())
519            .ok_or_else(|| Error::new(ErrorKind::InvalidInput, "file name must be valid UTF-8"))?;
520
521        let mut file = File::open(&file)?;
522
523        let size = file.metadata()?.len() as usize;
524
525        p.send(self, &mut file, name, size, |p| {
526            on_progress(p, size);
527        })?;
528        let perfix = self.perfix.clone();
529        self.wait_for_reply(&perfix)
530    }
531
532    fn wait_for_load_crc(&mut self) -> Result<bool> {
533        let mut reply = Vec::new();
534        loop {
535            let byte = self.read_byte()?;
536            reply.push(byte);
537            print_raw(&[byte]);
538
539            if reply.ends_with(b"C") {
540                return Ok(true);
541            }
542            let res = String::from_utf8_lossy(&reply);
543            if res.contains("try 'help'") {
544                return Err(Error::new(
545                    ErrorKind::InvalidData,
546                    format!("U-Boot loady failed: {res}"),
547                ));
548            }
549        }
550    }
551}
552
553impl Read for UbootShell {
554    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
555        self.rx().read(buf)
556    }
557}
558
559impl Write for UbootShell {
560    fn write(&mut self, buf: &[u8]) -> Result<usize> {
561        self.tx().write(buf)
562    }
563
564    fn flush(&mut self) -> Result<()> {
565        self.tx().flush()
566    }
567}
568
569fn parse_int(line: &str) -> Option<usize> {
570    let mut line = line.trim();
571    let mut radix = 10;
572    if line.starts_with("0x") {
573        line = &line[2..];
574        radix = 16;
575    }
576    u64::from_str_radix(line, radix).ok().map(|o| o as _)
577}
578
579fn print_raw(buff: &[u8]) {
580    #[cfg(target_os = "windows")]
581    print_raw_win(buff);
582    #[cfg(not(target_os = "windows"))]
583    stdout().write_all(buff).unwrap();
584}
585
586#[cfg(target_os = "windows")]
587fn print_raw_win(buff: &[u8]) {
588    use std::sync::Mutex;
589    static PRINT_BUFF: Mutex<Vec<u8>> = Mutex::new(Vec::new());
590
591    let mut g = PRINT_BUFF.lock().unwrap();
592
593    g.extend_from_slice(buff);
594
595    if g.ends_with(b"\n") {
596        let s = String::from_utf8_lossy(&g[..]);
597        println!("{}", s.trim());
598        g.clear();
599    }
600}