opentitanlib/util/
testing.rs1use 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#[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}
45pub 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}