Skip to main content

uboot_shell/
ymodem.rs

1//! Async YMODEM file transfer protocol implementation.
2
3use std::io::{Error, ErrorKind, Result};
4
5use futures::{
6    AsyncReadExt, AsyncWriteExt,
7    io::{AllowStdIo, AsyncRead, AsyncWrite},
8};
9
10use crate::crc::crc16_ccitt;
11
12const SOH: u8 = 0x01;
13const STX: u8 = 0x02;
14const EOT: u8 = 0x04;
15const ACK: u8 = 0x06;
16const NAK: u8 = 0x15;
17const EOF: u8 = 0x1A;
18const CRC: u8 = 0x43;
19
20pub struct Ymodem {
21    crc_mode: bool,
22    blk: u8,
23    retries: usize,
24}
25
26impl Ymodem {
27    pub fn new(crc_mode: bool) -> Self {
28        Self {
29            crc_mode,
30            blk: 0,
31            retries: 10,
32        }
33    }
34
35    fn nak(&self) -> u8 {
36        if self.crc_mode { CRC } else { NAK }
37    }
38
39    async fn getc<D: AsyncRead + Unpin>(&mut self, dev: &mut D) -> Result<u8> {
40        let mut buff = [0u8; 1];
41        dev.read_exact(&mut buff).await?;
42        Ok(buff[0])
43    }
44
45    async fn wait_for_start<D: AsyncRead + Unpin>(&mut self, dev: &mut D) -> Result<()> {
46        loop {
47            match self.getc(dev).await? {
48                NAK => {
49                    self.crc_mode = false;
50                    return Ok(());
51                }
52                CRC => {
53                    self.crc_mode = true;
54                    return Ok(());
55                }
56                _ => {}
57            }
58        }
59    }
60
61    pub async fn send<D, F>(
62        &mut self,
63        dev: &mut D,
64        file: &mut F,
65        name: &str,
66        size: usize,
67        on_progress: impl Fn(usize),
68    ) -> Result<()>
69    where
70        D: AsyncWrite + AsyncRead + Unpin,
71        F: AsyncRead + Unpin,
72    {
73        info!("Sending file: {name}");
74
75        self.send_header(dev, name, size).await?;
76
77        let mut buff = [0u8; 1024];
78        let mut send_size = 0;
79
80        loop {
81            let n = file.read(&mut buff).await?;
82            if n == 0 {
83                break;
84            }
85            self.send_blk(dev, &buff[..n], EOF, false).await?;
86            send_size += n;
87            on_progress(send_size);
88        }
89
90        dev.write_all(&[EOT]).await?;
91        dev.flush().await?;
92        self.wait_ack(dev).await?;
93
94        self.send_blk(dev, &[0], 0, true).await?;
95        self.wait_for_start(dev).await?;
96        Ok(())
97    }
98
99    async fn wait_ack<D: AsyncRead + Unpin>(&mut self, dev: &mut D) -> Result<()> {
100        let nak = self.nak();
101        loop {
102            let c = self.getc(dev).await?;
103            match c {
104                ACK => return Ok(()),
105                _ => {
106                    if c == nak {
107                        return Err(Error::new(ErrorKind::BrokenPipe, "NAK"));
108                    }
109                    let mut out = AllowStdIo::new(std::io::stdout());
110                    out.write_all(&[c]).await?;
111                }
112            }
113        }
114    }
115
116    async fn send_header<D: AsyncWrite + AsyncRead + Unpin>(
117        &mut self,
118        dev: &mut D,
119        name: &str,
120        size: usize,
121    ) -> Result<()> {
122        let mut buff = Vec::new();
123        buff.append(&mut name.as_bytes().to_vec());
124        buff.push(0);
125        buff.append(&mut format!("{size}").as_bytes().to_vec());
126        buff.push(0);
127        self.send_blk(dev, &buff, 0, false).await
128    }
129
130    async fn send_blk<D: AsyncWrite + AsyncRead + Unpin>(
131        &mut self,
132        dev: &mut D,
133        data: &[u8],
134        pad: u8,
135        last: bool,
136    ) -> Result<()> {
137        let (len, p) = if data.len() > 128 {
138            (1024, STX)
139        } else {
140            (128, SOH)
141        };
142        let blk = if last { 0 } else { self.blk };
143        let mut err = None;
144
145        loop {
146            if self.retries == 0 {
147                return Err(err.unwrap_or(Error::new(ErrorKind::BrokenPipe, "retry too much")));
148            }
149
150            dev.write_all(&[p, blk, !blk]).await?;
151
152            let mut buf = vec![pad; len];
153            buf[..data.len()].copy_from_slice(data);
154            dev.write_all(&buf).await?;
155
156            if self.crc_mode {
157                let chsum = crc16_ccitt(0, &buf);
158                let crc1 = (chsum >> 8) as u8;
159                let crc2 = (chsum & 0xff) as u8;
160                dev.write_all(&[crc1, crc2]).await?;
161            }
162            dev.flush().await?;
163
164            match self.wait_ack(dev).await {
165                Ok(_) => break,
166                Err(e) => {
167                    err = Some(e);
168                    self.retries -= 1;
169                }
170            }
171        }
172
173        if self.blk == u8::MAX {
174            self.blk = 0;
175        } else {
176            self.blk += 1;
177        }
178
179        Ok(())
180    }
181}