Dragonball VMM 侧 Vsock 设计与实现

pub struct Vsock<AS: GuestAddressSpace, M: VsockGenericMuxer = VsockMuxer> {
    cid: u64,
    queue_sizes: Arc<Vec<u16>>,
    device_info: VirtioDeviceInfo,
    subscriber_id: Option<SubscriberId>,
    muxer: Option<M>,
    phantom: PhantomData<AS>,
}

impl VsockDeviceMgr {
  // ...
  pub fn attach_devices(
        &mut self,
        ctx: &mut DeviceOpContext,
    ) -> std::result::Result<(), StartMicroVmError> {
      // ...
      for info in self.info_list.iter_mut() {
        // ...
        if let Some(uds_path) = info.config.uds_path.as_ref() {
          // (1)
          let unix_backend = VsockUnixStreamBackend::new(uds_path.clone())
            .map_err(VirtioError::VirtioVsockError)
            .map_err(StartMicroVmError::CreateVsockDevice)?;
          // (2)
          device
            .add_backend(Box::new(unix_backend), true)
            .map_err(VirtioError::VirtioVsockError)
            .map_err(StartMicroVmError::CreateVsockDevice)?;
        }
        // ...
}

impl VsockUnixStreamBackend {
    pub fn new(host_sock_path: String) -> Result<Self> {
        // ...
        let host_sock = UnixListener::bind(&host_sock_path)
            .and_then(|sock| sock.set_nonblocking(true).map(|_| sock))
            .map_err(VsockError::Backend)?;
        // ...
        Ok(VsockUnixStreamBackend {
            host_sock,
            host_sock_path,
        })
    }
}

impl VsockGenericMuxer for VsockMuxer {
    fn add_backend(&mut self, backend: Box<dyn VsockBackend>, is_peer_backend: bool) -> Result<()> {
      // backend_type 是 `VsockBackendType::UnixStream`
      let backend_type = backend.r#type();
        if self.backend_map.contains_key(&backend_type) {
            return Err(Error::BackendRegistered(backend_type));
        }
     // (2.1) 
        self.add_listener(
            backend.as_raw_fd(),
            EpollListener::Backend(backend_type.clone()),
        )?;
      // (2.2)
        self.backend_map.insert(backend_type.clone(), backend);
        if is_peer_backend {
            self.peer_backend = Some(backend_type);
        }
        Ok(())
    }
}

impl VsockMuxer {
  // ...
    fn handle_event(&mut self, fd: RawFd, event_set: epoll::Events) {
      // ...
        match self.listener_map.get_mut(&fd) {
            // ...
            Some(EpollListener::Backend(backend_type)) => {
                if let Some(backend) = self.backend_map.get_mut(backend_type) {
                    // ...
                    backend
                  // (4)
                        .accept()
                        .map_err(Error::BackendAccept)
                        .and_then(|stream| {
                          // (5)
                            self.add_listener(
                                stream.as_raw_fd(),
                                EpollListener::LocalStream(stream),
                            )
                        })
                        .unwrap_or_else(|err| {
                            warn!("vsock: unable to accept local connection: {:?}", err);
                        });
                } else {
                    error!("vsock: unsable to find specific backend {:?}", backend_type)
                }
            }
            // ...
        }
    }

pub trait VsockBackend: AsRawFd + Send {
    fn accept(&mut self) -> std::io::Result<Box<dyn VsockStream>>;
    fn connect(&self, dst_port: u32) -> std::io::Result<Box<dyn VsockStream>>;
    fn r#type(&self) -> VsockBackendType;
    fn as_any(&self) -> &dyn Any;
}

impl VsockBackend for VsockUnixStreamBackend {
    fn accept(&mut self) -> std::io::Result<Box<dyn VsockStream>> {
      // (4.1)
        let (stream, _) = self.host_sock.accept()?;
        stream.set_nonblocking(true)?;

        Ok(Box::new(stream))
    }
  // ...
}

fn handle_event(&mut self, fd: RawFd, event_set: epoll::Events) {
    // ...
    match self.listener_map.get_mut(&fd) {
        // (6)
        Some(EpollListener::LocalStream(_)) => {
          // (7)
            if let Some(EpollListener::LocalStream(mut stream)) = self.remove_listener(fd) {
              // (8)
                Self::read_local_stream_port(&mut stream)
                    .and_then(|read_port_result| match read_port_result {
                        ReadPortResult::Connect(peer_port) => {
                          // (9)
                            let local_port = self.allocate_local_port();
                          // (10)
                            self.add_connection(
                                ConnMapKey {
                                    local_port,
                                    peer_port,
                                },
                                VsockConnection::new_local_init(
                                    stream,
                                    uapi::VSOCK_HOST_CID,
                                    self.cid,
                                    local_port,
                                    peer_port,
                                ),
                            )
                        }
                        // ...
                    })
                    .unwrap_or_else(|err| {
                        info!("vsock: error adding local-init connection: {:?}", err);
                    })
            }
        }
        // ...
    }
}

pub fn new_local_init(
    stream: Box<dyn VsockStream>,
    local_cid: u64,
    peer_cid: u64,
    local_port: u32,
    peer_port: u32,
) -> Self {
    Self {
        local_cid,
        peer_cid,
        local_port,
        peer_port,
        stream,
        state: ConnState::LocalInit,
        tx_buf: TxBuf::default(),
        fwd_cnt: Wrapping(0),
        peer_buf_alloc: 0,
        peer_fwd_cnt: Wrapping(0),
        rx_cnt: Wrapping(0),
        last_fwd_cnt_to_peer: Wrapping(0),
      // (10.1)
        pending_rx: PendingRxSet::from(PendingRx::Request),
        expiry: None,
    }
}

fn add_connection(&mut self, key: ConnMapKey, conn: VsockConnection) -> Result<()> {
    // ...
  // (10.2)
    self.add_listener(
        conn.as_raw_fd(),
        EpollListener::Connection {
            key,
            evset: conn.get_polled_evset(),
            backend: conn.stream.backend_type(),
        },
    )
    .map(|_| {
      // (10.3)
        if conn.has_pending_rx() {
            self.rxq.push(MuxerRx::ConnRx(key));
        }
      // (10.4)
        self.conn_map.insert(key, conn);
    })
}

fn process_rx(&mut self, mem: &AS::M) {
    trace!("{}: epoll_handler::process_rx()", self.id);
    let mut raise_irq = false;
    {
      // (11.1)
        let rxvq = &mut self.config.queues[QUEUE_RX].queue_mut().lock();
        loop {
          // (11.2)
            let mut iter = match rxvq.iter(mem) {
                // ...
            };
          // (11.3)
            if let Some(mut desc_chain) = iter.next() {
              // (12)
                let used_len = match VsockPacket::from_rx_virtq_head(&mut desc_chain) {
                    Ok(mut pkt) => {
                      // (13)
                        if self.muxer.recv_pkt(&mut pkt).is_ok() {
                            pkt.hdr().len() as u32 + pkt.len()
                        }
                        // ...
                    }
                    // ...
                };
                raise_irq = true;
              // (11.4)
                let _ = rxvq.add_used(mem, desc_chain.head_index(), used_len);
            } else {
                break;
            }
        }
    }
    if raise_irq {
      // (11.5)
        if let Err(e) = self.signal_used_queue(QUEUE_RX) {
            error!("{}: failed to notify guest for RX queue, {:?}", self.id, e);
        }
    }
}

pub trait VsockChannel {
    fn recv_pkt(&mut self, pkt: &mut VsockPacket) -> Result<()>;
    fn send_pkt(&mut self, pkt: &VsockPacket) -> Result<()>;
    fn has_pending_rx(&self) -> bool;
}

impl VsockChannel for VsockMuxer {
  // ...
}

impl VsockChannel for VsockConnection {
  // ...
}

fn recv_pkt(&mut self, pkt: &mut VsockPacket) -> VsockResult<()> {
    // (13.1)
    while let Some(rx) = self.rxq.peek() {
        let res = match rx {
            // ...
            MuxerRx::ConnRx(key) => {
                let mut conn_res = Err(VsockError::NoData);
                let mut do_pop = true;
              // (13.3)
                self.apply_conn_mutation(key, |conn| {
                  // (14)
                    conn_res = conn.recv_pkt(pkt);
                  // (13.2)
                    do_pop = !conn.has_pending_rx();
                });
                if do_pop {
                    self.rxq.pop().unwrap();
                }
                conn_res
            }
        };
      // (13.4)
      if res.is_ok() {
            // ...
            return Ok(());
        }
    }
    Err(VsockError::NoData)
}

fn recv_pkt(&mut self, pkt: &mut VsockPacket) -> VsockResult<()> {
    // (14.1)
    self.init_pkt(pkt);

    // (14.2)
    if self.pending_rx.remove(PendingRx::Rst) {
        pkt.set_op(uapi::VSOCK_OP_RST);
        return Ok(());
    }

    // (14.3)
    if self.pending_rx.remove(PendingRx::Response) {
        self.state = ConnState::Established;
        pkt.set_op(uapi::VSOCK_OP_RESPONSE);
        return Ok(());
    }

    // (14.4)
    if self.pending_rx.remove(PendingRx::Request) {
        self.expiry =
            Some(Instant::now() + Duration::from_millis(defs::CONN_REQUEST_TIMEOUT_MS));
        pkt.set_op(uapi::VSOCK_OP_REQUEST);
        return Ok(());
    }

  // (14.5)
    if self.pending_rx.remove(PendingRx::Rw) {
        // ...
      let buf = pkt.buf_mut().ok_or(VsockError::PktBufMissing)?;
      // ...
        match self.stream.read(&mut buf[..max_len]) {
            Ok(read_cnt) => {
                if read_cnt == 0 {
                    // ...
                } else {
                    pkt.set_op(uapi::VSOCK_OP_RW).set_len(read_cnt as u32);
                }
                // ...
                return Ok(());
            }
            // ...
        };
    }
    // ...
    Err(VsockError::NoData)
}

fn init_pkt<'a>(&self, pkt: &'a mut VsockPacket) -> &'a mut VsockPacket {
    for b in pkt.hdr_mut() {
        *b = 0;
    }
    pkt.set_src_cid(self.local_cid)
        .set_dst_cid(self.peer_cid)
        .set_src_port(self.local_port)
        .set_dst_port(self.peer_port)
        .set_type(uapi::VSOCK_TYPE_STREAM)
        .set_buf_alloc(defs::CONN_TX_BUF_SIZE)
        .set_fwd_cnt(self.fwd_cnt.0)
}