1mod rx_injection;
4#[cfg(test)]
5mod tests;
6mod virtio_device;
7
8use std::collections::HashMap;
9use std::sync::{Arc, Mutex, RwLock};
10
11use arcbox_virtio_core::error::{Result, VirtioError};
12use arcbox_virtio_core::queue::VirtQueue;
13use arcbox_virtio_core::{DeviceCtx, virtio_bindings};
14
15use crate::addr::{HOST_CID, RESERVED_CID, VsockAddr, VsockHostConnections};
16use crate::backend::VsockBackend;
17use crate::connection::{ConnectionState, VsockConnection};
18use crate::manager::VsockConnectionManager;
19use crate::protocol::{VsockHeader, VsockOp};
20
21fn write_all_with_backoff(fd: i32, buf: &[u8]) -> usize {
38 const MAX_POLL_RETRIES: u32 = 16;
39 const POLL_TIMEOUT_MS: libc::c_int = 2; let mut offset = 0usize;
42 let mut eagain_retries = 0u32;
43
44 while offset < buf.len() {
45 let ret = unsafe {
48 libc::write(
49 fd,
50 buf[offset..].as_ptr().cast::<libc::c_void>(),
51 buf.len() - offset,
52 )
53 };
54
55 use std::cmp::Ordering;
56 match ret.cmp(&0) {
57 Ordering::Greater => {
58 offset += ret as usize;
59 eagain_retries = 0;
60 }
61 Ordering::Equal => {
62 break;
64 }
65 Ordering::Less => {
66 let err = std::io::Error::last_os_error();
67 match err.raw_os_error() {
68 Some(e) if e == libc::EAGAIN || e == libc::EWOULDBLOCK => {
69 if eagain_retries >= MAX_POLL_RETRIES {
70 tracing::warn!(
71 "Vsock: giving up after {MAX_POLL_RETRIES} EAGAIN retries at offset {offset}/{} on fd {fd}",
72 buf.len(),
73 );
74 break;
75 }
76 eagain_retries += 1;
77 let mut pfd = libc::pollfd {
79 fd,
80 events: libc::POLLOUT,
81 revents: 0,
82 };
83 let _ = unsafe { libc::poll(&mut pfd, 1, POLL_TIMEOUT_MS) };
85 }
86 Some(libc::EINTR) => {}
87 _ => {
88 tracing::warn!("Vsock: write to fd {fd} failed at offset {offset}: {err}");
89 break;
90 }
91 }
92 }
93 }
94 }
95
96 offset
97}
98
99#[derive(Debug, Clone)]
101pub struct VsockConfig {
102 pub guest_cid: u64,
104}
105
106impl Default for VsockConfig {
107 fn default() -> Self {
108 Self {
109 guest_cid: 3, }
111 }
112}
113
114pub struct VirtioVsock {
119 config: VsockConfig,
120 features: u64,
121 acked_features: u64,
122 backend: Option<Arc<Mutex<dyn VsockBackend>>>,
124 connections: RwLock<HashMap<(u32, u32), VsockConnection>>,
126 rx_queue: Option<VirtQueue>,
128 tx_queue: Option<VirtQueue>,
130 event_queue: Option<VirtQueue>,
132 host_connections: HashMap<u32, std::os::unix::io::RawFd>,
136 last_avail_idx_tx: usize,
138 last_avail_idx_rx: usize,
140 ctx: Option<DeviceCtx>,
144 conns: Option<Arc<Mutex<dyn VsockHostConnections>>>,
148 conn_mgr: Option<Arc<Mutex<VsockConnectionManager>>>,
155}
156
157impl VirtioVsock {
158 pub const FEATURE_STREAM: u64 = 1 << 0;
160 pub const FEATURE_SEQPACKET: u64 = 1 << 1;
162 pub const FEATURE_VERSION_1: u64 = 1 << virtio_bindings::virtio_config::VIRTIO_F_VERSION_1;
164
165 pub const HOST_CID: u64 = HOST_CID;
167 pub const RESERVED_CID: u64 = RESERVED_CID;
169
170 #[must_use]
172 pub fn new(config: VsockConfig) -> Self {
173 Self {
174 config,
175 features: Self::FEATURE_STREAM
176 | Self::FEATURE_VERSION_1
177 | arcbox_virtio_core::queue::VIRTIO_F_EVENT_IDX,
178 acked_features: 0,
179 backend: None,
180 connections: RwLock::new(HashMap::new()),
181 rx_queue: None,
182 tx_queue: None,
183 event_queue: None,
184 host_connections: HashMap::new(),
185 last_avail_idx_tx: 0,
186 last_avail_idx_rx: 0,
187 ctx: None,
188 conns: None,
189 conn_mgr: None,
190 }
191 }
192
193 #[must_use]
195 pub fn with_backend<B: VsockBackend + 'static>(config: VsockConfig, backend: B) -> Self {
196 Self {
197 config,
198 features: Self::FEATURE_STREAM
199 | Self::FEATURE_VERSION_1
200 | arcbox_virtio_core::queue::VIRTIO_F_EVENT_IDX,
201 acked_features: 0,
202 backend: Some(Arc::new(Mutex::new(backend))),
203 connections: RwLock::new(HashMap::new()),
204 rx_queue: None,
205 tx_queue: None,
206 event_queue: None,
207 host_connections: HashMap::new(),
208 last_avail_idx_tx: 0,
209 last_avail_idx_rx: 0,
210 ctx: None,
211 conns: None,
212 conn_mgr: None,
213 }
214 }
215
216 pub fn set_backend<B: VsockBackend + 'static>(&mut self, backend: B) {
218 self.backend = Some(Arc::new(Mutex::new(backend)));
219 }
220
221 pub fn bind_ctx(&mut self, ctx: DeviceCtx) {
224 self.ctx = Some(ctx);
225 }
226
227 pub fn bind_connections(&mut self, conns: Arc<Mutex<dyn VsockHostConnections>>) {
232 self.conns = Some(conns);
233 }
234
235 pub fn bind_connection_manager(&mut self, mgr: Arc<Mutex<VsockConnectionManager>>) {
240 self.conns = Some(mgr.clone());
241 self.conn_mgr = Some(mgr);
242 }
243
244 pub fn connections(&self) -> Option<Arc<Mutex<dyn VsockHostConnections>>> {
246 self.conns.clone()
247 }
248
249 #[must_use]
251 pub const fn guest_cid(&self) -> u64 {
252 self.config.guest_cid
253 }
254
255 pub fn handle_connect(&self, src_port: u32, dst_port: u32) -> Result<()> {
257 let local = VsockAddr::new(self.config.guest_cid, src_port);
258 let remote = VsockAddr::new(Self::HOST_CID, dst_port);
259
260 let mut conn = VsockConnection::new(local, remote);
261 conn.state = ConnectionState::Connecting;
262
263 if let Some(ref backend) = self.backend {
264 backend.lock().unwrap().on_connect(local)?;
265 conn.state = ConnectionState::Connected;
266 }
267
268 self.connections
269 .write()
270 .unwrap()
271 .insert((src_port, dst_port), conn);
272 tracing::debug!(
273 "Vsock connect: {}:{} -> {}:{}",
274 self.config.guest_cid,
275 src_port,
276 Self::HOST_CID,
277 dst_port
278 );
279
280 Ok(())
281 }
282
283 pub fn handle_send(&self, src_port: u32, dst_port: u32, data: &[u8]) -> Result<usize> {
285 let local = VsockAddr::new(self.config.guest_cid, src_port);
286
287 if let Some(ref backend) = self.backend {
288 backend.lock().unwrap().on_send(local, data)
289 } else {
290 let mut conns = self.connections.write().unwrap();
291 if let Some(conn) = conns.get_mut(&(src_port, dst_port)) {
292 conn.enqueue_tx(data);
293 Ok(data.len())
294 } else {
295 Err(VirtioError::InvalidOperation("Connection not found".into()))
296 }
297 }
298 }
299
300 pub fn handle_recv(&self, src_port: u32, dst_port: u32, buf: &mut [u8]) -> Result<usize> {
302 let local = VsockAddr::new(self.config.guest_cid, src_port);
303
304 if let Some(ref backend) = self.backend {
305 backend.lock().unwrap().on_recv(local, buf)
306 } else {
307 let mut conns = self.connections.write().unwrap();
308 if let Some(conn) = conns.get_mut(&(src_port, dst_port)) {
309 let data = conn.dequeue_rx(buf.len());
310 buf[..data.len()].copy_from_slice(&data);
311 Ok(data.len())
312 } else {
313 Err(VirtioError::InvalidOperation("Connection not found".into()))
314 }
315 }
316 }
317
318 pub fn handle_close(&self, src_port: u32, dst_port: u32) -> Result<()> {
320 let local = VsockAddr::new(self.config.guest_cid, src_port);
321
322 if let Some(ref backend) = self.backend {
323 backend.lock().unwrap().on_close(local)?;
324 }
325
326 self.connections
327 .write()
328 .unwrap()
329 .remove(&(src_port, dst_port));
330 tracing::debug!("Vsock close: {}:{}", self.config.guest_cid, src_port);
331
332 Ok(())
333 }
334
335 #[must_use]
337 pub fn connection_count(&self) -> usize {
338 self.connections.read().unwrap().len()
339 }
340
341 pub fn tx_queue_mut(&mut self) -> Option<&mut VirtQueue> {
343 self.tx_queue.as_mut()
344 }
345
346 pub fn rx_queue_mut(&mut self) -> Option<&mut VirtQueue> {
348 self.rx_queue.as_mut()
349 }
350
351 fn handle_tx_packet_with_fds(
353 &self,
354 hdr: &VsockHeader,
355 payload: &[u8],
356 connections: Option<&mut dyn VsockHostConnections>,
357 ) {
358 let src_cid = { hdr.src_cid };
360 let dst_cid = { hdr.dst_cid };
361 let src_port = { hdr.src_port };
362 let dst_port = { hdr.dst_port };
363 let buf_alloc = { hdr.buf_alloc };
364 let fwd_cnt = { hdr.fwd_cnt };
365 let flags = { hdr.flags };
366
367 match hdr.operation() {
368 Some(VsockOp::Request) => {
369 tracing::debug!(
370 "Vsock TX: OP_REQUEST src={}:{} dst={}:{}",
371 src_cid,
372 src_port,
373 dst_cid,
374 dst_port,
375 );
376 }
377 Some(VsockOp::Response) => {
378 tracing::info!(
381 "Vsock TX: OP_RESPONSE — connection established (guest_port={}, host_port={})",
382 src_port,
383 dst_port,
384 );
385 if let Some(conns) = connections {
386 conns.update_peer_credit(src_port, dst_port, buf_alloc, fwd_cnt);
387 conns.mark_connected(src_port, dst_port);
388 }
389 }
390 Some(VsockOp::Rw) => {
391 if let Some(conns) = connections {
393 conns.update_peer_credit(src_port, dst_port, buf_alloc, fwd_cnt);
394 if let Some(fd) = conns.fd_for(src_port, dst_port) {
395 if !payload.is_empty() {
396 let total = payload.len();
397 let forwarded = write_all_with_backoff(fd, payload);
398 if forwarded > 0 {
399 tracing::debug!(
400 "Vsock TX: OP_RW guest_port={} host_port={} -> fd {fd}, {}/{} bytes",
401 src_port,
402 dst_port,
403 forwarded,
404 total,
405 );
406 #[allow(clippy::cast_possible_truncation)]
413 {
414 conns.advance_fwd_cnt(src_port, dst_port, forwarded as u32);
415 }
416 }
417 if forwarded < total {
418 tracing::warn!(
419 "Vsock TX: truncated write guest_port={} host_port={}: only {}/{} bytes forwarded (ABX-365)",
420 src_port,
421 dst_port,
422 forwarded,
423 total,
424 );
425 }
426 }
427 } else {
428 tracing::warn!(
429 "Vsock TX: OP_RW no host fd for guest_port={} host_port={}",
430 src_port,
431 dst_port,
432 );
433 }
434 }
435 }
436 Some(VsockOp::Shutdown) => {
437 tracing::debug!(
438 "Vsock TX: OP_SHUTDOWN guest_port={} host_port={} flags=0x{:x}",
439 src_port,
440 dst_port,
441 flags,
442 );
443 if let Some(conns) = connections {
444 conns.handle_shutdown(src_port, dst_port, flags);
448 }
449 }
450 Some(VsockOp::Rst) => {
451 tracing::debug!(
452 "Vsock TX: OP_RST guest_port={} host_port={}",
453 src_port,
454 dst_port,
455 );
456 if let Some(conns) = connections {
457 conns.remove_connection(src_port, dst_port);
458 }
459 }
460 Some(VsockOp::CreditUpdate) => {
461 tracing::trace!(
462 "Vsock TX: OP_CREDIT_UPDATE guest_port={} host_port={} buf_alloc={} fwd_cnt={}",
463 src_port,
464 dst_port,
465 buf_alloc,
466 fwd_cnt,
467 );
468 if let Some(conns) = connections {
469 conns.update_peer_credit(src_port, dst_port, buf_alloc, fwd_cnt);
470 }
471 }
472 Some(VsockOp::CreditRequest) => {
473 tracing::trace!(
474 "Vsock TX: OP_CREDIT_REQUEST guest_port={} host_port={}",
475 src_port,
476 dst_port,
477 );
478 if let Some(conns) = connections {
479 conns.update_peer_credit(src_port, dst_port, buf_alloc, fwd_cnt);
480 conns.enqueue_credit_update(src_port, dst_port);
481 }
482 }
483 _ => {}
484 }
485 }
486
487 pub fn add_host_connection(&mut self, guest_port: u32, fd: std::os::unix::io::RawFd) {
491 tracing::info!("Vsock: host connection for guest port {guest_port} -> fd {fd}");
492 self.host_connections.insert(guest_port, fd);
493 }
494
495 pub fn process_tx_queue(&mut self, memory: &mut [u8]) -> Result<Vec<(u16, u32)>> {
506 let mut raw_packets: Vec<(u16, Vec<u8>)> = Vec::new();
508
509 {
510 let queue = self
511 .tx_queue
512 .as_mut()
513 .ok_or_else(|| VirtioError::NotReady("TX queue not ready".into()))?;
514
515 while let Some((head_idx, chain)) = queue.pop_avail() {
516 let mut data = Vec::new();
517
518 for desc in chain {
519 if !desc.is_write_only() {
520 let start = desc.addr as usize;
522 let end = start + desc.len as usize;
523 if end <= memory.len() {
524 data.extend_from_slice(&memory[start..end]);
525 }
526 }
527 }
528
529 raw_packets.push((head_idx, data));
530 }
531 }
532
533 let mut completions = Vec::new();
535 let mut rx_inject: Vec<(VsockHeader, Vec<u8>)> = Vec::new();
537
538 for (head_idx, data) in &raw_packets {
539 if data.len() < VsockHeader::SIZE {
540 tracing::warn!(
541 "Vsock TX: descriptor {} too short ({} bytes), skipping",
542 head_idx,
543 data.len()
544 );
545 completions.push((*head_idx, 0u32));
546 continue;
547 }
548
549 let header = match VsockHeader::from_bytes(&data[..VsockHeader::SIZE]) {
550 Some(h) => h,
551 None => {
552 tracing::warn!(
553 "Vsock TX: failed to parse header for descriptor {}",
554 head_idx
555 );
556 completions.push((*head_idx, 0u32));
557 continue;
558 }
559 };
560
561 let payload_len = { header.len } as usize;
562 let payload = if payload_len > 0 && data.len() > VsockHeader::SIZE {
563 let avail = data.len() - VsockHeader::SIZE;
564 &data[VsockHeader::SIZE..VsockHeader::SIZE + payload_len.min(avail)]
565 } else {
566 &[] as &[u8]
567 };
568
569 let src_port = { header.src_port };
570 let dst_port = { header.dst_port };
571
572 match header.operation() {
573 Some(VsockOp::Request) => {
574 tracing::debug!(
575 "Vsock TX: OP_REQUEST from port {} to port {}",
576 src_port,
577 dst_port
578 );
579 match self.handle_connect(src_port, dst_port) {
580 Ok(()) => {
581 let resp = VsockHeader::new(
583 VsockAddr::new(Self::HOST_CID, dst_port),
584 VsockAddr::new(self.config.guest_cid, src_port),
585 VsockOp::Response,
586 );
587 rx_inject.push((resp, Vec::new()));
588 }
589 Err(e) => {
590 tracing::warn!("Vsock TX: connect failed: {}", e);
591 let rst = VsockHeader::new(
593 VsockAddr::new(Self::HOST_CID, dst_port),
594 VsockAddr::new(self.config.guest_cid, src_port),
595 VsockOp::Rst,
596 );
597 rx_inject.push((rst, Vec::new()));
598 }
599 }
600 }
601 Some(VsockOp::Response) => {
602 tracing::debug!(
604 "Vsock TX: OP_RESPONSE from port {} to port {}",
605 src_port,
606 dst_port
607 );
608 let mut conns = self.connections.write().unwrap();
609 if let Some(conn) = conns.get_mut(&(src_port, dst_port)) {
610 conn.state = ConnectionState::Connected;
611 }
612 }
613 Some(VsockOp::Rw) => {
614 tracing::trace!(
615 "Vsock TX: OP_RW {} bytes from port {} to port {}",
616 payload.len(),
617 src_port,
618 dst_port
619 );
620 if let Err(e) = self.handle_send(src_port, dst_port, payload) {
621 tracing::warn!("Vsock TX: send failed: {}", e);
622 }
623 }
624 Some(VsockOp::Shutdown) => {
625 tracing::debug!(
626 "Vsock TX: OP_SHUTDOWN from port {} to port {}",
627 src_port,
628 dst_port
629 );
630 if let Err(e) = self.handle_close(src_port, dst_port) {
631 tracing::warn!("Vsock TX: close failed: {}", e);
632 }
633 let rst = VsockHeader::new(
635 VsockAddr::new(Self::HOST_CID, dst_port),
636 VsockAddr::new(self.config.guest_cid, src_port),
637 VsockOp::Rst,
638 );
639 rx_inject.push((rst, Vec::new()));
640 }
641 Some(VsockOp::Rst) => {
642 tracing::debug!(
643 "Vsock TX: OP_RST from port {} to port {}",
644 src_port,
645 dst_port
646 );
647 let _ = self.handle_close(src_port, dst_port);
648 }
649 Some(VsockOp::CreditUpdate) => {
650 let buf_alloc = { header.buf_alloc };
651 let fwd_cnt = { header.fwd_cnt };
652 tracing::trace!(
653 "Vsock TX: OP_CREDIT_UPDATE port {} buf_alloc={} fwd_cnt={}",
654 src_port,
655 buf_alloc,
656 fwd_cnt
657 );
658 let mut conns = self.connections.write().unwrap();
659 if let Some(conn) = conns.get_mut(&(src_port, dst_port)) {
660 conn.update_peer_credit(buf_alloc, fwd_cnt);
661 }
662 }
663 Some(VsockOp::CreditRequest) => {
664 tracing::trace!(
665 "Vsock TX: OP_CREDIT_REQUEST from port {} to port {}",
666 src_port,
667 dst_port
668 );
669 let conns = self.connections.read().unwrap();
671 if let Some(conn) = conns.get(&(src_port, dst_port)) {
672 let mut update = VsockHeader::new(
673 VsockAddr::new(Self::HOST_CID, dst_port),
674 VsockAddr::new(self.config.guest_cid, src_port),
675 VsockOp::CreditUpdate,
676 );
677 update.buf_alloc = conn.buf_alloc;
678 update.fwd_cnt = conn.fwd_cnt;
679 rx_inject.push((update, Vec::new()));
680 }
681 }
682 Some(VsockOp::Invalid) | None => {
683 let raw_op = { header.op };
684 tracing::warn!(
685 "Vsock TX: unknown/invalid op {} from port {}",
686 raw_op,
687 src_port
688 );
689 }
690 }
691
692 completions.push((*head_idx, data.len() as u32));
693 }
694
695 for (hdr, payload) in rx_inject {
697 if let Err(e) = self.inject_rx_packet(&hdr, &payload, memory) {
698 tracing::warn!("Vsock: failed to inject RX packet: {}", e);
699 }
700 }
701
702 Ok(completions)
703 }
704
705 pub fn process_queue(&mut self, queue_idx: u16, memory: &mut [u8]) -> Result<Vec<(u16, u32)>> {
716 match queue_idx {
717 1 => self.process_tx_queue(memory),
718 _ => Ok(Vec::new()),
719 }
720 }
721
722 pub fn inject_rx_packet(
734 &mut self,
735 header: &VsockHeader,
736 data: &[u8],
737 memory: &mut [u8],
738 ) -> Result<()> {
739 let queue = self
740 .rx_queue
741 .as_mut()
742 .ok_or_else(|| VirtioError::NotReady("RX queue not ready".into()))?;
743
744 let (head_idx, chain) = queue
745 .pop_avail()
746 .ok_or_else(|| VirtioError::InvalidQueue("No available RX descriptors".into()))?;
747
748 let header_bytes = header.to_bytes();
749 let total_len = header_bytes.len() + data.len();
750 let mut frame = Vec::with_capacity(total_len);
751 frame.extend_from_slice(&header_bytes);
752 frame.extend_from_slice(data);
753
754 let mut written = 0usize;
755 for desc in chain {
756 if !desc.is_write_only() {
757 continue;
758 }
759 let start = desc.addr as usize;
760 let remaining = frame.len().saturating_sub(written);
761 let to_write = remaining.min(desc.len as usize);
762 if to_write == 0 {
763 continue;
764 }
765 let end = start + to_write;
766 if end > memory.len() {
767 return Err(VirtioError::MemoryError(
768 "RX descriptor points outside guest memory".into(),
769 ));
770 }
771 memory[start..end].copy_from_slice(&frame[written..written + to_write]);
772 written += to_write;
773 }
774
775 queue.push_used(head_idx, written as u32);
776 Ok(())
777 }
778}