1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
|
use std::io::{self, Read as _};
use std::process::Output;
use std::thread;
use std::time::{Duration, Instant};
const POLL_INTERVAL: Duration = Duration::from_millis(50);
pub(super) fn wait_with_output_timeout(
mut child: std::process::Child,
timeout: Duration,
) -> io::Result<Output> {
let stdout = child
.stdout
.take()
.ok_or_else(|| io::Error::other("SSH command stdout was not captured"))?;
let stderr = child
.stderr
.take()
.ok_or_else(|| io::Error::other("SSH command stderr was not captured"))?;
let stdout = thread::spawn(move || {
let mut stdout = stdout;
let mut bytes = Vec::new();
stdout.read_to_end(&mut bytes).map(|_| bytes)
});
let stderr = thread::spawn(move || {
let mut stderr = stderr;
let mut bytes = Vec::new();
stderr.read_to_end(&mut bytes).map(|_| bytes)
});
let started = Instant::now();
let status = loop {
match child.try_wait() {
Ok(Some(status)) => break status,
Ok(None) => {}
Err(error) => {
let _ = child.kill();
let _ = child.wait();
let _ = stdout.join();
let _ = stderr.join();
return Err(error);
}
}
if started.elapsed() >= timeout {
let _ = child.kill();
let _ = child.wait();
let _ = stdout.join();
let _ = stderr.join();
return Err(io::Error::new(
io::ErrorKind::TimedOut,
"noninteractive SSH command timed out",
));
}
thread::sleep(POLL_INTERVAL);
};
let stdout = stdout
.join()
.map_err(|_| io::Error::other("SSH stdout reader panicked"))??;
let stderr = stderr
.join()
.map_err(|_| io::Error::other("SSH stderr reader panicked"))??;
Ok(Output {
status,
stdout,
stderr,
})
}
#[cfg(all(test, unix))]
mod tests {
use super::*;
use std::process::{Command, Stdio};
#[test]
fn timeout_kills_the_child() {
let mut command = Command::new("sh");
command
.arg("-c")
.arg("exec sleep 10")
.stdout(Stdio::piped())
.stderr(Stdio::piped());
let started = Instant::now();
let error = wait_with_output_timeout(command.spawn().unwrap(), Duration::from_millis(25))
.unwrap_err();
assert_eq!(error.kind(), io::ErrorKind::TimedOut);
assert!(started.elapsed() < Duration::from_secs(1));
}
}
|
