diff --git a/alioth/src/virtio/vu/backend.rs b/alioth/src/virtio/vu/backend.rs index a672e89e..8a648319 100644 --- a/alioth/src/virtio/vu/backend.rs +++ b/alioth/src/virtio/vu/backend.rs @@ -154,7 +154,7 @@ impl VuBackend { conn.set_nonblocking(false)?; let queue_num = dev.queue_regs.len(); Ok(VuBackend { - session: VuSession { conn }, + session: VuSession::from(conn), channel: None, dev, memory, diff --git a/alioth/src/virtio/vu/conn.rs b/alioth/src/virtio/vu/conn.rs index 7a3b342d..5efbbed6 100644 --- a/alioth/src/virtio/vu/conn.rs +++ b/alioth/src/virtio/vu/conn.rs @@ -17,6 +17,7 @@ use std::os::fd::{AsFd, BorrowedFd, FromRawFd, OwnedFd}; use std::os::unix::net::UnixStream; use std::path::Path; +use parking_lot::Mutex; use snafu::ResultExt; use zerocopy::{FromBytes, FromZeros, Immutable, IntoBytes}; @@ -179,7 +180,15 @@ fn recv_config(mut conn: &UnixStream, buf: &mut [u8]) -> Result { #[derive(Debug)] pub struct VuSession { - pub conn: UnixStream, + conn: Mutex, +} + +impl From for VuSession { + fn from(conn: UnixStream) -> Self { + VuSession { + conn: Mutex::new(conn), + } + } } impl VuSession { @@ -187,7 +196,7 @@ impl VuSession { let conn = UnixStream::connect(&path).context(error::AccessSocket { path: path.as_ref(), })?; - Ok(VuSession { conn }) + Ok(Self::from(conn)) } fn send(&self, req: VuFrontMsg, payload: &T, fds: &[BorrowedFd]) -> Result @@ -195,22 +204,26 @@ impl VuSession { T: IntoBytes + Immutable, R: FromBytes + IntoBytes, { - send(&self.conn, req.raw(), payload, &[], &mut [], fds) + let conn = self.conn.lock(); + send(&conn, req.raw(), payload, &[], &mut [], fds) } pub fn recv_payload(&self) -> Result where T: IntoBytes + Immutable + FromBytes, { - recv_with_fds(&self.conn, &mut []) + let conn = self.conn.lock(); + recv_with_fds(&conn, &mut []) } pub fn recv_config(&self, buf: &mut [u8]) -> Result { - recv_config(&self.conn, buf) + let conn = self.conn.lock(); + recv_config(&conn, buf) } pub fn recv_msg(&self, fds: &mut [Option]) -> Result { - recv_with_fds(&self.conn, fds) + let conn = self.conn.lock(); + recv_with_fds(&conn, fds) } pub fn reply( @@ -219,11 +232,13 @@ impl VuSession { payload: &T, fds: &[BorrowedFd], ) -> Result<()> { - reply(&self.conn, req.raw(), payload, fds) + let conn = self.conn.lock(); + reply(&conn, req.raw(), payload, fds) } pub fn reply_config(&self, config: &DeviceConfig, buf: &[u8]) -> Result<()> { - reply_config(&self.conn, config, buf) + let conn = self.conn.lock(); + reply_config(&conn, config, buf) } pub fn get_features(&self) -> Result { @@ -262,7 +277,8 @@ impl VuSession { let in_ = [0; MAX_CONFIG_SIZE]; let len = buf.len(); let req = VuFrontMsg::GET_CONFIG.raw(); - send(&self.conn, req, payload, &in_[..len], buf, &[]) + let conn = self.conn.lock(); + send(&conn, req, payload, &in_[..len], buf, &[]) } pub fn set_config(&self, payload: &DeviceConfig, buf: &[u8]) -> Result<()> @@ -270,7 +286,8 @@ impl VuSession { DeviceConfig: IntoBytes, { let req = VuFrontMsg::SET_CONFIG.raw(); - send(&self.conn, req, payload, buf, &mut [], &[]) + let conn = self.conn.lock(); + send(&conn, req, payload, buf, &mut [], &[]) } pub fn get_virtq_base(&self, payload: &VirtqState) -> Result { @@ -371,3 +388,7 @@ impl VuChannel { self.send(VuBackMsg::SHMEM_UNMAP, payload, &[]) } } + +#[cfg(test)] +#[path = "conn_test.rs"] +mod tests; diff --git a/alioth/src/virtio/vu/conn_test.rs b/alioth/src/virtio/vu/conn_test.rs new file mode 100644 index 00000000..ee9ead84 --- /dev/null +++ b/alioth/src/virtio/vu/conn_test.rs @@ -0,0 +1,82 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::os::unix::net::UnixStream; +use std::sync::Arc; + +use zerocopy::IntoBytes; + +use crate::virtio::vu::bindings::{DeviceConfig, MAX_CONFIG_SIZE, VuFrontMsg}; +use crate::virtio::vu::conn::VuSession; + +#[test] +fn test_vu_session_concurrent_requests() { + let (client_conn, server_conn) = UnixStream::pair().unwrap(); + let client = Arc::new(VuSession::from(client_conn)); + let server = VuSession::from(server_conn); + + let server_thread = std::thread::spawn(move || { + let mut fds = [const { None }; 8]; + while let Ok(msg) = server.recv_msg(&mut fds) { + match VuFrontMsg::from(msg.request) { + VuFrontMsg::GET_FEATURES => { + server + .reply(VuFrontMsg::GET_FEATURES, &0xdead_beef_u64, &[]) + .unwrap(); + } + VuFrontMsg::GET_CONFIG => { + let mut buf = [0u8; MAX_CONFIG_SIZE]; + let cfg = server.recv_config(&mut buf).unwrap(); + let val = 0x0102_0304_0506_0708_u64; + server + .reply_config(&cfg, &val.as_bytes()[..cfg.size as usize]) + .unwrap(); + } + VuFrontMsg::SET_STATUS => { + let _status: u64 = server.recv_payload().unwrap(); + server.reply(VuFrontMsg::SET_STATUS, &0u64, &[]).unwrap(); + } + req => panic!("unexpected request: {req:?}"), + } + } + }); + + let mut handles = Vec::new(); + for i in 0..4 { + let session = client.clone(); + handles.push(std::thread::spawn(move || { + for _ in 0..50 { + if i % 2 == 0 { + assert_eq!(session.get_features().unwrap(), 0xdead_beef); + session.set_status(&0).unwrap(); + } else { + let req = DeviceConfig { + offset: 0, + size: 8, + flags: 0, + }; + let mut out = 0u64; + session.get_config(&req, out.as_mut_bytes()).unwrap(); + assert_eq!(out, 0x0102_0304_0506_0708); + } + } + })); + } + + for h in handles { + h.join().unwrap(); + } + drop(client); + server_thread.join().unwrap(); +}