From eea3b94554cf55ce3e8ac5a72fde2cfb48942a34 Mon Sep 17 00:00:00 2001 From: Changyuan Lyu Date: Wed, 22 Jul 2026 23:26:50 -0700 Subject: [PATCH] fix(vsock): drain host data before reporting EOF Keep a host-initiated UDS connection alive after the peer closes until its buffered data has been delivered to guest RX buffers. Report the close with a VSOCK RST only after the data is drained, and resume pending reads when RX descriptors become available. Signed-off-by: Changyuan Lyu --- alioth/src/virtio/dev/vsock/uds_vsock.rs | 109 +++++++- alioth/src/virtio/dev/vsock/uds_vsock_test.rs | 260 ++++++++++++++++++ alioth/src/virtio/queue/queue.rs | 4 + 3 files changed, 364 insertions(+), 9 deletions(-) diff --git a/alioth/src/virtio/dev/vsock/uds_vsock.rs b/alioth/src/virtio/dev/vsock/uds_vsock.rs index 7b9c41ec..336e35d6 100644 --- a/alioth/src/virtio/dev/vsock/uds_vsock.rs +++ b/alioth/src/virtio/dev/vsock/uds_vsock.rs @@ -161,6 +161,7 @@ impl UdsVsock { reader, writer: BufWriter::new(writer), buf_alloc: buf_size as u32, + eof: false, }; self.connections.insert((host_port, port), conn); let count = self.host_ports.entry(host_port).or_default(); @@ -228,6 +229,7 @@ impl UdsVsock { fn handle_tx_response<'m, Q, S>( &mut self, hdr: &VsockHeader, + registry: &Registry, rx_q: &mut Queue<'_, 'm, Q>, irq_sender: &S, ) -> Result<()> @@ -262,7 +264,7 @@ impl UdsVsock { "{}: host:{host_port} -> vm:{guest_port}: established", self.name ); - self.transfer_rx_data(host_port, guest_port, rx_q, irq_sender) + self.process_rx_data(host_port, guest_port, registry, rx_q, irq_sender) } fn remove_conn(&mut self, host_port: u32, guest_port: u32, registry: &Registry) -> Result<()> { @@ -384,6 +386,7 @@ impl UdsVsock { reader: BufReader::new(reader), writer: BufWriter::new(writer), buf_alloc: buf_size as u32, + eof: false, state: ConnState::Established { fwd_cnt: Wrapping(0), }, @@ -451,7 +454,7 @@ impl UdsVsock { ); match hdr.op { VsockOp::REQUEST => self.handle_tx_request(hdr, registry, irq_sender, rx_q), - VsockOp::RESPONSE => self.handle_tx_response(hdr, rx_q, irq_sender), + VsockOp::RESPONSE => self.handle_tx_response(hdr, registry, rx_q, irq_sender), VsockOp::RST => self.handle_tx_rst(hdr, registry), VsockOp::RW => self.transfer_tx_data(hdr, body, readable), VsockOp::CREDIT_UPDATE => { @@ -506,9 +509,10 @@ impl UdsVsock { hdr: &mut VsockHeader, conn: &mut BufReader, buffers: &mut [IoSliceMut], - ) -> Result { + ) -> Result<(usize, bool)> { let mut nskip = 0; let mut nread = 0; + let mut eof = false; for buf in buffers.iter_mut() { let r = if HEADER_SIZE > nskip { let Some((_, data)) = buf.split_at_mut_checked(HEADER_SIZE - nskip) else { @@ -524,7 +528,10 @@ impl UdsVsock { conn.read(buf) }; let n = match r { - Ok(0) => break, + Ok(0) => { + eof = true; + break; + } Ok(n) => n, Err(e) if e.kind() == ErrorKind::WouldBlock => break, Err(e) => Err(e)?, @@ -537,7 +544,7 @@ impl UdsVsock { hdr.len = nread as u32; let mut hdr_buf = hdr.as_bytes(); let _ = hdr_buf.read_vectored(buffers); - Ok(nread) + Ok((nread, eof)) } let rx_idx = VsockVirtq::RX.raw(); @@ -548,6 +555,9 @@ impl UdsVsock { ); return Ok(()); }; + if conn.eof { + return Ok(()); + } let ConnState::Established { fwd_cnt } = conn.state else { log::error!("{}: unexpected state {:?}", self.name, conn.state); return Ok(()); @@ -564,7 +574,11 @@ impl UdsVsock { ..Default::default() }; rx_q.handle_desc(rx_idx, irq_sender, |desc| { - let nread = copy_to_rx(&mut hdr, &mut conn.reader, &mut desc.writable)? as u32; + if conn.eof { + return Ok(Status::Break); + } + let (nread, read_eof) = copy_to_rx(&mut hdr, &mut conn.reader, &mut desc.writable)?; + conn.eof |= read_eof; if nread == 0 { return Ok(Status::Break); } @@ -573,12 +587,76 @@ impl UdsVsock { self.name ); Ok(Status::Done { - len: nread + HEADER_SIZE as u32, + len: (nread + HEADER_SIZE) as u32, }) })?; Ok(()) } + fn process_rx_data<'m, Q, S>( + &mut self, + host_port: u32, + guest_port: u32, + registry: &Registry, + rx_q: &mut Queue<'_, 'm, Q>, + irq_sender: &S, + ) -> Result<()> + where + Q: VirtQueue<'m>, + S: IrqSender, + { + self.transfer_rx_data(host_port, guest_port, rx_q, irq_sender)?; + + let eof = self + .connections + .get(&(host_port, guest_port)) + .is_some_and(|conn| conn.eof); + if eof && rx_q.desc_avail() { + let hdr = VsockHeader { + src_cid: self.config.guest_cid, + dst_cid: VSOCK_CID_HOST, + src_port: guest_port, + dst_port: host_port, + type_: SOCKET_TYPE, + ..Default::default() + }; + self.respond_rst(&hdr, irq_sender, rx_q)?; + self.remove_conn(host_port, guest_port, registry)?; + } + Ok(()) + } + + fn flush_rx_data<'m, Q, S>( + &mut self, + registry: &Registry, + rx_q: &mut Queue<'_, 'm, Q>, + irq_sender: &S, + ) -> Result<()> + where + Q: VirtQueue<'m>, + S: IrqSender, + { + if self.connections.is_empty() { + return Ok(()); + } + let mut ports: Vec<_> = self + .connections + .iter() + .map(|(ports, conn)| (*ports, conn.eof)) + .collect(); + // Sort connections to process those with EOF first. + // !eof maps true to false, and false to true. Since false < true, + // this puts eof=true connections at the front of the list. + ports.sort_by_key(|(_, eof)| !eof); + for ((host_port, guest_port), _) in ports { + if !rx_q.desc_avail() { + break; + } + self.process_rx_data(host_port, guest_port, registry, rx_q, irq_sender)?; + } + Ok(()) + } + fn transfer_tx_data( &mut self, hdr: &VsockHeader, @@ -665,6 +743,7 @@ pub struct Connection { reader: BufReader, writer: BufWriter, buf_alloc: u32, + eof: bool, } impl UdsVsock { @@ -770,7 +849,7 @@ impl VirtioMio for UdsVsock { self.handle_conn_request(token, socket, rx_q, irq_sender) } else if let Some(port_pair) = self.ports.get(&token) { let (host_port, guest_port) = port_pair.to_owned(); - self.transfer_rx_data(host_port, guest_port, rx_q, irq_sender) + self.process_rx_data(host_port, guest_port, registry, rx_q, irq_sender) } else { log::error!("{}: invalid token: {token:#x?}", self.name); Ok(()) @@ -791,7 +870,19 @@ impl VirtioMio for UdsVsock { let name = &self.name; match index { VsockVirtq::TX => self.handle_tx(active_mio)?, - VsockVirtq::RX => log::debug!("{name}: queue RX buffer available"), + VsockVirtq::RX => { + log::debug!("{name}: queue RX buffer available"); + let registry = active_mio.poll.registry(); + let irq_sender = active_mio.irq_sender; + let Some(Some(rx_q)) = active_mio.queues.get_mut(VsockVirtq::RX.raw() as usize) + else { + return error::InvalidQueueIndex { + index: VsockVirtq::RX.raw(), + } + .fail(); + }; + self.flush_rx_data(registry, rx_q, irq_sender)?; + } VsockVirtq::EVENT => log::debug!("{name}: queue EVENT buffer available"), _ => log::error!("{name}: unknown queue index {index:?}"), } diff --git a/alioth/src/virtio/dev/vsock/uds_vsock_test.rs b/alioth/src/virtio/dev/vsock/uds_vsock_test.rs index fa2a9df0..60d360a5 100644 --- a/alioth/src/virtio/dev/vsock/uds_vsock_test.rs +++ b/alioth/src/virtio/dev/vsock/uds_vsock_test.rs @@ -16,6 +16,7 @@ use std::io::{BufRead, BufReader, ErrorKind, Read, Write}; use std::mem::size_of; use std::os::unix::net::{UnixListener, UnixStream}; use std::sync::Arc; +use std::thread; use std::time::Duration; use assert_matches::assert_matches; @@ -386,3 +387,262 @@ fn vsock_conn_test() { notifier.notify().unwrap(); handle.join().unwrap(); } + +#[test] +fn vsock_host_close_test() { + let ram_bus = Arc::new(fixture_ram_bus()); + let ram = ram_bus.lock_layout(); + let regs: Arc<[QueueReg]> = Arc::from(fixture_queues(3)); + let reg_tx = ®s[VsockVirtq::TX.raw() as usize]; + let reg_rx = ®s[VsockVirtq::RX.raw() as usize]; + let mut rx_q = GuestQueue::new( + SplitQueue::new(reg_rx, &ram, false).unwrap().unwrap(), + reg_rx, + ); + let mut tx_q = GuestQueue::new( + SplitQueue::new(reg_tx, &ram, false).unwrap().unwrap(), + reg_tx, + ); + + let temp_dir = TempDir::new().unwrap(); + let sock_path = temp_dir.path().join("vsock.sock"); + + const GUEST_CID: u32 = 3; + let param = UdsVsockSpec { + cid: GUEST_CID, + path: sock_path.clone().into(), + }; + let dev = param.build("vsock").unwrap(); + + let (tx, rx) = flume::unbounded(); + let (handle, notifier) = dev.spawn_worker(rx, ram_bus.clone(), regs).unwrap(); + let (irq_tx, irq_rx) = flume::unbounded(); + let irq_sender = Arc::new(FakeIrqSender { q_tx: irq_tx }); + let start_param = StartParam { + feature: VirtioFeature::VERSION_1.bits(), + irq_sender, + ioeventfds: Option::>::None, + }; + tx.send(WakeEvent::Start { param: start_param }).unwrap(); + + let rx_buf_addr = DATA_ADDR; + let tx_buf_addr = DATA_ADDR + 4096; + + // Establish a host-initiated connection + let mut h2g_stream = UnixStream::connect(&sock_path).unwrap(); + h2g_stream.set_nonblocking(true).unwrap(); + + let buf_id = rx_q.add_desc(&[], &[(rx_buf_addr, 4096)]); + const H2G_GUEST_PORT: u32 = 1025; + writeln!(h2g_stream, "CONNECT {H2G_GUEST_PORT}").unwrap(); + assert_eq!( + irq_rx.recv_timeout(Duration::from_secs(1)).unwrap(), + VsockVirtq::RX.raw() + ); + let used = rx_q.get_used().unwrap(); + assert_eq!(used.id, buf_id); + assert_eq!(used.len as usize, size_of::()); + + let mut hdr = VsockHeader::new_zeroed(); + ram.read(rx_buf_addr, hdr.as_mut_bytes()).unwrap(); + assert_eq!(hdr.op, VsockOp::REQUEST); + let h2g_host_port = hdr.src_port; + + let resp_hdr = VsockHeader { + src_cid: GUEST_CID, + dst_cid: VSOCK_CID_HOST, + src_port: H2G_GUEST_PORT, + dst_port: h2g_host_port, + op: VsockOp::RESPONSE, + type_: VsockType::STREAM, + ..Default::default() + }; + send_to_tx( + &resp_hdr, + &[], + &ram, + tx_buf_addr, + &mut tx_q, + &tx, + ¬ifier, + &irq_rx, + false, + ); + let mut reader = BufReader::new(&h2g_stream); + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + assert_eq!(line, format!("OK {h2g_host_port}\n")); + + // Provide RX descriptor first, then close host socket + let rx_buf_id = rx_q.add_desc(&[], &[(rx_buf_addr, 4096)]); + tx.send(WakeEvent::Notify { + q_index: VsockVirtq::RX.raw(), + }) + .unwrap(); + notifier.notify().unwrap(); + + drop(h2g_stream); // EOF to alioth + + // Verify guest receives RST + assert_eq!( + irq_rx.recv_timeout(Duration::from_secs(1)).unwrap(), + VsockVirtq::RX.raw() + ); + let used = rx_q.get_used().unwrap(); + assert_eq!(used.id, rx_buf_id); + assert_eq!(used.len as usize, size_of::()); + + let mut hdr = VsockHeader::new_zeroed(); + ram.read(rx_buf_addr, hdr.as_mut_bytes()).unwrap(); + assert_eq!(hdr.src_cid, VSOCK_CID_HOST); + assert_eq!(hdr.dst_cid, GUEST_CID); + assert_eq!(hdr.src_port, h2g_host_port); + assert_eq!(hdr.dst_port, H2G_GUEST_PORT); + assert_eq!(hdr.op, VsockOp::RST); + assert_eq!(hdr.type_, VsockType::STREAM); + + tx.send(WakeEvent::Shutdown).unwrap(); + notifier.notify().unwrap(); + handle.join().unwrap(); +} + +#[test] +fn vsock_host_close_no_desc_test() { + let ram_bus = Arc::new(fixture_ram_bus()); + let ram = ram_bus.lock_layout(); + let regs: Arc<[QueueReg]> = Arc::from(fixture_queues(3)); + let reg_tx = ®s[VsockVirtq::TX.raw() as usize]; + let reg_rx = ®s[VsockVirtq::RX.raw() as usize]; + let mut rx_q = GuestQueue::new( + SplitQueue::new(reg_rx, &ram, false).unwrap().unwrap(), + reg_rx, + ); + let mut tx_q = GuestQueue::new( + SplitQueue::new(reg_tx, &ram, false).unwrap().unwrap(), + reg_tx, + ); + + let temp_dir = TempDir::new().unwrap(); + let sock_path = temp_dir.path().join("vsock.sock"); + + const GUEST_CID: u32 = 3; + let param = UdsVsockSpec { + cid: GUEST_CID, + path: sock_path.clone().into(), + }; + let dev = param.build("vsock").unwrap(); + + let (tx, rx) = flume::unbounded(); + let (handle, notifier) = dev.spawn_worker(rx, ram_bus.clone(), regs).unwrap(); + let (irq_tx, irq_rx) = flume::unbounded(); + let irq_sender = Arc::new(FakeIrqSender { q_tx: irq_tx }); + let start_param = StartParam { + feature: VirtioFeature::VERSION_1.bits(), + irq_sender, + ioeventfds: Option::>::None, + }; + tx.send(WakeEvent::Start { param: start_param }).unwrap(); + + let rx_buf_addr = DATA_ADDR; + let tx_buf_addr = DATA_ADDR + 4096; + + // Establish a host-initiated connection + let mut h2g_stream = UnixStream::connect(&sock_path).unwrap(); + h2g_stream.set_nonblocking(true).unwrap(); + + let buf_id = rx_q.add_desc(&[], &[(rx_buf_addr, 4096)]); + const H2G_GUEST_PORT: u32 = 1025; + writeln!(h2g_stream, "CONNECT {H2G_GUEST_PORT}").unwrap(); + assert_eq!( + irq_rx.recv_timeout(Duration::from_secs(1)).unwrap(), + VsockVirtq::RX.raw() + ); + let used = rx_q.get_used().unwrap(); + assert_eq!(used.id, buf_id); + assert_eq!(used.len as usize, size_of::()); + + let mut hdr = VsockHeader::new_zeroed(); + ram.read(rx_buf_addr, hdr.as_mut_bytes()).unwrap(); + assert_eq!(hdr.op, VsockOp::REQUEST); + let h2g_host_port = hdr.src_port; + + let resp_hdr = VsockHeader { + src_cid: GUEST_CID, + dst_cid: VSOCK_CID_HOST, + src_port: H2G_GUEST_PORT, + dst_port: h2g_host_port, + op: VsockOp::RESPONSE, + type_: VsockType::STREAM, + ..Default::default() + }; + send_to_tx( + &resp_hdr, + &[], + &ram, + tx_buf_addr, + &mut tx_q, + &tx, + ¬ifier, + &irq_rx, + false, + ); + let mut reader = BufReader::new(&h2g_stream); + let mut line = String::new(); + reader.read_line(&mut line).unwrap(); + assert_eq!(line, format!("OK {h2g_host_port}\n")); + + // Write data and close the host socket WITHOUT providing an RX descriptor. + // The data must be delivered before the guest sees the final RST. + const DATA: &[u8] = b"drain before close"; + h2g_stream.write_all(DATA).unwrap(); + drop(h2g_stream); // EOF to alioth + + // Let the worker observe the socket event before an RX descriptor becomes + // available. + thread::sleep(Duration::from_millis(50)); + + // The first descriptor drains the host data. + let data_buf_id = rx_q.add_desc(&[], &[(rx_buf_addr, 4096)]); + tx.send(WakeEvent::Notify { + q_index: VsockVirtq::RX.raw(), + }) + .unwrap(); + notifier.notify().unwrap(); + + assert_eq!( + irq_rx.recv_timeout(Duration::from_secs(1)).unwrap(), + VsockVirtq::RX.raw() + ); + let used = rx_q.get_used().unwrap(); + assert_eq!(used.id, data_buf_id); + assert_eq!(used.len as usize, size_of::() + DATA.len()); + let mut hdr = VsockHeader::new_zeroed(); + ram.read(rx_buf_addr, hdr.as_mut_bytes()).unwrap(); + assert_eq!(hdr.op, VsockOp::RW); + assert_eq!(hdr.len as usize, DATA.len()); + let mut data = vec![0; DATA.len()]; + ram.read(rx_buf_addr + size_of::() as u64, &mut data) + .unwrap(); + assert_eq!(data, DATA); + + // The next descriptor observes EOF and receives the final RST. + let rst_buf_id = rx_q.add_desc(&[], &[(rx_buf_addr, 4096)]); + tx.send(WakeEvent::Notify { + q_index: VsockVirtq::RX.raw(), + }) + .unwrap(); + notifier.notify().unwrap(); + assert_eq!( + irq_rx.recv_timeout(Duration::from_secs(1)).unwrap(), + VsockVirtq::RX.raw() + ); + let used = rx_q.get_used().unwrap(); + assert_eq!(used.id, rst_buf_id); + assert_eq!(used.len as usize, size_of::()); + ram.read(rx_buf_addr, hdr.as_mut_bytes()).unwrap(); + assert_eq!(hdr.op, VsockOp::RST); + + tx.send(WakeEvent::Shutdown).unwrap(); + notifier.notify().unwrap(); + handle.join().unwrap(); +} diff --git a/alioth/src/virtio/queue/queue.rs b/alioth/src/virtio/queue/queue.rs index e83bc666..a4b47d20 100644 --- a/alioth/src/virtio/queue/queue.rs +++ b/alioth/src/virtio/queue/queue.rs @@ -108,6 +108,10 @@ where self.reg } + pub fn desc_avail(&self) -> bool { + self.q.desc_avail(self.avail) + } + fn push_used(&mut self, chain: DescChain, len: u32) { self.q.set_used(self.used, chain.id, len); self.used = self.q.index_add(self.used, chain.delta);