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 = device.coverage();
57        let device: &dyn ConsoleDevice = if quiet { &device } else { &device.logged() };
58
59        let timeout = self.timeout;
60        let rx = async {
61            loop {
62                self.uart_read(device).await?;
63                if self
64                    .exit_success
65                    .as_ref()
66                    .map(|rx| rx.is_match(&self.buffer))
67                    == Some(true)
68                {
69                    return Ok(ExitStatus::ExitSuccess);
70                }
71                if self
72                    .exit_failure
73                    .as_ref()
74                    .map(|rx| rx.is_match(&self.buffer))
75                    == Some(true)
76                {
77                    return Ok(ExitStatus::ExitFailure);
78                }
79            }
80        };
81        let timeout = async {
82            if let Some(timeout) = timeout {
83                tokio::time::sleep(timeout).await;
84            } else {
85                std::future::pending().await
86            }
87        };
88
89        tokio::select! {
90            v = rx => v,
91            _ = timeout => Ok(ExitStatus::Timeout),
92        }
93    }
94
95    /// Returns `true` if any regular expressions are used to match the streamed output.  If so,
96    /// then this struct will keep a window of the most recent output, and take care not to read
97    /// any more characters from the underlying stream should one of the regular expressions
98    /// match.
99    fn uses_regex(&self) -> bool {
100        self.exit_success.is_some() || self.exit_failure.is_some()
101    }
102
103    // Maintain a buffer for the exit regexes to match against.
104    fn append_buffer(&mut self, data: &[u8]) {
105        self.buffer.push_str(&String::from_utf8_lossy(data));
106        while self.buffer.len() > UartConsole::BUFFER_LEN {
107            self.buffer.remove(0);
108        }
109    }
110
111    // Read from the console device and process the data read.
112    async fn uart_read<T>(&mut self, device: &T) -> Result<()>
113    where
114        T: ConsoleDevice + ?Sized,
115    {
116        let mut ch = 0;
117
118        // Read one byte at a time to avoid the risk of consuming output past a match.
119        let len =
120            std::future::poll_fn(|cx| device.poll_read(cx, std::slice::from_mut(&mut ch))).await?;
121
122        if len == 0 {
123            return Ok(());
124        }
125
126        if self.uses_regex() {
127            self.append_buffer(std::slice::from_ref(&ch));
128        }
129        Ok(())
130    }
131
132    pub fn captures(&self, status: ExitStatus) -> Option<Captures<'_>> {
133        match status {
134            ExitStatus::ExitSuccess => self
135                .exit_success
136                .as_ref()
137                .and_then(|rx| rx.captures(&self.buffer)),
138            ExitStatus::ExitFailure => self
139                .exit_failure
140                .as_ref()
141                .and_then(|rx| rx.captures(&self.buffer)),
142            _ => None,
143        }
144    }
145
146    /// Wait on the console until the regex matches the input.
147    ///
148    /// The input is processed one byte at a time, and is accumulated until match happens.
149    pub fn wait_for_bytes<T>(device: &T, rx: &str, timeout: Duration) -> Result<Vec<String>>
150    where
151        T: ConsoleDevice + ?Sized,
152    {
153        let mut console = UartConsole::new(Some(timeout), Some(Regex::new(rx)?), None);
154        let result = console.interact(device, false)?;
155        println!();
156        match result {
157            ExitStatus::ExitSuccess => {
158                let caps = console.captures(ExitStatus::ExitSuccess).expect("capture");
159                let mut vec = Vec::new();
160                for c in caps.iter() {
161                    match c {
162                        None => vec.push(String::new()),
163                        _ => vec.push(c.unwrap().as_str().to_owned()),
164                    }
165                }
166                Ok(vec)
167            }
168            ExitStatus::Timeout => Err(ConsoleError::TimedOut)?,
169            _ => Err(anyhow!("Impossible result: {:?}", result)),
170        }
171    }
172
173    /// Wait on the console until the regex matches the input.
174    ///
175    /// The input is processed one line at a time. Lines that does not match the input is discarded.
176    pub fn wait_for<T>(device: &T, rx: &str, timeout: Duration) -> Result<Vec<String>>
177    where
178        T: ConsoleDevice + ?Sized,
179    {
180        device
181            .coverage()
182            .logged()
183            .wait_for_line(Regex::new(rx)?, timeout)
184    }
185
186    /// Wait on the console until the coverage profile end or skip anchor is received.
187    pub fn wait_for_coverage<T>(device: &T, timeout: Duration) -> Result<()>
188    where
189        T: ConsoleDevice + ?Sized,
190    {
191        device.coverage().logged().wait_for_coverage(timeout)
192    }
193}
opentitanlib/uart/console.rs

opentitanlib/uart/
console.rs
