From 436075a704b10788af44ce6c71108c3c01093348 Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Wed, 21 Oct 2026 23:56:32 +0100 Subject: [PATCH] Migrate to Tokio asynchronous runtime --- Cargo.lock | 78 ++---- Cargo.toml | 8 +- lib/host.rs | 102 ++++--- lib/mod.rs | 1 - lib/poller.rs | 124 --------- lib/proxy/control.rs | 624 +++++++++++++++++++++---------------------- lib/proxy/host.rs | 2 +- lib/proxy/mod.rs | 212 ++++++--------- lib/proxy/vm.rs | 2 +- lib/vm.rs | 98 ++++--- src/main.rs | 22 +- 11 files changed, 512 insertions(+), 761 deletions(-) delete mode 100644 lib/poller.rs diff --git a/Cargo.lock b/Cargo.lock index 6584f1b..d14297c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -436,15 +436,6 @@ dependencies = [ "memchr", ] -[[package]] -name = "concurrent-queue" -version = "2.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" -dependencies = [ - "crossbeam-utils", -] - [[package]] name = "const-oid" version = "0.10.2" @@ -492,12 +483,6 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" -[[package]] -name = "crossbeam-utils" -version = "0.8.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22ec99545bb0ed0ea7bb9b8e1e9122ea386ff8a48c0922e43f36d45ab09e0e80" - [[package]] name = "crypto-common" version = "0.2.1" @@ -940,12 +925,6 @@ version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d231dfb89cfffdbc30e7fc41579ed6066ad03abda9e567ccafae602b97ec5024" -[[package]] -name = "hermit-abi" -version = "0.5.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f154ce46856750ed433c8649605bf7ed2de3bc35fd9d2a9f30cddd873c80cb08" - [[package]] name = "hex" version = "0.4.3" @@ -1419,12 +1398,6 @@ version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "78b3ae25bc7c8c38cec158d1f2757ee79e9b3740fbc7ccf0e59e4b08d793fa89" -[[package]] -name = "linux-raw-sys" -version = "0.9.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12" - [[package]] name = "litemap" version = "0.7.4" @@ -1516,7 +1489,7 @@ version = "1.0.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "80e04d1dcff3aae0704555fe5fee3bcfaf3d1fdf8a7e521d5b9d2b42acb52cec" dependencies = [ - "hermit-abi 0.3.9", + "hermit-abi", "libc", "log", "wasi", @@ -1794,20 +1767,6 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "953ec861398dccce10c670dfeaf3ec4911ca479e9c02154b3a215178c5f566f2" -[[package]] -name = "polling" -version = "3.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d0e4f59085d47d8241c88ead0f274e8a0cb551f3625263c05eb8dd897c34218" -dependencies = [ - "cfg-if", - "concurrent-queue", - "hermit-abi 0.5.1", - "pin-project-lite", - "rustix 1.0.7", - "windows-sys 0.61.0", -] - [[package]] name = "portable-atomic" version = "1.9.0" @@ -2007,23 +1966,10 @@ dependencies = [ "bitflags 2.13.2", "errno", "libc", - "linux-raw-sys 0.4.14", + "linux-raw-sys", "windows-sys 0.52.0", ] -[[package]] -name = "rustix" -version = "1.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c71e83d6afe7ff64890ec6b71d6a69bb8a610ab78ce364b3352876bb4c801266" -dependencies = [ - "bitflags 2.13.2", - "errno", - "libc", - "linux-raw-sys 0.9.4", - "windows-sys 0.59.0", -] - [[package]] name = "rustls-pemfile" version = "2.2.0" @@ -2430,9 +2376,11 @@ name = "softnet" version = "0.1.0" dependencies = [ "anyhow", + "bytes", "clap", "coarsetime", "dhcproto", + "futures-util", "ip_network", "ipnet", "jsonrpsee-types", @@ -2441,9 +2389,7 @@ dependencies = [ "mac_address", "network-interface", "nix 0.31.3", - "num_enum", "oslog", - "polling", "prefix-trie", "privdrop", "rand", @@ -2454,6 +2400,8 @@ dependencies = [ "serial_test", "smoltcp", "system-configuration", + "tokio", + "tokio-util", "uzers", "vmnet", ] @@ -2542,7 +2490,7 @@ dependencies = [ "cfg-if", "fastrand", "once_cell", - "rustix 0.38.37", + "rustix", "windows-sys 0.59.0", ] @@ -2636,9 +2584,21 @@ dependencies = [ "pin-project-lite", "signal-hook-registry", "socket2 0.5.7", + "tokio-macros", "windows-sys 0.52.0", ] +[[package]] +name = "tokio-macros" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "tokio-native-tls" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index eef6741..e3ae862 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,10 @@ debug = true [dependencies] smoltcp = "0" libc = "0" -polling = "3" +tokio = { version = "1", features = ["rt", "sync", "net", "signal", "time", "macros", "io-util"] } +tokio-util = { version = "0.7", features = ["codec"] } +bytes = "1" +futures-util = { version = "0.3", default-features = false, features = ["std", "sink"] } dhcproto = { git = "https://github.com/bluecatengineering/dhcproto.git", branch = "master" } vmnet = "0.7.0" clap = { version = "4", features = ["derive"] } @@ -24,10 +27,9 @@ anyhow = { version = "1", features = ["backtrace"] } ip_network = "0" uzers = "0" system-configuration = "0" -num_enum = "0" sentry = { version = "0", features = ["debug-images"] } sentry-anyhow = { version = "0", features = ["backtrace"] } -nix = { version = "0", features = ["signal", "socket"] } +nix = { version = "0", features = ["socket"] } network-interface = "2" prefix-trie = "0" ipnet = "2" diff --git a/lib/host.rs b/lib/host.rs index 57aab90..255d33f 100644 --- a/lib/host.rs +++ b/lib/host.rs @@ -4,9 +4,9 @@ use clap::ValueEnum; use log::info; use smoltcp::wire::EthernetAddress; use std::net::Ipv4Addr; -use std::os::unix::io::{AsRawFd, RawFd}; -use std::os::unix::net::UnixDatagram; +use std::sync::Arc; use std::sync::mpsc::{SyncSender, sync_channel}; +use tokio::sync::Notify; use vmnet::network::Mode; use vmnet::parameters::{Parameter, ParameterKind}; use vmnet::port_forwarding::{AddressFamily, Protocol}; @@ -26,8 +26,9 @@ pub enum NetType { pub struct Host { interface: vmnet::Interface, - new_packets_rx: UnixDatagram, - callback_can_continue_tx: SyncSender<()>, + ready: Arc, + draining: bool, + callback_can_continue_tx: Option>, pub vm_ip: Ipv4Addr, pub subnet_mask: Ipv4Addr, pub gateway_ip: smoltcp::wire::Ipv4Address, @@ -98,33 +99,26 @@ impl Host { )); }; - // Set up a socketpair() to emulate polling of the vmnet interface - let (new_packets_tx, new_packets_rx) = UnixDatagram::pair()?; - new_packets_rx.set_nonblocking(true)?; - - let (callback_can_continue_tx, callback_can_continue_rx) = sync_channel(0); + // Coordinate packet readiness and callback acknowledgments + let ready = Arc::new(Notify::new()); + let ready_for_callback = Arc::clone(&ready); + let (callback_can_continue_tx, callback_can_continue_rx) = sync_channel(1); interface .set_event_callback(Events::PACKETS_AVAILABLE, move |_mask, _params| { - // Send a dummy datagram to make the other end of socketpair() readable - // and ignore the error as this merely a signalling channel to wake up - // the poller - new_packets_tx.send(&[0; 1]).ok(); - - // Wait for the permission to continue to avoid - // wasting CPU cycles or in case of termination, - // to unblock this Block[1] and allow - // vmnet.framework to terminate - // - // [1]: https://en.wikipedia.org/wiki/Blocks_(C_language_extension) - callback_can_continue_rx.recv().unwrap(); + // Wake the Tokio reader to drain the available packets + ready_for_callback.notify_one(); + + // Park this GCD worker until drained; returning early makes vmnet repeat callbacks + let _ = callback_can_continue_rx.recv(); }) .context("failed to set vmnet interface's event callback")?; Ok(Host { interface, - new_packets_rx, - callback_can_continue_tx, + ready, + draining: false, + callback_can_continue_tx: Some(callback_can_continue_tx), vm_ip, subnet_mask, gateway_ip, @@ -158,22 +152,26 @@ impl Host { .map_err(|err| anyhow!("failed to add port forwarding rule {details}: {err}")) } - pub fn read(&mut self, batch: &mut Batch, bufs: &mut [Vec]) -> vmnet::Result { - // Dequeue dummy datagram from the socket (if any) - // to free up buffer space and reduce false-positives - // when polling - let mut buf_to_be_discarded: [u8; 1] = [0; 1]; - let _ = self.new_packets_rx.recv(&mut buf_to_be_discarded); - - let result = self.interface.read_batch(batch, bufs); - - if let Err(vmnet::Error::VmnetReadNothing) = result { - // We've emptied everything, unlock the callback - // so that it will be able to pick up new events - let _ = self.callback_can_continue_tx.send(()); + pub async fn read(&mut self, batch: &mut Batch, bufs: &mut [Vec]) -> vmnet::Result { + loop { + // One callback can cover several batches; stay ready until an empty read + if !self.draining { + self.ready.notified().await; + self.draining = true; + } + + match self.interface.read_batch(batch, bufs) { + Err(vmnet::Error::VmnetReadNothing) => { + self.draining = false; + + if let Some(tx) = &self.callback_can_continue_tx { + // Buffer the acknowledgment so the reader never waits for the callback + let _ = tx.try_send(()); + } + } + result => return result, + } } - - result } pub fn write(&mut self, buf: &[u8]) -> vmnet::Result { @@ -181,21 +179,21 @@ impl Host { } pub fn finalize(&mut self) -> Result<()> { - // First make sure our callback won't be scheduled again after it finishes - self.interface - .clear_event_callback() - .context("failed to clear vmnet interface's event callback")?; + if self.finalized { + return Ok(()); + } - // Now let the callback finish - let _ = self.callback_can_continue_tx.send(()); + // First make sure our callback won't be scheduled again after it finishes + let clear_result = self.interface.clear_event_callback(); - self.interface - .finalize() - .context("failed to finalize vmnet's interface")?; + // Disconnecting releases the callback during shutdown without waiting for a receiver + self.callback_can_continue_tx.take(); - self.finalized = true; + let stop_result = self.interface.finalize(); + self.finalized = stop_result.is_ok(); - Ok(()) + clear_result.context("failed to clear vmnet interface's event callback")?; + stop_result.context("failed to finalize vmnet's interface") } } @@ -206,9 +204,3 @@ impl Drop for Host { } } } - -impl AsRawFd for Host { - fn as_raw_fd(&self) -> RawFd { - self.new_packets_rx.as_raw_fd() - } -} diff --git a/lib/mod.rs b/lib/mod.rs index 8051f75..3039ce3 100644 --- a/lib/mod.rs +++ b/lib/mod.rs @@ -1,7 +1,6 @@ mod dhcp_server; mod host; pub use host::NetType; -mod poller; pub mod proxy; mod subnet_finder; mod vm; diff --git a/lib/poller.rs b/lib/poller.rs deleted file mode 100644 index 7c7dc1b..0000000 --- a/lib/poller.rs +++ /dev/null @@ -1,124 +0,0 @@ -use anyhow::Result; -use num_enum::IntoPrimitive; -use polling::PollMode; -use polling::os::kqueue::PollerKqueueExt; -use std::os::fd::{AsRawFd, BorrowedFd}; -use std::os::unix::io::RawFd; -use std::time::Duration; - -pub struct Poller<'poller> { - poller: polling::Poller, - events: polling::Events, - timeout: Duration, - last_periodic_tick: coarsetime::Instant, - vm_fd: BorrowedFd<'poller>, - host_fd: BorrowedFd<'poller>, - control_fd: Option>, -} - -#[derive(IntoPrimitive)] -#[repr(usize)] -enum EventKey { - VM, - Host, - Control, - Interrupt, -} - -impl Poller<'_> { - pub fn new<'poller>( - vm_fd: RawFd, - host_fd: RawFd, - control_fd: Option, - timeout: Duration, - ) -> Result> { - let poller = polling::Poller::new()?; - - Ok(Poller { - poller, - events: polling::Events::new(), - timeout, - last_periodic_tick: coarsetime::Instant::recent(), - vm_fd: unsafe { BorrowedFd::borrow_raw(vm_fd) }, - host_fd: unsafe { BorrowedFd::borrow_raw(host_fd) }, - control_fd: control_fd.map(|fd| unsafe { BorrowedFd::borrow_raw(fd) }), - }) - } - - pub fn arm(&self) -> Result<()> { - unsafe { - self.poller.add_with_mode( - self.vm_fd.as_raw_fd(), - self.vm_interest(), - PollMode::Edge, - )?; - - if let Some(control_fd) = self.control_fd { - self.poller.add_with_mode( - control_fd.as_raw_fd(), - polling::Event::all(EventKey::Control.into()), - PollMode::Edge, - )?; - } - self.poller.add_with_mode( - self.host_fd.as_raw_fd(), - self.host_interest(), - PollMode::Edge, - )?; - } - - let interrupt_signal = polling::os::kqueue::Signal(libc::SIGINT); - self.poller - .add_filter(interrupt_signal, EventKey::Interrupt.into(), PollMode::Edge)?; - - Ok(()) - } - - pub fn rearm(&mut self) { - self.events.clear(); - } - - pub fn wait(&mut self) -> Result<(bool, bool, bool)> { - self.poller.wait(&mut self.events, Some(self.timeout))?; - - let vm_readable = self - .events - .iter() - .any(|ev| ev.key == Into::::into(EventKey::VM)); - let host_readable = self - .events - .iter() - .any(|ev| ev.key == Into::::into(EventKey::Host)); - let interrupt = self - .events - .iter() - .any(|ev| ev.key == Into::::into(EventKey::Interrupt)); - Ok((vm_readable, host_readable, interrupt)) - } - - pub fn periodic_tick_due(&mut self) -> bool { - let due = self.last_periodic_tick.elapsed_since_recent() >= self.timeout.into(); - - if due { - self.last_periodic_tick = coarsetime::Instant::recent(); - } - - due - } - - pub fn remove_control(&mut self) -> Result<()> { - if let Some(control_fd) = self.control_fd.take() { - self.poller.delete(control_fd)?; - } - - Ok(()) - } - - fn vm_interest(&self) -> polling::Event { - polling::Event::readable(EventKey::VM.into()) - } - - fn host_interest(&self) -> polling::Event { - polling::Event::readable(EventKey::Host.into()) - } -} diff --git a/lib/proxy/control.rs b/lib/proxy/control.rs index 219e496..b4fd440 100644 --- a/lib/proxy/control.rs +++ b/lib/proxy/control.rs @@ -1,5 +1,7 @@ use super::{Rule, Rules}; use anyhow::{Context, Result, bail}; +use bytes::BytesMut; +use futures_util::{SinkExt, StreamExt}; use jsonrpsee_types::{ ErrorObjectOwned, Id, Request, Response, ResponsePayload, error::{ @@ -7,20 +9,21 @@ use jsonrpsee_types::{ METHOD_NOT_FOUND_CODE as METHOD_NOT_FOUND, PARSE_ERROR_CODE as PARSE_ERROR, }, }; +use nix::sys::socket::{Shutdown, shutdown}; use serde::Deserialize; use serde_json::{Value, json}; use smoltcp::wire::Ipv4Address; -use std::io::{self, ErrorKind, Read, Write}; +use std::future::poll_fn; +use std::io; use std::mem::{size_of, zeroed}; -use std::net::Shutdown; use std::os::fd::{AsRawFd, FromRawFd, RawFd}; use std::os::unix::net::UnixStream; +use tokio_util::codec::{Decoder, Encoder, Framed, LinesCodec, LinesCodecError}; const MAX_REQUEST_BYTES: usize = 1024 * 1024; const MAX_PENDING_RESPONSE_BYTES: usize = 4 * MAX_REQUEST_BYTES; const MAX_RULES: usize = 4096; const MAX_IDENTIFIER_BYTES: usize = 256; -const MAX_SERVICE_BYTES: usize = MAX_REQUEST_BYTES; pub(super) struct Policy { allow: Vec, @@ -124,13 +127,7 @@ pub(super) fn normalize_rules(mut rules: Vec) -> Vec { pub(super) struct Control { policy: Policy, - stream: UnixStream, - input: Vec, - output: Vec, - output_offset: usize, - discarding_input: bool, - input_closed: bool, - policy_changed: bool, + stream: Framed, } impl Control { @@ -148,200 +145,119 @@ impl Control { Ok(Control { policy: Policy::new(gateway_ip, allow, block), - stream, - input: Vec::new(), - output: Vec::new(), - output_offset: 0, - discarding_input: false, - input_closed: false, - policy_changed: false, + stream: Framed::new( + tokio::net::UnixStream::from_std(stream)?, + ControlCodec(LinesCodec::new_with_max_length(MAX_REQUEST_BYTES)), + ), }) } - pub(super) fn service(&mut self, rules: &mut Rules) -> Result { - if !self.flush()? { - return Ok(false); - } - - if !self.output.is_empty() { - return Ok(true); - } + /// Cancellation only interrupts buffered I/O; policy changes finish in the same poll. + /// Returns None on EOF, or whether the completed request changed the policy. + pub(super) async fn service(&mut self, rules: &mut Rules) -> Result> { + // Flush pending responses before reading; Framed retains partial I/O across cancellation + self.stream.flush().await?; + poll_fn(|cx| self.stream.poll_ready_unpin(cx)).await?; - if !self.process_input(rules)? { - return Ok(false); - } + let Some(frame) = self.stream.next().await.transpose()? else { + return Ok(None); + }; - if !self.output.is_empty() { - return Ok(true); - } + self.respond(frame, rules).map(Some) + } - if self.input_closed { - return Ok(false); - } + fn respond( + &mut self, + frame: Result, + rules: &mut Rules, + ) -> Result { + // Reply to malformed frames without changing the active policy + let (response, update) = match frame { + Ok(line) => handle_request(&self.policy, rules, line.as_bytes()), + Err(error) => ( + error_response( + Id::Null, + PARSE_ERROR, + match error { + LinesCodecError::MaxLineLengthExceeded => { + "request exceeds the maximum frame size" + } + LinesCodecError::Io(_) => "invalid JSON-RPC frame", + }, + ), + None, + ), + }; - let mut buf = [0; 8192]; - let mut bytes_read = 0; + // service() reserved sink capacity; encoding checks the byte limit before any update + self.stream.start_send_unpin(response)?; - while bytes_read < MAX_SERVICE_BYTES { - match self.stream.read(&mut buf) { - Ok(0) => { - self.input_closed = true; - break; - } - Ok(n) => { - bytes_read += n; - self.input.extend_from_slice(&buf[..n]); - if !self.process_input(rules)? { - return Ok(false); - } - - if !self.output.is_empty() { - return Ok(true); - } - } - Err(err) if err.kind() == ErrorKind::WouldBlock => break, - Err(err) - if matches!( - err.kind(), - ErrorKind::BrokenPipe | ErrorKind::ConnectionReset - ) => - { - return Ok(false); - } - Err(err) => return Err(err).context("failed to read the control socket"), - } + // Keep the update in this poll with no intervening await + if let Some(update) = update + && let Some(updated_rules) = self.policy.apply(update) + { + *rules = updated_rules; + return Ok(true); } - Ok(!self.input_closed) + Ok(false) } +} - pub(super) fn shutdown(&self) -> Result<()> { - self.stream - .shutdown(Shutdown::Both) - .context("failed to shut down the control socket") - } +impl Drop for Control { + fn drop(&mut self) { + // Shut down the inherited connection even while the original descriptor remains open + let fd = self.stream.get_ref().as_raw_fd(); - /// Returns whether the policy changed and clears the change flag. - pub(super) fn policy_changed(&mut self) -> bool { - std::mem::take(&mut self.policy_changed) + // Shut down writes first because macOS can reject SHUT_RDWR after a peer half-close + let _ = shutdown(fd, Shutdown::Write); + let _ = shutdown(fd, Shutdown::Read); } +} - fn process_input(&mut self, rules: &mut Rules) -> Result { - loop { - if self.discarding_input { - if let Some(newline) = self.input.iter().position(|byte| *byte == b'\n') { - self.input.drain(..=newline); - self.discarding_input = false; - continue; - } - - self.input.clear(); - return Ok(true); - } - - let Some(newline) = self.input.iter().position(|byte| *byte == b'\n') else { - if self.input.len() > MAX_REQUEST_BYTES { - self.input.clear(); - self.discarding_input = true; - self.enqueue(error_response( - Id::Null, - PARSE_ERROR, - "request exceeds the maximum frame size", - ))?; - - return self.flush(); - } - - return Ok(true); - }; - - let line = self.input.drain(..=newline).collect::>(); - - if newline > MAX_REQUEST_BYTES { - self.enqueue(error_response( - Id::Null, - PARSE_ERROR, - "request exceeds the maximum frame size", - ))?; - - if !self.flush()? { - return Ok(false); - } - - if !self.output.is_empty() { - return Ok(true); - } - - continue; - } +struct ControlCodec(LinesCodec); - let (response, update) = handle_request(&self.policy, rules, &line[..newline]); - self.enqueue(response)?; +impl Decoder for ControlCodec { + type Item = Result; + type Error = io::Error; - if let Some(update) = update - && let Some(updated_rules) = self.policy.apply(update) - { - *rules = updated_rules; - self.policy_changed = true; - } + fn decode(&mut self, input: &mut BytesMut) -> io::Result> { + // Return malformed frames as RPC errors so later requests can still be processed + Ok(match self.0.decode(input) { + Ok(line) => line.map(Ok), + Err(error) => Some(Err(error)), + }) + } - if !self.flush()? { - return Ok(false); - } + fn decode_eof(&mut self, input: &mut BytesMut) -> io::Result> { + let frame = self.decode(input)?; - if !self.output.is_empty() { - return Ok(true); - } + // Unlike LinesCodec, discard an unterminated final request without applying it + if frame.is_none() { + input.clear(); } + + Ok(frame) } +} - fn enqueue(&mut self, response: Value) -> Result<()> { - if self.output_offset != 0 { - self.output.drain(..self.output_offset); - self.output_offset = 0; - } +impl Encoder for ControlCodec { + type Error = io::Error; - let mut response = - serde_json::to_vec(&response).context("failed to encode RPC response")?; + fn encode(&mut self, response: Value, output: &mut BytesMut) -> io::Result<()> { + let mut response = serde_json::to_vec(&response).map_err(io::Error::other)?; response.push(b'\n'); - if self.output.len() + response.len() > MAX_PENDING_RESPONSE_BYTES { - bail!("control response queue exceeded {MAX_PENDING_RESPONSE_BYTES} bytes"); + // Bound the complete encoded response before appending it + if output.len() + response.len() > MAX_PENDING_RESPONSE_BYTES { + return Err(io::Error::other(format!( + "control response queue exceeded {MAX_PENDING_RESPONSE_BYTES} bytes" + ))); } - self.output.extend(response); + output.extend_from_slice(&response); Ok(()) } - - fn flush(&mut self) -> Result { - while self.output_offset < self.output.len() { - match self.stream.write(&self.output[self.output_offset..]) { - Ok(0) => return Ok(false), - Ok(n) => self.output_offset += n, - Err(err) if err.kind() == ErrorKind::WouldBlock => return Ok(true), - Err(err) - if matches!( - err.kind(), - ErrorKind::BrokenPipe | ErrorKind::ConnectionReset - ) => - { - return Ok(false); - } - Err(err) => return Err(err).context("failed to write the control socket"), - } - } - - self.output.clear(); - self.output_offset = 0; - - Ok(true) - } -} - -impl AsRawFd for Control { - fn as_raw_fd(&self) -> RawFd { - self.stream.as_raw_fd() - } } #[derive(Deserialize)] @@ -533,18 +449,21 @@ fn validate_control_fd(control_fd: RawFd) -> Result<()> { #[cfg(test)] mod tests { use super::{ - Control, INVALID_PARAMS, INVALID_REQUEST, MAX_PENDING_RESPONSE_BYTES, MAX_REQUEST_BYTES, - MAX_RULES, METHOD_NOT_FOUND, PARSE_ERROR, Policy, handle_request, + Control, ControlCodec, INVALID_PARAMS, INVALID_REQUEST, MAX_PENDING_RESPONSE_BYTES, + MAX_REQUEST_BYTES, MAX_RULES, METHOD_NOT_FOUND, PARSE_ERROR, Policy, handle_request, }; use crate::proxy::{Direction, PolicyDecision, Rule, Rules}; + use bytes::BytesMut; + use futures_util::SinkExt; use serde_json::{Value, json}; use smoltcp::wire::Ipv4Address; use std::fs::File; use std::io::{Read, Write}; - use std::net::{Shutdown, TcpListener}; + use std::net::Shutdown; use std::os::fd::{AsRawFd, RawFd}; use std::os::unix::net::{UnixDatagram, UnixStream}; use std::time::Duration; + use tokio_util::codec::{Decoder, Encoder, LinesCodec}; struct TestPolicy { state: Policy, @@ -867,15 +786,11 @@ mod tests { assert_eq!(policy.result(), before); } - #[test] - fn newline_delimited_control_socket_handles_multiple_requests_and_eof() { + #[tokio::test(flavor = "current_thread")] + async fn newline_delimited_control_socket_handles_multiple_requests_and_eof() { let (mut client, server) = UnixStream::pair().unwrap(); - client - .set_read_timeout(Some(Duration::from_secs(1))) - .unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); - client .write_all( concat!( @@ -885,100 +800,101 @@ mod tests { .as_bytes(), ) .unwrap(); + client.shutdown(Shutdown::Write).unwrap(); - assert!(control.service(&mut rules).unwrap()); - let mut response = [0; 2048]; - let n = client.read(&mut response).unwrap(); - let lines = std::str::from_utf8(&response[..n]) - .unwrap() + assert_eq!(control.service(&mut rules).await.unwrap(), Some(false)); + assert_eq!(control.service(&mut rules).await.unwrap(), Some(true)); + assert_eq!(control.service(&mut rules).await.unwrap(), None); + assert_eq!(control.policy.allow, vec!["@host".parse().unwrap()]); + drop(control); + + let mut output = String::new(); + client.read_to_string(&mut output).unwrap(); + let responses = output .lines() .map(|line| serde_json::from_str::(line).unwrap()) .collect::>(); - assert_eq!(lines.len(), 2); - assert_eq!(lines[0]["id"], 1); - assert_eq!(lines[1]["id"], 2); - assert_eq!(lines[1]["result"]["allow"], json!(["@host"])); - assert_eq!(control.policy.allow, vec!["@host".parse().unwrap()]); - - let before = control.policy.result(rules.len()); - drop(client); - assert!(!control.service(&mut rules).unwrap()); - assert_eq!(control.policy.result(rules.len()), before); + assert_eq!(responses.len(), 2); + assert_eq!(responses[0]["id"], 1); + assert_eq!(responses[1]["id"], 2); + assert_eq!(responses[1]["result"]["allow"], json!(["@host"])); } - #[test] - fn policy_change_detection_uses_normalized_policy() { + #[tokio::test(flavor = "current_thread")] + async fn policy_change_detection_uses_normalized_policy() { let (mut client, server) = UnixStream::pair().unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); - - client - .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"out 10.1.2.3/8\"],\"block\":[\"in 10.0.0.0/8\",\"out 10.0.0.0/8\"]}}\n") - .unwrap(); - assert!(control.service(&mut rules).unwrap()); - assert!(control.policy_changed()); - - client - .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"out 10.0.0.0/8\"],\"block\":[\"in 10.0.0.0/8\",\"out 10.0.0.0/8\"]}}\n") + for (allow, changed) in [("out 10.1.2.3/8", true), ("out 10.0.0.0/8", false)] { + writeln!( + client, + "{}", + json!({ + "jsonrpc": "2.0", "id": 1, "method": "softnet.policy.set", + "params": {"allow": [allow], "block": ["in 10.0.0.0/8", "out 10.0.0.0/8"]} + }) + ) .unwrap(); - assert!(control.service(&mut rules).unwrap()); - assert!(!control.policy_changed()); - + assert_eq!(control.service(&mut rules).await.unwrap(), Some(changed)); + } client .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":3,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[],\"block\":[\"in 10.0.0.0/8\",\"out 10.0.0.0/8\"]}}\n") .unwrap(); - assert!(control.service(&mut rules).unwrap()); - assert!(control.policy_changed()); + assert_eq!(control.service(&mut rules).await.unwrap(), Some(true)); assert!(control.policy.allow.is_empty()); } - #[test] - fn write_side_eof_flushes_the_final_policy_response() { + #[tokio::test(flavor = "current_thread")] + async fn write_side_eof_flushes_the_final_policy_response() { let (mut client, server) = UnixStream::pair().unwrap(); - client - .set_read_timeout(Some(Duration::from_secs(1))) - .unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); - client .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"10.0.0.0/8\"],\"block\":[]}}\n") .unwrap(); client.shutdown(Shutdown::Write).unwrap(); - assert!(!control.service(&mut rules).unwrap()); + assert_eq!(control.service(&mut rules).await.unwrap(), Some(true)); + assert_eq!(control.service(&mut rules).await.unwrap(), None); + assert!(control.stream.write_buffer().is_empty()); + drop(control); - let mut response = [0; 1024]; - let n = client.read(&mut response).unwrap(); - let response = serde_json::from_slice::(&response[..n - 1]).unwrap(); + let mut response = String::new(); + client.read_to_string(&mut response).unwrap(); + let response = serde_json::from_str::(&response).unwrap(); assert_eq!(response["id"], 1); assert_eq!(response["result"]["allow"], json!(["10.0.0.0/8"])); - assert!(control.output.is_empty()); } - #[test] - fn oversized_frame_is_discarded_and_following_frame_is_processed() { - let (mut client, server) = UnixStream::pair().unwrap(); - client - .set_read_timeout(Some(Duration::from_secs(1))) - .unwrap(); + #[tokio::test(flavor = "current_thread")] + async fn oversized_frame_is_discarded_and_following_frame_is_processed() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + + let (client, server) = UnixStream::pair().unwrap(); + client.set_nonblocking(true).unwrap(); + let mut client = tokio::net::UnixStream::from_std(client).unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); - - control.input = vec![b'x'; MAX_REQUEST_BYTES + 1]; - control.process_input(&mut rules).unwrap(); - assert!(control.discarding_input); - - control.input.extend_from_slice( + let mut input = vec![b'x'; MAX_REQUEST_BYTES + 1]; + input.extend_from_slice( b"still-too-long\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"softnet.policy.get\"}\n", ); - control.process_input(&mut rules).unwrap(); - assert!(!control.discarding_input); - - let mut output = [0; 2048]; - let n = client.read(&mut output).unwrap(); - let responses = std::str::from_utf8(&output[..n]) - .unwrap() + let (output, ()) = tokio::join!( + async { + client.write_all(&input).await.unwrap(); + client.shutdown().await.unwrap(); + let mut output = String::new(); + client.read_to_string(&mut output).await.unwrap(); + output + }, + async { + assert_eq!(control.service(&mut rules).await.unwrap(), Some(false)); + assert_eq!(control.service(&mut rules).await.unwrap(), Some(false)); + assert_eq!(control.service(&mut rules).await.unwrap(), None); + drop(control); + } + ); + let responses = output .lines() .map(|line| serde_json::from_str::(line).unwrap()) .collect::>(); @@ -987,105 +903,112 @@ mod tests { assert_eq!(responses[1]["id"], 2); } - #[test] - fn fragmented_frame_does_not_apply_until_the_newline_arrives() { + #[tokio::test(flavor = "current_thread")] + async fn fragmented_frame_survives_cancelled_service() { let (mut client, server) = UnixStream::pair().unwrap(); - client - .set_read_timeout(Some(Duration::from_secs(1))) - .unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); let before = control.policy.result(rules.len()); - client - .write_all( - b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"@host\"],", - ) + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"@host\"],") .unwrap(); - assert!(control.service(&mut rules).unwrap()); + assert!( + tokio::time::timeout(Duration::from_millis(10), control.service(&mut rules)) + .await + .is_err() + ); assert_eq!(control.policy.result(rules.len()), before); - assert!(control.output.is_empty()); + assert!(control.stream.write_buffer().is_empty()); client.write_all(b"\"block\":[]}}\n").unwrap(); - assert!(control.service(&mut rules).unwrap()); - + assert_eq!(control.service(&mut rules).await.unwrap(), Some(true)); + control.stream.flush().await.unwrap(); let mut response = [0; 1024]; let n = client.read(&mut response).unwrap(); let response = serde_json::from_slice::(&response[..n - 1]).unwrap(); assert_eq!(response["id"], 1); assert_eq!(response["result"]["allow"], json!(["@host"])); - assert_eq!(control.policy.allow, vec!["@host".parse().unwrap()]); } #[test] - fn response_backpressure_keeps_the_pending_queue_bounded() { - let (_client, server) = UnixStream::pair().unwrap(); - let mut control = control(server.as_raw_fd()).unwrap(); - let response = json!({"jsonrpc": "2.0", "id": 1, "result": "x".repeat(MAX_REQUEST_BYTES)}); - let mut bounded = false; - - for _ in 0..8 { - match control.enqueue(response.clone()) { - Ok(()) => assert!(control.flush().unwrap()), - Err(error) => { - assert!( - error - .to_string() - .contains("control response queue exceeded") - ); - bounded = true; - break; - } - } - } - - assert!(bounded); - assert!(control.output.len() - control.output_offset <= 4 * MAX_REQUEST_BYTES); + fn codec_enforces_exact_request_and_response_limits() { + let mut codec = ControlCodec(LinesCodec::new_with_max_length(MAX_REQUEST_BYTES)); + let mut input = BytesMut::from(vec![b'x'; MAX_REQUEST_BYTES].as_slice()); + input.extend_from_slice(b"\n"); + assert_eq!( + codec.decode(&mut input).unwrap().unwrap().unwrap().len(), + MAX_REQUEST_BYTES + ); + input.resize(MAX_REQUEST_BYTES + 1, b'x'); + assert!(codec.decode(&mut input).unwrap().unwrap().is_err()); + input.extend_from_slice(b"discarded\nvalid\n"); + assert_eq!(codec.decode(&mut input).unwrap().unwrap().unwrap(), "valid"); + + let response = json!({"result": "x".repeat(MAX_REQUEST_BYTES)}); + let mut output = BytesMut::new(); + while codec.encode(response.clone(), &mut output).is_ok() {} + assert!(output.len() <= MAX_PENDING_RESPONSE_BYTES); + assert!(output.len() >= 3 * MAX_REQUEST_BYTES); } - #[test] - fn pipelined_policy_responses_stop_before_queue_overflow() { - let (_client, server) = UnixStream::pair().unwrap(); + #[tokio::test(flavor = "current_thread")] + async fn pipelined_policy_responses_stop_before_queue_overflow() { + use tokio::io::AsyncWriteExt; + + let (client, server) = UnixStream::pair().unwrap(); + client.set_nonblocking(true).unwrap(); + let mut client = tokio::net::UnixStream::from_std(client).unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); let allow = (0..MAX_RULES) .map(|index| format!("10.{}.{}.0/24", index / 256, index % 256)) .collect::>(); let mut input = serde_json::to_vec(&json!({ - "jsonrpc": "2.0", - "id": 0, - "method": "softnet.policy.set", + "jsonrpc": "2.0", "id": 0, "method": "softnet.policy.set", "params": {"allow": allow, "block": []} })) .unwrap(); input.push(b'\n'); - for id in 1..=100 { input.extend_from_slice( format!("{{\"jsonrpc\":\"2.0\",\"id\":{id},\"method\":\"softnet.policy.get\"}}\n") .as_bytes(), ); } - - control.input = input; - assert!(control.process_input(&mut rules).unwrap()); + let (written, processed) = tokio::join!(client.write_all(&input), async { + let mut processed = 0; + while let Ok(result) = + tokio::time::timeout(Duration::from_millis(50), control.service(&mut rules)).await + { + assert!(result.unwrap().is_some()); + processed += 1; + assert!(processed <= 101); + } + processed + }); + written.unwrap(); + assert!((1..101).contains(&processed)); assert_eq!(rules.len(), MAX_RULES); - assert!(!control.input.is_empty()); - assert!(!control.output.is_empty()); - assert!(control.output.len() - control.output_offset <= MAX_PENDING_RESPONSE_BYTES); + assert!(!control.stream.write_buffer().is_empty()); + assert!(control.stream.write_buffer().len() <= MAX_PENDING_RESPONSE_BYTES); } - #[test] - fn response_queue_overflow_does_not_apply_a_policy_update() { + #[tokio::test(flavor = "current_thread")] + async fn response_queue_overflow_does_not_apply_a_policy_update() { let (_client, server) = UnixStream::pair().unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); let before = control.policy.result(rules.len()); - - control.output = vec![b'x'; MAX_PENDING_RESPONSE_BYTES]; - control.input = b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"10.0.0.0/8\"],\"block\":[]}}\n".to_vec(); - - let error = control.process_input(&mut rules).unwrap_err(); + control + .stream + .write_buffer_mut() + .resize(MAX_PENDING_RESPONSE_BYTES, b'x'); + let error = control + .respond( + Ok("{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"10.0.0.0/8\"],\"block\":[]}}".into()), + &mut rules, + ) + .unwrap_err(); assert!( error .to_string() @@ -1094,39 +1017,96 @@ mod tests { assert_eq!(control.policy.result(rules.len()), before); } - #[test] - fn response_backpressure_stops_consuming_policy_updates() { + #[tokio::test(flavor = "current_thread")] + async fn backpressure_survives_cancellation_and_stops_policy_updates() { + use tokio::io::{AsyncBufReadExt, BufReader}; + let (mut client, server) = UnixStream::pair().unwrap(); let mut control = control(server.as_raw_fd()).unwrap(); let mut rules = Rules::default(); let before = control.policy.result(rules.len()); - - control.output = vec![b'x'; MAX_REQUEST_BYTES]; - assert!(control.flush().unwrap()); - assert!(!control.output.is_empty()); - + let response = json!({"jsonrpc": "2.0", "id": 0, "result": "x".repeat(MAX_REQUEST_BYTES)}); + control.stream.start_send_unpin(response.clone()).unwrap(); client .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"10.0.0.0/8\"],\"block\":[]}}\n") .unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(10), control.service(&mut rules)) + .await + .is_err() + ); + assert_eq!(control.policy.result(rules.len()), before); + assert!(control.stream.read_buffer().is_empty()); + assert!(!control.stream.write_buffer().is_empty()); + client.set_nonblocking(true).unwrap(); + let mut client = BufReader::new(tokio::net::UnixStream::from_std(client).unwrap()); + let mut line = String::new(); + let (changed, _) = tokio::join!(control.service(&mut rules), client.read_line(&mut line)); + assert_eq!(changed.unwrap(), Some(true)); + assert_eq!(serde_json::from_str::(&line).unwrap(), response); + line.clear(); + let (flushed, read) = tokio::join!(control.stream.flush(), client.read_line(&mut line)); + flushed.unwrap(); + read.unwrap(); + assert_eq!(serde_json::from_str::(&line).unwrap()["id"], 1); + } - assert!(control.service(&mut rules).unwrap()); + #[tokio::test(flavor = "current_thread")] + async fn unterminated_request_at_eof_does_not_change_policy() { + let (mut client, server) = UnixStream::pair().unwrap(); + let mut control = control(server.as_raw_fd()).unwrap(); + let mut rules = Rules::default(); + let before = control.policy.result(rules.len()); + client + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"method\":\"softnet.policy.set\",\"params\":{\"allow\":[\"@host\"],\"block\":[]}}") + .unwrap(); + client.shutdown(Shutdown::Write).unwrap(); + assert_eq!(control.service(&mut rules).await.unwrap(), None); assert_eq!(control.policy.result(rules.len()), before); - assert!(control.input.is_empty()); + assert!(control.stream.write_buffer().is_empty()); } - #[test] - fn validates_control_descriptor_without_taking_ownership() { + #[tokio::test(flavor = "current_thread")] + async fn invalid_utf8_frame_does_not_prevent_following_request() { + let (mut client, server) = UnixStream::pair().unwrap(); + let mut control = control(server.as_raw_fd()).unwrap(); + let mut rules = Rules::default(); + client + .write_all(b"\xff\n{\"jsonrpc\":\"2.0\",\"id\":2,\"method\":\"softnet.policy.get\"}\n") + .unwrap(); + client.shutdown(Shutdown::Write).unwrap(); + assert_eq!(control.service(&mut rules).await.unwrap(), Some(false)); + assert_eq!(control.service(&mut rules).await.unwrap(), Some(false)); + assert_eq!(control.service(&mut rules).await.unwrap(), None); + drop(control); + let mut output = String::new(); + client.read_to_string(&mut output).unwrap(); + let responses = output + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .collect::>(); + assert_eq!(responses.len(), 2); + assert_eq!(responses[0]["error"]["code"], PARSE_ERROR); + assert_eq!(responses[1]["id"], 2); + } + + #[tokio::test(flavor = "current_thread")] + async fn validates_control_descriptor_without_taking_ownership() { let file = File::open("/dev/null").unwrap(); let error = control(file.as_raw_fd()).err().unwrap(); assert!(error.to_string().contains("is not a socket")); assert!(file.metadata().is_ok()); - let (datagram, _) = UnixDatagram::pair().unwrap(); let error = control(datagram.as_raw_fd()).err().unwrap(); assert!(error.to_string().contains("not a Unix stream socket")); - - let listener = TcpListener::bind("127.0.0.1:0").unwrap(); - let error = control(listener.as_raw_fd()).err().unwrap(); + let internet = nix::sys::socket::socket( + nix::sys::socket::AddressFamily::Inet, + nix::sys::socket::SockType::Stream, + nix::sys::socket::SockFlag::empty(), + None, + ) + .unwrap(); + let error = control(internet.as_raw_fd()).err().unwrap(); assert!(error.to_string().contains("not a Unix socket")); let (stream, _peer) = UnixStream::pair().unwrap(); @@ -1135,17 +1115,11 @@ mod tests { assert!(unsafe { libc::fcntl(stream.as_raw_fd(), libc::F_GETFD) != -1 }); } - #[test] - fn shutdown_signals_eof_while_the_original_descriptor_remains_open() { + #[tokio::test(flavor = "current_thread")] + async fn dropping_control_signals_eof_while_the_original_descriptor_remains_open() { let (mut client, server) = UnixStream::pair().unwrap(); - client - .set_read_timeout(Some(Duration::from_secs(1))) - .unwrap(); let control = control(server.as_raw_fd()).unwrap(); - - control.shutdown().unwrap(); drop(control); - assert!(unsafe { libc::fcntl(server.as_raw_fd(), libc::F_GETFD) != -1 }); let mut response = [0; 1]; assert_eq!(client.read(&mut response).unwrap(), 0); diff --git a/lib/proxy/host.rs b/lib/proxy/host.rs index 7c6e1f6..0e8f88d 100644 --- a/lib/proxy/host.rs +++ b/lib/proxy/host.rs @@ -5,7 +5,7 @@ use dhcproto::v4::{CLIENT_PORT, SERVER_PORT}; use smoltcp::phy::ChecksumCapabilities; use smoltcp::wire::{EthernetFrame, EthernetProtocol, IpProtocol, Ipv4Packet, Ipv4Repr, UdpPacket}; -impl Proxy<'_> { +impl Proxy { pub(crate) fn process_frame_from_host(&mut self, frame: &EthernetFrame<&[u8]>) -> Result<()> { if self.allowed_from_host(frame).is_none() { // Block packet by not forwarding it to the VM diff --git a/lib/proxy/mod.rs b/lib/proxy/mod.rs index 62c2d30..8c1ef59 100644 --- a/lib/proxy/mod.rs +++ b/lib/proxy/mod.rs @@ -9,26 +9,25 @@ mod vm; use crate::dhcp_server::{DHCP_SERVER_MAC, DhcpServer}; use crate::host::Host; use crate::host::NetType; -use crate::poller::Poller; use crate::vm::VM; use anyhow::{Context, Result, bail}; use control::{Control, normalize_rules}; pub use exposed_port::ExposedPort; use flows::{FlowTable, PendingFlow}; +use futures_util::future::OptionFuture; use ipnet::Ipv4Net; use mac_address::MacAddress; pub use rule::{Direction, Rule, Target}; pub(crate) use rules::{PolicyDecision, Rules}; use smoltcp::wire::{EthernetAddress, EthernetFrame, Ipv4Address}; -use std::io::ErrorKind; -use std::os::unix::io::{AsRawFd, RawFd}; +use std::os::unix::io::RawFd; use std::time::Duration; +use tokio::time::{MissedTickBehavior, interval}; use vmnet::Batch; -pub struct Proxy<'proxy> { +pub struct Proxy { vm: VM, host: Host, - poller: Poller<'proxy>, vm_mac_address: smoltcp::wire::EthernetAddress, dhcp_server: DhcpServer, rules: Rules, @@ -37,8 +36,9 @@ pub struct Proxy<'proxy> { enobufs_encountered: bool, } -impl Proxy<'_> { - pub fn new<'proxy>( +impl Proxy { + /// Requires an entered Tokio runtime during construction. + pub fn new( vm_fd: RawFd, vm_mac_address: MacAddress, vm_net_type: NetType, @@ -46,7 +46,7 @@ impl Proxy<'_> { block: Vec, exposed_ports: Vec, control_fd: Option, - ) -> Result> { + ) -> Result { // Ensure that VM's MAC address won't conflict with our DHCP server's MAC address let vm_mac_address = EthernetAddress(vm_mac_address.bytes()); if vm_mac_address == DHCP_SERVER_MAC { @@ -62,18 +62,11 @@ impl Proxy<'_> { vm_mac_address, !allow.contains(&Rule::Stateless(Target::Prefix(Ipv4Net::default()))), )?; - let poller_timeout = Duration::from_millis(100); let control = control_fd .map(|control_fd| { Control::new(control_fd, host.gateway_ip, allow.clone(), block.clone()) }) .transpose()?; - let poller = Poller::new( - vm.as_raw_fd(), - host.as_raw_fd(), - control.as_ref().map(AsRawFd::as_raw_fd), - poller_timeout, - )?; let rules = Rules::new(host.gateway_ip, &allow, &block); @@ -99,7 +92,6 @@ impl Proxy<'_> { Ok(Proxy { vm, host, - poller, vm_mac_address, dhcp_server, rules, @@ -109,7 +101,7 @@ impl Proxy<'_> { }) } - pub fn run(&mut self) -> Result<()> { + pub async fn run(&mut self) -> Result<()> { // Create a single buffer from reading from the VM let mut buf: Vec = vec![0; self.host.max_packet_size as usize]; @@ -120,129 +112,93 @@ impl Proxy<'_> { ]; let mut batch = Batch::preallocate(bufs.len()); - self.poller.arm()?; + // Check VM connectivity periodically, skipping missed ticks + let mut maintenance = interval(Duration::from_millis(100)); + maintenance.set_missed_tick_behavior(MissedTickBehavior::Skip); + // Service packet I/O, policy changes, and maintenance on the same runtime loop { - let (vm_readable, host_readable, interrupt) = self.poller.wait()?; - - // Update coarse time for flows - coarsetime::Instant::update(); - - // Service control on every wake (including timeouts) so a bounded read or a pending - // response continues making progress even when no new edge is generated. - self.service_control(); - - if vm_readable { - self.read_from_vm(buf.as_mut_slice())?; - } - - if host_readable { - self.read_from_host(&mut batch, &mut bufs)?; - } + tokio::select! { + result = self.vm.read(&mut buf) => { + let n = result?; - // Graceful termination - if interrupt { - return Ok(()); - } - - // Periodic maintenance - if self.poller.periodic_tick_due() { - // kqueue(2) does not report peer disconnects for Unix datagram sockets - if !self.vm.is_connected()? { - return Ok(()); - } - } - - self.poller.rearm(); - } - } - - fn read_from_vm(&mut self, buf: &mut [u8]) -> Result<()> { - let mut packets_read = 0; - - loop { - match self.vm.read(buf) { - Ok(n) => { - // Update coarse time for flows + // Update the coarse time once the packet arrives coarsetime::Instant::update(); if let Ok(frame) = EthernetFrame::new_checked(&buf[..n]) { self.process_frame_from_vm(frame)?; } - packets_read += 1; - if packets_read == 128 { - self.service_control(); - packets_read = 0; + // Limit VM bursts to 128 packets (including the one we read above), + // so that other select branches will have a chance to run + for _ in 0..127 { + match self.vm.try_read(&mut buf) { + Ok(n) => { + coarsetime::Instant::update(); + + if let Ok(frame) = EthernetFrame::new_checked(&buf[..n]) { + self.process_frame_from_vm(frame)?; + } + + tokio::task::consume_budget().await; + } + Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => break, + Err(err) => return Err(err.into()), + } } } - Err(err) => { - if err.kind() == ErrorKind::WouldBlock { - return Ok(()); - } - return Err(err.into()); - } - } - } - } + result = self.host.read(&mut batch, &mut bufs) => { + let pktcnt = result?; - fn read_from_host(&mut self, batch: &mut Batch, bufs: &mut [Vec]) -> Result<()> { - loop { - match self.host.read(batch, bufs) { - Ok(pktcnt) => { - // Update coarse time for flows + // Update the coarse time once the packet arrives coarsetime::Instant::update(); - for buf in batch.packet_sized_bufs(bufs).take(pktcnt) { + for buf in batch.packet_sized_bufs(&bufs).take(pktcnt) { if let Ok(pkt) = EthernetFrame::new_checked(buf) { self.process_frame_from_host(&pkt)?; } - } - self.service_control(); - } - Err(err) => { - if let vmnet::Error::VmnetReadNothing = err { - return Ok(()); + tokio::task::consume_budget().await; } - - return Err(err.into()); } - } - } - } - - fn service_control(&mut self) { - let Some(control) = self.control.as_mut() else { - return; - }; - - let keep_open = match control.service(&mut self.rules) { - Ok(keep_open) => keep_open, - Err(err) => { - log::warn!("disabling Softnet control socket: {err:#}"); - false - } - }; - // Invalidate tracked flows whenever the policy changes - if control.policy_changed() { - self.flows = self.rules.has_stateful().then(FlowTable::new); - } + // Skip this branch when no control socket is present + Some(result) = OptionFuture::from( + self.control + .as_mut() + .map(|control| control.service(&mut self.rules)) + ) => { + match result { + Ok(Some(true)) => { + // Invalidate tracked flows whenever the policy changes + self.flows = self.rules.has_stateful().then(FlowTable::new); + } + Ok(Some(false)) => { + // Policy unchanged; keep tracked flows intact + } + Ok(None) => { + // Control stream reached EOF; stop polling it + self.control = None; + } + Err(err) => { + // Control service failed; stop polling it + self.control = None; - if keep_open { - return; - } + log::warn!("disabling Softnet control socket: {err:#}"); + } + } - if let Err(err) = self.poller.remove_control() { - log::warn!("failed to remove Softnet control socket from the poller: {err:#}"); - } + coarsetime::Instant::update(); + } - if let Some(control) = self.control.take() - && let Err(err) = control.shutdown() - { - log::warn!("failed to shut down Softnet control socket: {err:#}"); + _ = maintenance.tick() => { + // Unix datagram sockets do not report peer closure as readable + if !self.vm.is_connected()? { + return Ok(()); + } + } + } } } @@ -285,9 +241,9 @@ mod tests { use std::os::fd::AsRawFd; use std::str::FromStr; - #[test] + #[tokio::test] #[serial] - fn test_blocking_takes_precedence() { + async fn test_blocking_takes_precedence() { let mut proxy = create_proxy(vec!["66.66.0.0/16"], vec!["66.66.0.0/16"]); let vm_ip = proxy.host.vm_ip; @@ -296,9 +252,9 @@ mod tests { assert!(allowed_from_vm_ipv4(&mut proxy, vm_ip, "66.66.66.66").is_none()); } - #[test] + #[tokio::test] #[serial] - fn test_longest_prefix_match_wins() { + async fn test_longest_prefix_match_wins() { let mut proxy = create_proxy(vec!["33.33.33.33/32"], vec!["33.33.33.0/24"]); let vm_ip = proxy.host.vm_ip; @@ -309,9 +265,9 @@ mod tests { assert!(allowed_from_vm_ipv4(&mut proxy, vm_ip, "33.33.33.34").is_none()); } - #[test] + #[tokio::test] #[serial] - fn test_allow_host() { + async fn test_allow_host() { let mut proxy = create_proxy(vec!["@host"], vec!["0.0.0.0/0"]); let vm_ip = proxy.host.vm_ip; @@ -325,9 +281,9 @@ mod tests { assert!(allowed_from_vm_ipv4(&mut proxy, vm_ip, &gateway_ip).is_some()); } - #[test] + #[tokio::test] #[serial] - fn test_bare_default_block_applies_in_both_directions_in_stateful_mode() { + async fn test_bare_default_block_applies_in_both_directions_in_stateful_mode() { let mut proxy = create_proxy(vec!["in 192.0.2.0/24"], vec!["0.0.0.0/0"]); let vm_ip = proxy.host.vm_ip; let fallback_peer = Ipv4Address::new(203, 0, 113, 1); @@ -350,9 +306,9 @@ mod tests { assert!(proxy.allowed_from_vm_ipv4(allowed_reply).is_some()); } - #[test] + #[tokio::test] #[serial] - fn test_directional_egress_block_does_not_block_reply_to_unmatched_inbound_flow() { + async fn test_directional_egress_block_does_not_block_reply_to_unmatched_inbound_flow() { let mut proxy = create_proxy(vec![], vec!["out 203.0.113.0/24"]); let vm_ip = proxy.host.vm_ip; let peer = Ipv4Address::new(203, 0, 113, 1); @@ -366,9 +322,9 @@ mod tests { assert!(proxy.allowed_from_vm_ipv4(reply).is_some()); } - #[test] + #[tokio::test] #[serial] - fn test_directional_ingress_block_does_not_block_reply_to_bare_outbound_allow() { + async fn test_directional_ingress_block_does_not_block_reply_to_bare_outbound_allow() { let mut proxy = create_proxy(vec!["203.0.113.0/24"], vec!["in 203.0.113.0/24"]); let vm_ip = proxy.host.vm_ip; let peer = Ipv4Address::new(203, 0, 113, 1); @@ -382,7 +338,7 @@ mod tests { assert!(proxy.allowed_from_host_ipv4(&reply).is_some()); } - fn create_proxy<'test>(allow: Vec<&str>, block: Vec<&str>) -> Proxy<'test> { + fn create_proxy(allow: Vec<&str>, block: Vec<&str>) -> Proxy { let (vm_fd, _) = socketpair( AddressFamily::Unix, SockType::Datagram, diff --git a/lib/proxy/vm.rs b/lib/proxy/vm.rs index eb2b5c0..a1f2737 100644 --- a/lib/proxy/vm.rs +++ b/lib/proxy/vm.rs @@ -11,7 +11,7 @@ use smoltcp::wire::{ const IPV4_HEADER_LEN_WITHOUT_OPTIONS: u8 = 20; -impl Proxy<'_> { +impl Proxy { pub(crate) fn process_frame_from_vm(&mut self, frame: EthernetFrame<&[u8]>) -> Result<()> { if self.allowed_from_vm(&frame).is_none() { // Block packet by not forwarding it to the host diff --git a/lib/vm.rs b/lib/vm.rs index cadab07..d535472 100644 --- a/lib/vm.rs +++ b/lib/vm.rs @@ -1,11 +1,13 @@ use anyhow::{Context, Result, bail}; use std::io; use std::mem::{size_of, zeroed}; -use std::os::fd::{AsRawFd, FromRawFd, RawFd}; +use std::os::fd::{FromRawFd, RawFd}; use std::os::unix::net::UnixDatagram; +use tokio::io::Interest; +use tokio::io::unix::AsyncFd; pub struct VM { - sock: UnixDatagram, + sock: AsyncFd, } impl VM { @@ -16,19 +18,27 @@ impl VM { let sock = unsafe { UnixDatagram::from_raw_fd(vm_fd) }; sock.set_nonblocking(true)?; - Ok(VM { sock }) + Ok(VM { + sock: AsyncFd::new(sock)?, + }) } pub fn write(&self, pkt: &[u8]) -> std::io::Result { - self.sock.send(pkt) + self.sock.get_ref().send(pkt) } - pub fn read(&self, buf: &mut [u8]) -> std::io::Result { - self.sock.recv(buf) + pub fn try_read(&self, buf: &mut [u8]) -> io::Result { + self.sock.get_ref().recv(buf) + } + + pub async fn read(&self, buf: &mut [u8]) -> io::Result { + self.sock + .async_io(Interest::READABLE, |sock| sock.recv(buf)) + .await } pub fn is_connected(&self) -> io::Result { - match self.sock.peer_addr() { + match self.sock.get_ref().peer_addr() { Ok(_) => Ok(true), Err(error) if error.kind() == io::ErrorKind::NotConnected => Ok(false), Err(error) => Err(error), @@ -114,16 +124,9 @@ fn validate_vm_fd(vm_fd: RawFd) -> Result<()> { Ok(()) } -impl AsRawFd for VM { - fn as_raw_fd(&self) -> RawFd { - self.sock.as_raw_fd() - } -} - #[cfg(test)] mod tests { use super::VM; - use polling::{Event, Events, PollMode, Poller}; use std::fs::File; use std::net::UdpSocket; use std::os::fd::AsRawFd; @@ -180,8 +183,8 @@ mod tests { assert!(error.to_string().contains("not a Unix socket")); } - #[test] - fn test_new_does_not_close_original_fd_when_vm_is_dropped() { + #[tokio::test] + async fn test_new_does_not_close_original_fd_when_vm_is_dropped() { let (socket, _peer) = UnixDatagram::pair().unwrap(); let vm = VM::new(socket.as_raw_fd()).unwrap(); drop(vm); @@ -198,61 +201,52 @@ mod tests { assert!(socket_fd_is_open); } - #[test] - fn test_connected_socket_has_peer() { + #[tokio::test] + async fn test_connected_socket_has_peer() { let (socket, _peer) = UnixDatagram::pair().unwrap(); let vm = VM::new(socket.as_raw_fd()).unwrap(); assert!(vm.is_connected().unwrap()); } - #[test] - fn test_disconnected_peer_is_detected_without_kqueue_event() { + #[tokio::test] + async fn test_disconnected_peer_is_detected_without_readiness() { let (socket, peer) = UnixDatagram::pair().unwrap(); let vm = VM::new(socket.as_raw_fd()).unwrap(); - let poller = Poller::new().unwrap(); - let mut events = Events::new(); - - unsafe { - poller - .add_with_mode(vm.as_raw_fd(), Event::readable(0), PollMode::Edge) - .unwrap(); - } + let mut buf = [0]; drop(peer); - poller - .wait(&mut events, Some(Duration::from_millis(20))) - .unwrap(); - - assert!(events.is_empty()); + assert!( + tokio::time::timeout(Duration::from_millis(20), vm.read(&mut buf)) + .await + .is_err() + ); assert!(!vm.is_connected().unwrap()); } - #[test] - fn test_disconnected_peer_is_detected_when_another_socket_wakes_kqueue() { + #[tokio::test] + async fn test_disconnected_peer_is_detected_when_another_socket_is_readable() { let (socket, peer) = UnixDatagram::pair().unwrap(); let (host, host_peer) = UnixDatagram::pair().unwrap(); let vm = VM::new(socket.as_raw_fd()).unwrap(); - let poller = Poller::new().unwrap(); - let mut events = Events::new(); - - unsafe { - poller - .add_with_mode(vm.as_raw_fd(), Event::readable(0), PollMode::Edge) - .unwrap(); - poller - .add_with_mode(host.as_raw_fd(), Event::readable(1), PollMode::Edge) - .unwrap(); - } + host.set_nonblocking(true).unwrap(); + let host = tokio::net::UnixDatagram::from_std(host).unwrap(); + let mut buf = [0]; drop(peer); host_peer.send(&[1]).unwrap(); - poller - .wait(&mut events, Some(Duration::from_millis(20))) - .unwrap(); - - assert!(events.iter().any(|event| event.key == 1)); - assert!(!events.iter().any(|event| event.key == 0)); + assert_eq!( + tokio::time::timeout(Duration::from_secs(2), host.recv(&mut buf)) + .await + .unwrap() + .unwrap(), + 1 + ); + assert!( + tokio::time::timeout(Duration::from_millis(20), vm.read(&mut buf)) + .await + .is_err() + ); assert!(!vm.is_connected().unwrap()); } } diff --git a/src/main.rs b/src/main.rs index a5324b6..43bab8a 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,7 +1,6 @@ use anyhow::{Context, anyhow}; use clap::Parser; use log::LevelFilter; -use nix::sys::signal::{SigHandler, Signal, signal}; use oslog::OsLogger; use privdrop::PrivDrop; use softnet::NetType; @@ -23,6 +22,7 @@ use system_configuration::sys::preferences::{ SCPreferencesApplyChanges, SCPreferencesCommitChanges, SCPreferencesLock, SCPreferencesSetValue, SCPreferencesUnlock, }; +use tokio::signal::unix::{SignalKind, signal}; use uzers::{get_current_groupname, get_current_username, get_effective_uid}; const VERSION: &str = match option_env!("SOFTNET_VERSION") { @@ -168,20 +168,13 @@ fn main() -> ExitCode { } } -fn try_main() -> anyhow::Result<()> { +#[tokio::main(flavor = "current_thread")] +async fn try_main() -> anyhow::Result<()> { // Initialize logger OsLogger::new("org.cirruslabs.softnet") .level_filter(LevelFilter::Info) .init()?; - // The default signal(3)[1] action for SIGINT is to interrupt program, - // but we want to handle SIGINT ourselves, so we ignore it. The kqueue(2)'s[2] - // EVFILT_SIGNAL will receive it anyways, because it has lower precedence. - // - // [1]: https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man3/signal.3.html - // [2]: https://developer.apple.com/library/archive/documentation/System/Conceptual/ManPages_iPhoneOS/man2/kqueue.2.html - unsafe { signal(Signal::SIGINT, SigHandler::SigIgn) }?; - let args: Args = Args::parse(); // No need to run anything, just return @@ -248,8 +241,13 @@ fn try_main() -> anyhow::Result<()> { .apply() .context("failed to drop privileges")?; - // Run proxy - proxy.run() + // Run the proxy until it finishes or the process receives SIGINT + let mut interrupt = signal(SignalKind::interrupt())?; + + tokio::select! { + result = proxy.run() => result, + _ = interrupt.recv() => Ok(()), + } } fn parse_vm_fd(value: &str) -> Result {