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 history: Vec<u8> = Vec::new();
48 const CTRL_C: u8 = 0x03;
49
50 let mut last = Instant::now();
51 let mut is_interrupt_found = false;
52 let mut is_shell_ok = false;
53 const INT_STR: &str = "<INTERRUPT>";
54 const INT: &[u8] = INT_STR.as_bytes();
55
56 trace!("wait for `{INT_STR}`");
57 loop {
58 match self.read_byte() {
59 Ok(ch) => {
60 if ch == b'\n' && history.last() != Some(&b'\r') {
61 print_raw(b"\r");
62 history.push(b'\r');
63 }
64 history.push(ch);
65
66 print_raw(&[ch]);
67
68 if history.ends_with(INT) && !is_interrupt_found {
69 let line = history.split(|n| *n == b'\n').next_back().unwrap();
70 let s = String::from_utf8_lossy(line);
71 self.perfix = s.trim().replace(INT_STR, "").trim().to_string();
72 is_interrupt_found = true;
73 trace!("clear");
74 let _ = self.rx().read_to_end(&mut vec![]);
75 trace!("test shell");
76 let ret = self.tx().write_all("testshell\r\n".as_bytes());
77 trace!("test shell ret {:?}", ret);
78 }
79
80 if is_interrupt_found && history.ends_with("\'help\'".as_bytes()) {
81 is_shell_ok = true;
82 }
83
84 if is_shell_ok && history.ends_with(self.perfix.as_bytes()) {
85 return Ok(());
86 }
87
88 if last.elapsed() > Duration::from_millis(20) && !is_interrupt_found {
89 let _ = self.tx().write_all(&[CTRL_C]);
90 last = Instant::now();
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
104 fn read_byte(&mut self) -> Result<u8> {
105 let mut buff = [0u8; 1];
106 let time_out = Duration::from_secs(5);
107 let start = Instant::now();
108
109 loop {
110 match self.rx().read_exact(&mut buff) {
111 Ok(_) => return Ok(buff[0]),
112 Err(e) => {
113 if e.kind() == ErrorKind::TimedOut {
114 if start.elapsed() > time_out {
115 return Err(std::io::Error::new(
116 std::io::ErrorKind::TimedOut,
117 "Timeout",
118 ));
119 }
120 } else {
121 return Err(e);
122 }
123 }
124 }
125 }
126 }
127
128 pub fn wait_for_reply(&mut self, val: &str) -> Result<String> {
129 let mut reply = Vec::new();
130
131 trace!("wait for `{}`", val);
132 loop {
133 let byte = self.read_byte()?;
134 reply.push(byte);
135 print_raw(&[byte]);
136
137 if reply.ends_with(val.as_bytes()) {
138 break;
139 }
140 }
141 Ok(String::from_utf8_lossy(&reply)
142 .trim()
143 .trim_end_matches(&self.perfix)
144 .to_string())
145 }
146
147 pub fn cmd_without_reply(&mut self, cmd: &str) -> Result<()> {
148 self.tx().write_all(cmd.as_bytes())?;
149 self.tx().write_all("\r\n".as_bytes())?;
150 self.tx().flush()?;
151 self.wait_for_reply(cmd)?;
152 trace!("cmd ok");
153 Ok(())
154 }
155
156 pub fn cmd(&mut self, cmd: &str) -> Result<String> {
157 self.cmd_without_reply(cmd)?;
158 let perfix = self.perfix.clone();
159 self.wait_for_reply(&perfix)
160 }
161
162 pub fn set_env(&mut self, name: impl Into<String>, value: impl Into<String>) -> Result<()> {
163 self.cmd(&format!("setenv {} {}", name.into(), value.into()))?;
164 Ok(())
165 }
166
167 pub fn env(&mut self, name: impl Into<String>) -> Result<String> {
168 let name = name.into();
169 let s = self.cmd(&format!("echo ${}", name))?;
170 if s.is_empty() {
171 return Err(Error::new(
172 ErrorKind::NotFound,
173 format!("env {} not found", name),
174 ));
175 }
176 Ok(s)
177 }
178
179 pub fn env_int(&mut self, name: impl Into<String>) -> Result<usize> {
180 let name = name.into();
181 let line = self.env(&name)?;
182 parse_int(&line).ok_or(Error::new(
183 ErrorKind::InvalidData,
184 format!("env {name} is not a number"),
185 ))
186 }
187
188 pub fn loady(
189 &mut self,
190 addr: usize,
191 file: impl Into<PathBuf>,
192 on_progress: impl Fn(usize, usize),
193 ) -> Result<String> {
194 self.cmd_without_reply(&format!("loady {:#x}", addr,))?;
195 let crc = self.wait_for_load_crc()?;
196 let mut p = ymodem::Ymodem::new(crc);
197
198 let file = file.into();
199 let name = file.file_name().unwrap().to_str().unwrap();
200
201 let mut file = File::open(&file).unwrap();
202
203 let size = file.metadata().unwrap().len() as usize;
204
205 p.send(self, &mut file, name, size, |p| {
206 on_progress(p, size);
207 })?;
208 let perfix = self.perfix.clone();
209 self.wait_for_reply(&perfix)
210 }
211
212 fn wait_for_load_crc(&mut self) -> Result<bool> {
213 let mut reply = Vec::new();
214 loop {
215 let byte = self.read_byte()?;
216 reply.push(byte);
217 print_raw(&[byte]);
218
219 if reply.ends_with(b"C") {
220 return Ok(true);
221 }
222 let res = String::from_utf8_lossy(&reply);
223 if res.contains("try 'help'") {
224 panic!("{}", res);
225 }
226 }
227 }
228}
229
230impl Read for UbootShell {
231 fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
232 self.rx().read(buf)
233 }
234}
235
236impl Write for UbootShell {
237 fn write(&mut self, buf: &[u8]) -> Result<usize> {
238 self.tx().write(buf)
239 }
240
241 fn flush(&mut self) -> Result<()> {
242 self.tx().flush()
243 }
244}
245
246fn parse_int(line: &str) -> Option<usize> {
247 let mut line = line.trim();
248 let mut radix = 10;
249 if line.starts_with("0x") {
250 line = &line[2..];
251 radix = 16;
252 }
253 u64::from_str_radix(line, radix).ok().map(|o| o as _)
254}
255
256fn print_raw(buff: &[u8]) {
257 stdout().write_all(buff).unwrap();
258}