1use std::{
2 fs::File,
3 io::*,
4 path::PathBuf,
5 time::{Duration, Instant},
6};
7
8use colored::Colorize;
9
10mod crc;
11mod ymodem;
12
13macro_rules! trace {
14 ($($arg:tt)*) => {{
15 println!("\r\n{}", &std::fmt::format(format_args!($($arg)*)).bright_black());
16 }};
17}
18
19pub struct UbootShell {
20 pub tx: Option<Box<dyn Write + Send>>,
21 pub rx: Option<Box<dyn Read + Send>>,
22 perfix: String,
23}
24
25impl UbootShell {
26 pub fn new(tx: impl Write + Send + 'static, rx: impl Read + Send + 'static) -> Result<Self> {
28 let mut s = Self {
29 tx: Some(Box::new(tx)),
30 rx: Some(Box::new(rx)),
31 perfix: "".to_string(),
32 };
33 s.wait_for_shell()?;
34 trace!("shell ready, perfix: `{}`", s.perfix);
35 Ok(s)
36 }
37
38 fn rx(&mut self) -> &mut Box<dyn Read + Send> {
39 self.rx.as_mut().unwrap()
40 }
41
42 fn tx(&mut self) -> &mut Box<dyn Write + Send> {
43 self.tx.as_mut().unwrap()
44 }
45
46 fn wait_for_shell(&mut self) -> Result<()> {
47 let mut buf = [0u8; 1];
48 let mut history: Vec<u8> = Vec::new();
49 const CTRL_C: u8 = 0x03;
50
51 let mut last = Instant::now();
52
53 loop {
54 match self.rx().read(&mut buf) {
55 Ok(n) => {
56 if n == 1 {
57 let ch = buf[0];
58 if ch == b'\n' && history.last() != Some(&b'\r') {
59 print_raw(b"\r");
60 history.push(b'\r');
61 }
62 history.push(ch);
63
64 print_raw(&buf);
65
66 if history.ends_with(c"<INTERRUPT>".to_bytes()) {
67 let line = history.split(|n| *n == b'\n').next_back().unwrap();
68 let s = String::from_utf8_lossy(line);
69 self.perfix = s.trim().replace("<INTERRUPT>", "").trim().to_string();
70
71 return Ok(());
72 }
73
74 if last.elapsed() > Duration::from_millis(20) {
75 let _ = self.tx().write_all(&[CTRL_C]);
76 last = Instant::now();
77 }
78 }
79 }
80
81 Err(ref e) if e.kind() == ErrorKind::TimedOut => {
82 continue;
83 }
84 Err(e) => {
85 return Err(e);
86 }
87 }
88 }
89 }
90
91 pub fn wait_for_reply(&mut self, val: &str) -> Result<String> {
92 let mut reply = Vec::new();
93 let mut buff = [0u8; 1];
94
95 trace!("wait for `{}`", val);
96 loop {
97 self.rx().read_exact(&mut buff)?;
98 reply.push(buff[0]);
99 print_raw(&buff);
100
101 if reply.ends_with(val.as_bytes()) {
102 break;
103 }
104 }
105 Ok(String::from_utf8_lossy(&reply)
106 .trim()
107 .trim_end_matches(&self.perfix)
108 .to_string())
109 }
110
111 pub fn cmd_without_reply(&mut self, cmd: &str) -> Result<()> {
112 self.tx().write_all(cmd.as_bytes())?;
113 self.tx().write_all("\r\n".as_bytes())?;
114 self.tx().flush()?;
115 self.wait_for_reply(cmd)?;
116 trace!("cmd ok");
117 Ok(())
118 }
119
120 pub fn cmd(&mut self, cmd: &str) -> Result<String> {
121 self.cmd_without_reply(cmd)?;
122 let perfix = self.perfix.clone();
123 self.wait_for_reply(&perfix)
124 }
125
126 pub fn set_env(&mut self, name: impl Into<String>, value: impl Into<String>) -> Result<()> {
127 self.cmd(&format!("setenv {} {}", name.into(), value.into()))?;
128 Ok(())
129 }
130
131 pub fn env(&mut self, name: impl Into<String>) -> Result<String> {
132 let name = name.into();
133 let s = self.cmd(&format!("echo ${}", name))?;
134 if s.is_empty() {
135 return Err(Error::new(
136 ErrorKind::NotFound,
137 format!("env {} not found", name),
138 ));
139 }
140 Ok(s)
141 }
142
143 pub fn env_int(&mut self, name: impl Into<String>) -> Result<usize> {
144 let name = name.into();
145 let line = self.env(&name)?;
146 parse_int(&line).ok_or(Error::new(
147 ErrorKind::InvalidData,
148 format!("env {name} is not a number"),
149 ))
150 }
151
152 pub fn loady(
153 &mut self,
154 addr: usize,
155 file: impl Into<PathBuf>,
156 on_progress: impl Fn(usize, usize),
157 ) -> Result<String> {
158 self.cmd_without_reply(&format!("loady {:#x}", addr,))?;
159 let crc = self.wait_for_load_crc()?;
160 let mut p = ymodem::Ymodem::new(crc);
161
162 let file = file.into();
163 let name = file.file_name().unwrap().to_str().unwrap();
164
165 let mut file = File::open(&file).unwrap();
166
167 let size = file.metadata().unwrap().len() as usize;
168
169 p.send(self, &mut file, name, size, |p| {
170 on_progress(p, size);
171 })?;
172 let perfix = self.perfix.clone();
173 self.wait_for_reply(&perfix)
174 }
175
176 fn wait_for_load_crc(&mut self) -> Result<bool> {
177 let mut reply = Vec::new();
178 let mut buff = [0u8; 1];
179 loop {
180 self.rx().read_exact(&mut buff)?;
181 reply.push(buff[0]);
182 let _ = stdout().write_all(&buff);
183
184 if reply.ends_with(b"C") {
185 return Ok(true);
186 }
187 let res = String::from_utf8_lossy(&reply);
188 if res.contains("try 'help'") {
189 panic!("{}", res);
190 }
191 }
192 }
193}
194
195impl Read for UbootShell {
196 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
197 self.rx().read(buf)
198 }
199}
200
201impl Write for UbootShell {
202 fn write(&mut self, buf: &[u8]) -> Result<usize> {
203 self.tx().write(buf)
204 }
205
206 fn flush(&mut self) -> Result<()> {
207 self.tx().flush()
208 }
209}
210
211fn parse_int(line: &str) -> Option<usize> {
212 let mut line = line.trim();
213 let mut radix = 10;
214 if line.starts_with("0x") {
215 line = &line[2..];
216 radix = 16;
217 }
218 u64::from_str_radix(line, radix).ok().map(|o| o as _)
219}
220
221fn print_raw(buff: &[u8]) {
222 stdout().write_all(buff).unwrap();
223}