opentitanlib/uart/
console.rs

1// Copyright lowRISC contributors (OpenTitan project).
2// Licensed under the Apache License, Version 2.0, see LICENSE for details.
3// SPDX-License-Identifier: Apache-2.0
4
5use anyhow::{Result, anyhow};
6use regex::{Captures, Regex};
7use std::fs::File;
8use std::io::Write;
9use std::os::fd::AsFd;
10use std::time::{Duration, Instant, SystemTime};
11
12use tokio::io::AsyncReadExt;
13
14use crate::io::console::{ConsoleDevice, ConsoleError};
15
16#[derive(Default)]
17pub struct UartConsole {
18    pub logfile: Option<File>,
19    pub timeout: Option<Duration>,
20    pub deadline: Option<Instant>,
21    pub exit_success: Option<Regex>,
22    pub exit_failure: Option<Regex>,
23    pub timestamp: bool,
24    pub buffer: String,
25    pub newline: bool,
26    pub break_en: bool,
27}
28
29#[derive(Clone, Copy, Debug, PartialEq, Eq)]
30pub enum ExitStatus {
31    None,
32    CtrlC,
33    Timeout,
34    ExitSuccess,
35    ExitFailure,
36}
37
38// Creates a vtable for implementors of Read and AsFd traits.
39pub trait ReadAsFd: tokio::io::AsyncRead + AsFd + std::marker::Unpin {}
40impl<T: tokio::io::AsyncRead + AsFd + std::marker::Unpin> ReadAsFd for T {}
41
42impl UartConsole {
43    const CTRL_B: u8 = 2;
44    const CTRL_C: u8 = 3;
45    const BUFFER_LEN: usize = 32768;
46
47    // Runs an interactive console until CTRL_C is received.
48    pub fn interact<T>(
49        &mut self,
50        device: &T,
51        stdin: Option<&mut dyn ReadAsFd>,
52        stdout: Option<&mut dyn Write>,
53    ) -> Result<ExitStatus>
54    where
55        T: ConsoleDevice + ?Sized,
56    {
57        if let Some(timeout) = &self.timeout {
58            self.deadline = Some(Instant::now() + *timeout);
59        }
60        crate::util::runtime::block_on(self.interact_async(device, stdin, stdout))
61    }
62
63    // Runs an interactive console until CTRL_C is received.  Uses `mio` library to simultaneously
64    // wait for data from UART or from stdin, without need for timeouts and repeated calls.
65    async fn interact_async<T>(
66        &mut self,
67        device: &T,
68        mut stdin: Option<&mut dyn ReadAsFd>,
69        mut stdout: Option<&mut dyn Write>,
70    ) -> Result<ExitStatus>
71    where
72        T: ConsoleDevice + ?Sized,
73    {
74        let mut break_en = self.break_en;
75        let deadline = self.deadline;
76        let tx = async {
77            if let Some(stdin) = stdin.as_mut() {
78                Self::process_input(&mut break_en, device, stdin).await
79            } else {
80                std::future::pending().await
81            }
82        };
83        let rx = async {
84            loop {
85                self.uart_read(device, &mut stdout).await?;
86                if self
87                    .exit_success
88                    .as_ref()
89                    .map(|rx| rx.is_match(&self.buffer))
90                    == Some(true)
91                {
92                    return Ok(ExitStatus::ExitSuccess);
93                }
94                if self
95                    .exit_failure
96                    .as_ref()
97                    .map(|rx| rx.is_match(&self.buffer))
98                    == Some(true)
99                {
100                    return Ok(ExitStatus::ExitFailure);
101                }
102            }
103        };
104        let deadline = async {
105            if let Some(deadline) = deadline {
106                tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await;
107            } else {
108                std::future::pending().await
109            }
110        };
111
112        let r = tokio::select! {
113            v = tx => v,
114            v = rx => v,
115            _ = deadline => Ok(ExitStatus::Timeout),
116        };
117        self.break_en = break_en;
118        r
119    }
120
121    /// Returns `true` if any regular expressions are used to match the streamed output.  If so,
122    /// then this struct will keep a window of the most recent output, and take care not to read
123    /// any more characters from the underlying stream should one of the regular expressions
124    /// match.
125    fn uses_regex(&self) -> bool {
126        self.exit_success.is_some() || self.exit_failure.is_some()
127    }
128
129    // Maintain a buffer for the exit regexes to match against.
130    fn append_buffer(&mut self, data: &[u8]) {
131        self.buffer.push_str(&String::from_utf8_lossy(data));
132        while self.buffer.len() > UartConsole::BUFFER_LEN {
133            self.buffer.remove(0);
134        }
135    }
136
137    // Read from the console device and process the data read.
138    async fn uart_read<T>(&mut self, device: &T, stdout: &mut Option<&mut dyn Write>) -> Result<()>
139    where
140        T: ConsoleDevice + ?Sized,
141    {
142        let mut buf = [0u8; 1024];
143        let effective_buf = if self.uses_regex() {
144            // Read one byte at a time when matching, to avoid the risk of consuming output past a
145            // match.
146            &mut buf[..1]
147        } else {
148            &mut buf
149        };
150        let len = std::future::poll_fn(|cx| device.console_poll_read(cx, effective_buf)).await?;
151        for i in 0..len {
152            if self.timestamp && self.newline {
153                let t = humantime::format_rfc3339_millis(SystemTime::now());
154                stdout.as_mut().map_or(Ok(()), |out| {
155                    out.write_fmt(format_args!("[{}  console]", t))
156                })?;
157                self.newline = false;
158            }
159            self.newline = buf[i] == b'\n';
160            stdout
161                .as_mut()
162                .map_or(Ok(()), |out| out.write_all(&buf[i..i + 1]))?;
163        }
164        stdout.as_mut().map_or(Ok(()), |out| out.flush())?;
165
166        // If we're logging, save it to the logfile.
167        self.logfile
168            .as_mut()
169            .map_or(Ok(()), |f| f.write_all(&buf[..len]))?;
170        if self.uses_regex() {
171            self.append_buffer(&buf[..len]);
172        }
173        Ok(())
174    }
175
176    async fn process_input<T>(
177        break_en: &mut bool,
178        device: &T,
179        stdin: &mut dyn ReadAsFd,
180    ) -> Result<ExitStatus>
181    where
182        T: ConsoleDevice + ?Sized,
183    {
184        loop {
185            let mut buf = [0u8; 256];
186            let len = stdin.read(&mut buf).await?;
187            if len == 1 {
188                if buf[0] == UartConsole::CTRL_C {
189                    return Ok(ExitStatus::CtrlC);
190                }
191                if buf[0] == UartConsole::CTRL_B {
192                    *break_en = !*break_en;
193                    eprint!(
194                        "\r\n{} break",
195                        if *break_en { "Setting" } else { "Clearing" }
196                    );
197                    let b = device.set_break(*break_en);
198                    if b.is_err() {
199                        eprint!(": {:?}", b);
200                    }
201                    eprint!("\r\n");
202                    continue;
203                }
204            }
205            if len > 0 {
206                device.console_write(&buf[..len])?;
207            }
208        }
209    }
210
211    pub fn captures(&self, status: ExitStatus) -> Option<Captures<'_>> {
212        match status {
213            ExitStatus::ExitSuccess => self
214                .exit_success
215                .as_ref()
216                .and_then(|rx| rx.captures(&self.buffer)),
217            ExitStatus::ExitFailure => self
218                .exit_failure
219                .as_ref()
220                .and_then(|rx| rx.captures(&self.buffer)),
221            _ => None,
222        }
223    }
224
225    pub fn wait_for<T>(device: &T, rx: &str, timeout: Duration) -> Result<Vec<String>>
226    where
227        T: ConsoleDevice + ?Sized,
228    {
229        let mut console = UartConsole {
230            timestamp: true,
231            newline: true,
232            timeout: Some(timeout),
233            exit_success: Some(Regex::new(rx)?),
234            ..Default::default()
235        };
236        let mut stdout = std::io::stdout();
237        let result = console.interact(device, None, Some(&mut stdout))?;
238        println!();
239        match result {
240            ExitStatus::ExitSuccess => {
241                let caps = console.captures(ExitStatus::ExitSuccess).expect("capture");
242                let mut vec = Vec::new();
243                for c in caps.iter() {
244                    match c {
245                        None => vec.push(String::new()),
246                        _ => vec.push(c.unwrap().as_str().to_owned()),
247                    }
248                }
249                Ok(vec)
250            }
251            ExitStatus::Timeout => Err(ConsoleError::GenericError("Timed Out".into()).into()),
252            _ => Err(anyhow!("Impossible result: {:?}", result)),
253        }
254    }
255}
opentitanlib/uart/console.rs

opentitanlib/uart/
console.rs
