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 std::time::Duration;
6
7use anyhow::{Result, anyhow};
8use regex::{Captures, Regex};
9
10use crate::io::console::{ConsoleDevice, ConsoleError, ConsoleExt};
11
12pub struct UartConsole {
13    timeout: Option<Duration>,
14    exit_success: Option<Regex>,
15    exit_failure: Option<Regex>,
16    buffer: String,
17}
18
19#[derive(Clone, Copy, Debug, PartialEq, Eq)]
20pub enum ExitStatus {
21    Timeout,
22    ExitSuccess,
23    ExitFailure,
24}
25
26impl UartConsole {
27    const BUFFER_LEN: usize = 32768;
28
29    pub fn new(
30        timeout: Option<Duration>,
31        exit_success: Option<Regex>,
32        exit_failure: Option<Regex>,
33    ) -> Self {
34        Self {
35            timeout,
36            exit_success,
37            exit_failure,
38            buffer: String::new(),
39        }
40    }
41
42    // Runs an interactive console until CTRL_C is received.
43    pub fn interact<T>(&mut self, device: &T, quiet: bool) -> Result<ExitStatus>
44    where
45        T: ConsoleDevice + ?Sized,
46    {
47        crate::util::runtime::block_on(self.interact_async(device, quiet))
48    }
49
50    // Runs an interactive console until CTRL_C is received.  Uses `mio` library to simultaneously
51    // wait for data from UART or from stdin, without need for timeouts and repeated calls.
52    pub async fn interact_async<T>(&mut self, device: &T, quiet: bool) -> Result<ExitStatus>
53    where
54        T: ConsoleDevice + ?Sized,
55    {
56        let device: &dyn ConsoleDevice = if quiet { &device } else { &device.logged() };
57
58        let timeout = self.timeout;
59        let rx = async {
60            loop {
61                self.uart_read(device).await?;
62                if self
63                    .exit_success
64                    .as_ref()
65                    .map(|rx| rx.is_match(&self.buffer))
66                    == Some(true)
67                {
68                    return Ok(ExitStatus::ExitSuccess);
69                }
70                if self
71                    .exit_failure
72                    .as_ref()
73                    .map(|rx| rx.is_match(&self.buffer))
74                    == Some(true)
75                {
76                    return Ok(ExitStatus::ExitFailure);
77                }
78            }
79        };
80        let timeout = async {
81            if let Some(timeout) = timeout {
82                tokio::time::sleep(timeout).await;
83            } else {
84                std::future::pending().await
85            }
86        };
87
88        tokio::select! {
89            v = rx => v,
90            _ = timeout => Ok(ExitStatus::Timeout),
91        }
92    }
93
94    /// Returns `true` if any regular expressions are used to match the streamed output.  If so,
95    /// then this struct will keep a window of the most recent output, and take care not to read
96    /// any more characters from the underlying stream should one of the regular expressions
97    /// match.
98    fn uses_regex(&self) -> bool {
99        self.exit_success.is_some() || self.exit_failure.is_some()
100    }
101
102    // Maintain a buffer for the exit regexes to match against.
103    fn append_buffer(&mut self, data: &[u8]) {
104        self.buffer.push_str(&String::from_utf8_lossy(data));
105        while self.buffer.len() > UartConsole::BUFFER_LEN {
106            self.buffer.remove(0);
107        }
108    }
109
110    // Read from the console device and process the data read.
111    async fn uart_read<T>(&mut self, device: &T) -> Result<()>
112    where
113        T: ConsoleDevice + ?Sized,
114    {
115        let mut ch = 0;
116
117        // Read one byte at a time to avoid the risk of consuming output past a match.
118        let len =
119            std::future::poll_fn(|cx| device.poll_read(cx, std::slice::from_mut(&mut ch))).await?;
120
121        if len == 0 {
122            return Ok(());
123        }
124
125        if self.uses_regex() {
126            self.append_buffer(std::slice::from_ref(&ch));
127        }
128        Ok(())
129    }
130
131    pub fn captures(&self, status: ExitStatus) -> Option<Captures<'_>> {
132        match status {
133            ExitStatus::ExitSuccess => self
134                .exit_success
135                .as_ref()
136                .and_then(|rx| rx.captures(&self.buffer)),
137            ExitStatus::ExitFailure => self
138                .exit_failure
139                .as_ref()
140                .and_then(|rx| rx.captures(&self.buffer)),
141            _ => None,
142        }
143    }
144
145    /// Wait on the console until the regex matches the input.
146    ///
147    /// The input is processed one byte at a time, and is accumulated until match happens.
148    pub fn wait_for_bytes<T>(device: &T, rx: &str, timeout: Duration) -> Result<Vec<String>>
149    where
150        T: ConsoleDevice + ?Sized,
151    {
152        let mut console = UartConsole::new(Some(timeout), Some(Regex::new(rx)?), None);
153        let result = console.interact(device, false)?;
154        println!();
155        match result {
156            ExitStatus::ExitSuccess => {
157                let caps = console.captures(ExitStatus::ExitSuccess).expect("capture");
158                let mut vec = Vec::new();
159                for c in caps.iter() {
160                    match c {
161                        None => vec.push(String::new()),
162                        _ => vec.push(c.unwrap().as_str().to_owned()),
163                    }
164                }
165                Ok(vec)
166            }
167            ExitStatus::Timeout => Err(ConsoleError::TimedOut)?,
168            _ => Err(anyhow!("Impossible result: {:?}", result)),
169        }
170    }
171
172    /// Wait on the console until the regex matches the input.
173    ///
174    /// The input is processed one line at a time. Lines that does not match the input is discarded.
175    pub fn wait_for<T>(device: &T, rx: &str, timeout: Duration) -> Result<Vec<String>>
176    where
177        T: ConsoleDevice + ?Sized,
178    {
179        device.logged().wait_for_line(Regex::new(rx)?, timeout)
180    }
181}
opentitanlib/uart/console.rs

opentitanlib/uart/
console.rs
