opentitanlib/util/
testing.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::cell::RefCell;
6use std::ffi::OsStr;
7use std::io::Write;
8use std::pin::Pin;
9use std::task::{Poll, ready};
10
11use anyhow::{Context, Result, anyhow};
12use tokio::io::AsyncRead;
13
14use crate::io::console::ConsoleDevice;
15use crate::util::runtime::MultiWaker;
16
17// The transfer state allows us to intentionally inject errors into
18// the data stream to test error handling.
19#[derive(Default)]
20pub struct TransferState {
21    nbytes: usize,
22    corrupt: Vec<usize>,
23}
24impl TransferState {
25    pub fn new(corrupt: &[usize]) -> Self {
26        TransferState {
27            nbytes: 0,
28            corrupt: corrupt.to_vec(),
29        }
30    }
31    fn maybe_corrupt(&mut self, buf: &mut [u8]) {
32        while !self.corrupt.is_empty() {
33            let mut index = self.corrupt[0];
34            if index >= self.nbytes && index < self.nbytes + buf.len() {
35                index -= self.nbytes;
36                buf[index] = !buf[index];
37            } else {
38                break;
39            }
40            self.corrupt.remove(0);
41        }
42        self.nbytes += buf.len();
43    }
44}
45// A convenience wrapper for spawning a child process and
46// interacting with its stdin/stdout.
47pub struct ChildConsole {
48    child: RefCell<std::process::Child>,
49    stdout: Option<RefCell<tokio::process::ChildStdout>>,
50    rd: RefCell<TransferState>,
51    wr: RefCell<TransferState>,
52    multi_waker: MultiWaker,
53}
54
55impl ChildConsole {
56    pub fn spawn_corrupt<S: AsRef<OsStr>>(
57        argv: &[S],
58        rd: TransferState,
59        wr: TransferState,
60    ) -> Result<Self> {
61        let mut command = std::process::Command::new(argv[0].as_ref());
62        for arg in &argv[1..] {
63            command.arg(arg.as_ref());
64        }
65        let mut child = command
66            .stdin(std::process::Stdio::piped())
67            .stdout(std::process::Stdio::piped())
68            .stderr(std::process::Stdio::inherit())
69            .spawn()?;
70        let _runtime_guard = crate::util::runtime().enter();
71        let stdout = match child.stdout.take() {
72            Some(v) => Some(RefCell::new(tokio::process::ChildStdout::from_std(v)?)),
73            None => None,
74        };
75        Ok(ChildConsole {
76            child: RefCell::new(child),
77            stdout,
78            rd: RefCell::new(rd),
79            wr: RefCell::new(wr),
80            multi_waker: MultiWaker::new(),
81        })
82    }
83
84    pub fn spawn<S: AsRef<OsStr>>(argv: &[S]) -> Result<Self> {
85        Self::spawn_corrupt(argv, Default::default(), Default::default())
86    }
87
88    pub fn wait(&self) -> Result<std::process::ExitStatus> {
89        let mut child = self.child.borrow_mut();
90        Ok(child.wait()?)
91    }
92}
93
94impl ConsoleDevice for ChildConsole {
95    fn poll_read(&self, cx: &mut std::task::Context<'_>, buf: &mut [u8]) -> Poll<Result<usize>> {
96        let mut stdout = self
97            .stdout
98            .as_ref()
99            .context("child has no stdout")?
100            .borrow_mut();
101        let mut read_buf = tokio::io::ReadBuf::new(buf);
102        ready!(
103            self.multi_waker
104                .poll_with(cx, |cx| Pin::new(&mut *stdout).poll_read(cx, &mut read_buf))
105        )?;
106        let n = read_buf.filled().len();
107        let mut rd = self.rd.borrow_mut();
108        rd.maybe_corrupt(read_buf.filled_mut());
109        Poll::Ready(Ok(n))
110    }
111
112    fn write(&self, buf: &[u8]) -> Result<()> {
113        let mut child = self.child.borrow_mut();
114        if let Some(stdin) = &mut child.stdin {
115            let mut data = buf.to_vec();
116            let mut wr = self.wr.borrow_mut();
117            wr.maybe_corrupt(&mut data);
118            stdin.write_all(&data)?;
119            Ok(())
120        } else {
121            Err(anyhow!("child has no stdin"))
122        }
123    }
124}