119 lines
4.1 KiB
Rust
119 lines
4.1 KiB
Rust
//! Absolute-deadline blocking TCP adapter for single-command clients.
|
|
#[cfg(not(target_os = "wasi"))]
|
|
use std::net::ToSocketAddrs;
|
|
use std::{
|
|
io::{self, Read, Write},
|
|
net::{SocketAddr, TcpStream},
|
|
time::Instant,
|
|
};
|
|
|
|
#[cfg(not(target_os = "wasi"))]
|
|
fn resolve(host: &str, port: u16, _deadline: Instant) -> io::Result<Vec<SocketAddr>> {
|
|
Ok((host, port).to_socket_addrs()?.take(8).collect())
|
|
}
|
|
#[cfg(target_os = "wasi")]
|
|
fn resolve(host: &str, port: u16, deadline: Instant) -> io::Result<Vec<SocketAddr>> {
|
|
use wasi::sockets::{ip_name_lookup, network::IpAddress};
|
|
if let Ok(ip) = host.parse::<std::net::IpAddr>() {
|
|
return Ok(vec![SocketAddr::new(ip, port)]);
|
|
}
|
|
let network = wasi::sockets::instance_network::instance_network();
|
|
let addresses = ip_name_lookup::resolve_addresses(&network, host)
|
|
.map_err(|_| io::Error::other("DNS resolution failed"))?;
|
|
let ready = addresses.subscribe();
|
|
let remaining = deadline
|
|
.checked_duration_since(Instant::now())
|
|
.ok_or_else(timeout)?;
|
|
let timer = wasi::clocks::monotonic_clock::subscribe_duration(
|
|
remaining.as_nanos().min(u64::MAX as u128) as u64,
|
|
);
|
|
let mut result = Vec::new();
|
|
loop {
|
|
if Instant::now() >= deadline {
|
|
return Err(timeout());
|
|
}
|
|
match addresses.resolve_next_address() {
|
|
Ok(Some(ip)) => {
|
|
let ip = match ip {
|
|
IpAddress::Ipv4((a, b, c, d)) => {
|
|
std::net::IpAddr::V4(std::net::Ipv4Addr::new(a, b, c, d))
|
|
}
|
|
IpAddress::Ipv6((a, b, c, d, e, f, g, h)) => {
|
|
std::net::IpAddr::V6(std::net::Ipv6Addr::new(a, b, c, d, e, f, g, h))
|
|
}
|
|
};
|
|
result.push(SocketAddr::new(ip, port));
|
|
if result.len() == 8 {
|
|
break;
|
|
}
|
|
}
|
|
Ok(None) => break,
|
|
Err(wasi::sockets::network::ErrorCode::WouldBlock) => {
|
|
if wasi::io::poll::poll(&[&ready, &timer]).contains(&1) {
|
|
return Err(timeout());
|
|
}
|
|
}
|
|
Err(_) => return Err(io::Error::other("DNS resolution failed")),
|
|
}
|
|
}
|
|
Ok(result)
|
|
}
|
|
|
|
pub(crate) fn timeout() -> io::Error {
|
|
io::Error::new(io::ErrorKind::TimedOut, "timeout")
|
|
}
|
|
/// Captured raw streams, sharing a single byte budget. None means no exit-status (e.g. signal).
|
|
#[derive(Default, Debug)]
|
|
pub struct CommandOutput {
|
|
pub exit_code: Option<u32>,
|
|
pub stdout: Vec<u8>,
|
|
pub stderr: Vec<u8>,
|
|
pub truncated: bool,
|
|
}
|
|
/// Each underlying operation checks time, including peers sending tiny packets indefinitely.
|
|
pub struct DeadlineTcp {
|
|
stream: TcpStream,
|
|
deadline: Instant,
|
|
}
|
|
impl DeadlineTcp {
|
|
pub fn connect(host: &str, port: u16, deadline: Instant) -> io::Result<Self> {
|
|
let addresses = resolve(host, port, deadline)?;
|
|
let mut last = io::Error::new(io::ErrorKind::ConnectionRefused, "connection refused");
|
|
for address in addresses {
|
|
let remaining = deadline
|
|
.checked_duration_since(Instant::now())
|
|
.ok_or_else(timeout)?;
|
|
match TcpStream::connect_timeout(&address, remaining) {
|
|
Ok(stream) => return Ok(Self { stream, deadline }),
|
|
Err(e) => last = e,
|
|
}
|
|
}
|
|
Err(last)
|
|
}
|
|
fn prepare(&self) -> io::Result<()> {
|
|
let remaining = self
|
|
.deadline
|
|
.checked_duration_since(Instant::now())
|
|
.filter(|d| !d.is_zero())
|
|
.ok_or_else(timeout)?;
|
|
self.stream.set_read_timeout(Some(remaining))?;
|
|
self.stream.set_write_timeout(Some(remaining))
|
|
}
|
|
}
|
|
impl Read for DeadlineTcp {
|
|
fn read(&mut self, b: &mut [u8]) -> io::Result<usize> {
|
|
self.prepare()?;
|
|
self.stream.read(b)
|
|
}
|
|
}
|
|
impl Write for DeadlineTcp {
|
|
fn write(&mut self, b: &[u8]) -> io::Result<usize> {
|
|
self.prepare()?;
|
|
self.stream.write(b)
|
|
}
|
|
fn flush(&mut self) -> io::Result<()> {
|
|
self.prepare()?;
|
|
self.stream.flush()
|
|
}
|
|
}
|