poly_msg_socket

we're gonna need this to send all of VirtioDevice over a socket

2 files changed, 56 insertions, 9 deletions

diff --git a/devices/src/virtio/controller.rs b/devices/src/virtio/controller.rs
index 815bb46..7b1b2cc 100644
--- a/devices/src/virtio/controller.rs
+++ b/devices/src/virtio/controller.rs
@@ -37,12 +37,13 @@ use super::resource_bridge::*;
 use super::{Interrupt, InterruptProxyEvent, Queue, VirtioDevice, TYPE_WL, VIRTIO_F_VERSION_1};
 use vm_control::{MaybeOwnedFd, VmMemoryControlRequestSocket};
 
-use msg_socket::{MsgError, MsgOnSocket, MsgReceiver, MsgSender, MsgSocket};
+use msg_socket::{MsgOnSocket, MsgReceiver, MsgSocket};
+use serde::{Deserialize, Serialize};
 use sys_util::net::UnixSeqpacket;
 use sys_util::{error, EventFd, GuestMemory, PollContext, PollToken, SharedMemory};
 
 #[derive(Debug, MsgOnSocket)]
-pub enum Request {
+pub enum MsgOnSocketRequest {
     Create {
         // wayland_paths: Map<String, PathBuf>,
         vm_socket: MaybeOwnedFd<UnixSeqpacket>,
@@ -60,12 +61,50 @@ pub enum Request {
     Kill,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub enum BincodeRequest {
+    WriteConfig { offset: u64, data: Vec<u8> },
+}
+
+pub type Request = poly_msg_socket::Value<MsgOnSocketRequest, BincodeRequest>;
+
+impl From<MsgOnSocketRequest> for Request {
+    fn from(request: MsgOnSocketRequest) -> Self {
+        Self::MsgOnSocket(request)
+    }
+}
+
+impl From<BincodeRequest> for Request {
+    fn from(request: BincodeRequest) -> Self {
+        Self::Bincode(request)
+    }
+}
+
 #[derive(Debug, MsgOnSocket)]
-pub enum Response {
+pub enum MsgOnSocketResponse {
     Kill,
 }
 
-type Socket = MsgSocket<Request, Response>;
+#[derive(Debug, Deserialize, Serialize)]
+pub struct BincodeResponse;
+
+pub type Response = poly_msg_socket::Value<MsgOnSocketResponse, BincodeResponse>;
+
+impl From<MsgOnSocketResponse> for Response {
+    fn from(response: MsgOnSocketResponse) -> Self {
+        Self::MsgOnSocket(response)
+    }
+}
+
+impl From<BincodeResponse> for Response {
+    fn from(response: BincodeResponse) -> Self {
+        Self::Bincode(response)
+    }
+}
+
+use poly_msg_socket::PolyMsgSocket;
+type Socket =
+    PolyMsgSocket<MsgOnSocketRequest, MsgOnSocketResponse, BincodeRequest, BincodeResponse>;
 
 const VIRTIO_WL_F_TRANS_FLAGS: u32 = 0x01;
 
@@ -96,11 +135,14 @@ impl Worker {
     }
 
     fn handle_response(&mut self) {
+        use poly_msg_socket::Value::*;
         match self.device_socket.recv() {
-            Ok(Response::Kill) => {
+            Ok(MsgOnSocket(MsgOnSocketResponse::Kill)) => {
                 self.shutdown = true;
             }
 
+            Ok(Bincode(BincodeResponse)) => unreachable!(),
+
             Err(e) => {
                 error!("recv failed: {:?}", e);
             }
@@ -122,7 +164,9 @@ impl Worker {
     }
 
     fn kill(&self) {
-        if let Err(e) = self.device_socket.send(&Request::Kill) {
+        if let Err(e) = self.device_socket.send(poly_msg_socket::Value::MsgOnSocket(
+            MsgOnSocketRequest::Kill,
+        )) {
             error!("failed to send Kill message: {}", e);
         }
     }
@@ -180,8 +224,8 @@ impl Controller {
         vm_socket: VmMemoryControlRequestSocket,
         resource_bridge: Option<ResourceRequestSocket>,
         socket: Socket,
-    ) -> Result<Controller, MsgError> {
-        socket.send(&Request::Create {
+    ) -> Result<Controller, poly_msg_socket::Error> {
+        socket.send(MsgOnSocketRequest::Create {
             // wayland_paths,
             vm_socket: MaybeOwnedFd::new_borrowed(&vm_socket),
             // resource_bridge,
@@ -267,7 +311,7 @@ impl VirtioDevice for Controller {
 
             let (ours, theirs) = UnixSeqpacket::pair().expect("pair failed");
 
-            if let Err(e) = socket.send(&Request::Activate {
+            if let Err(e) = socket.send(MsgOnSocketRequest::Activate {
                 shm: MaybeOwnedFd::new_borrowed(&mem),
                 interrupt: MaybeOwnedFd::new_borrowed(&theirs),
                 interrupt_resample_evt: MaybeOwnedFd::new_borrowed(interrupt.get_resample_evt()),
diff --git a/devices/src/virtio/virtio_pci_device.rs b/devices/src/virtio/virtio_pci_device.rs
index 0d44e48..7453e4c 100644
--- a/devices/src/virtio/virtio_pci_device.rs
+++ b/devices/src/virtio/virtio_pci_device.rs
@@ -398,6 +398,9 @@ impl PciDevice for VirtioPciDevice {
 
     fn keep_fds(&self) -> Vec<RawFd> {
         let mut fds = self.device.keep_fds();
+        fds.push(0);
+        fds.push(1);
+        fds.push(2);
         if let Some(interrupt_evt) = &self.interrupt_evt {
             fds.push(interrupt_evt.as_raw_fd());
         }