vmgs: Make VmgsClient MeshPayload (#2465)

As part of my work to support running ChipsetDevices in another process,
I ran into needing certain resources to be MeshPayload. The VmgsClient
is the second of these.

---------

Co-authored-by: Copilot <[email protected]>

Author: smalis-msft

Cargo.lock

@@ -9261,7 +9261,7 @@ version = "0.0.0"
 dependencies = [
  "async-trait",
  "inspect",
- "mesh_channel",
+ "mesh",
  "pal_async",
  "thiserror 2.0.16",
  "tracing",

vm/vmgs/vmgs/Cargo.toml

@@ -11,6 +11,7 @@ default = []
 
 inspect = ["vmgs_format/inspect", "dep:inspect", "dep:inspect_counters"]
 save_restore = ["dep:mesh_protobuf"]
+mesh = ["dep:mesh_protobuf"]
 
 # Use native windows crypto APIs
 encryption_win = ["dep:windows"]

vm/vmgs/vmgs/src/vmgs_impl.rs

@@ -47,6 +47,7 @@ enum LogOpType {
 
 /// Info about a specific VMGS file.
 #[derive(Debug)]
+#[cfg_attr(feature = "mesh", derive(mesh_protobuf::Protobuf))]
 pub struct VmgsFileInfo {
     /// Number of bytes allocated in the file.
     pub allocated_bytes: u64,

vm/vmgs/vmgs_broker/Cargo.toml

@@ -17,12 +17,12 @@ encryption_ossl = ["vmgs/encryption_ossl"]
 [dependencies]
 vmgs_format.workspace = true
 vmgs_resources.workspace = true
-vmgs = { workspace = true, features = ["inspect", "save_restore"] }
+vmgs = { workspace = true, features = ["inspect", "save_restore", "mesh"] }
 vm_resource.workspace = true
 vmcore.workspace = true
 
 inspect = { workspace = true, features = ["defer"] }
-mesh_channel.workspace = true
+mesh.workspace = true
 pal_async.workspace = true
 
 async-trait.workspace = true

vm/vmgs/vmgs_broker/src/broker.rs

@@ -1,19 +1,56 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-use mesh_channel::Receiver;
-use mesh_channel::rpc::Rpc;
+use mesh::MeshPayload;
+use mesh::Receiver;
+use mesh::error::RemoteError;
+use mesh::payload::Protobuf;
+use mesh::rpc::Rpc;
+use thiserror::Error;
 use vmgs::Vmgs;
 use vmgs::VmgsFileInfo;
 use vmgs_format::FileId;
 
+#[derive(Protobuf, Error, Debug)]
+pub enum VmgsBrokerError {
+    #[error("no allocated bytes for file id being read")]
+    FileInfoNotAllocated,
+    #[error(transparent)]
+    Other(RemoteError),
+}
+
+impl From<vmgs::Error> for VmgsBrokerError {
+    fn from(value: vmgs::Error) -> Self {
+        match value {
+            vmgs::Error::FileInfoNotAllocated => VmgsBrokerError::FileInfoNotAllocated,
+            other => VmgsBrokerError::Other(RemoteError::new(other)),
+        }
+    }
+}
+
+#[derive(Protobuf)]
+pub struct BrokerFileId(u32);
+
+impl From<FileId> for BrokerFileId {
+    fn from(value: FileId) -> Self {
+        BrokerFileId(value.0)
+    }
+}
+
+impl From<BrokerFileId> for FileId {
+    fn from(value: BrokerFileId) -> Self {
+        FileId(value.0)
+    }
+}
+
+#[derive(MeshPayload)]
 pub enum VmgsBrokerRpc {
     Inspect(inspect::Deferred),
-    GetFileInfo(Rpc<FileId, Result<VmgsFileInfo, vmgs::Error>>),
-    ReadFile(Rpc<FileId, Result<Vec<u8>, vmgs::Error>>),
-    WriteFile(Rpc<(FileId, Vec<u8>), Result<(), vmgs::Error>>),
+    GetFileInfo(Rpc<BrokerFileId, Result<VmgsFileInfo, VmgsBrokerError>>),
+    ReadFile(Rpc<BrokerFileId, Result<Vec<u8>, VmgsBrokerError>>),
+    WriteFile(Rpc<(BrokerFileId, Vec<u8>), Result<(), VmgsBrokerError>>),
     #[cfg(with_encryption)]
-    WriteFileEncrypted(Rpc<(FileId, Vec<u8>), Result<(), vmgs::Error>>),
+    WriteFileEncrypted(Rpc<(BrokerFileId, Vec<u8>), Result<(), VmgsBrokerError>>),
     Save(Rpc<(), vmgs::save_restore::state::SavedVmgsState>),
 }
 
@@ -41,21 +78,33 @@ impl VmgsBrokerTask {
             VmgsBrokerRpc::Inspect(req) => {
                 req.inspect(&self.vmgs);
             }
-            VmgsBrokerRpc::GetFileInfo(rpc) => {
-                rpc.handle_sync(|file_id| self.vmgs.get_file_info(file_id))
-            }
+            VmgsBrokerRpc::GetFileInfo(rpc) => rpc
+                .handle_sync(|file_id| self.vmgs.get_file_info(file_id.into()).map_err(Into::into)),
             VmgsBrokerRpc::ReadFile(rpc) => {
-                rpc.handle(async |file_id| self.vmgs.read_file(file_id).await)
-                    .await
+                rpc.handle(async |file_id| {
+                    self.vmgs
+                        .read_file(file_id.into())
+                        .await
+                        .map_err(Into::into)
+                })
+                .await
             }
             VmgsBrokerRpc::WriteFile(rpc) => {
-                rpc.handle(async |(file_id, buf)| self.vmgs.write_file(file_id, &buf).await)
-                    .await
+                rpc.handle(async |(file_id, buf)| {
+                    self.vmgs
+                        .write_file(file_id.into(), &buf)
+                        .await
+                        .map_err(Into::into)
+                })
+                .await
             }
             #[cfg(with_encryption)]
             VmgsBrokerRpc::WriteFileEncrypted(rpc) => {
                 rpc.handle(async |(file_id, buf)| {
-                    self.vmgs.write_file_encrypted(file_id, &buf).await
+                    self.vmgs
+                        .write_file_encrypted(file_id.into(), &buf)
+                        .await
+                        .map_err(Into::into)
                 })
                 .await
             }

vm/vmgs/vmgs_broker/src/client.rs

@@ -4,10 +4,12 @@
 //! The Vmgs worker will send messages to the Vmgs dispatch, allowing
 //! tasks to queue for the dispatcher to handle synchronously
 
+use crate::broker::VmgsBrokerError;
 use crate::broker::VmgsBrokerRpc;
 use inspect::Inspect;
-use mesh_channel::rpc::RpcError;
-use mesh_channel::rpc::RpcSend;
+use mesh::MeshPayload;
+use mesh::rpc::RpcError;
+use mesh::rpc::RpcSend;
 use thiserror::Error;
 use tracing::instrument;
 use vmgs::VmgsFileInfo;
@@ -19,14 +21,22 @@ use vmgs_format::FileId;
 pub enum VmgsClientError {
     /// VMGS broker is offline
     #[error("broker is offline")]
-    BrokerOffline(#[from] RpcError),
+    BrokerOffline(#[source] RpcError),
     /// VMGS error
     #[error("vmgs error")]
-    Vmgs(#[from] vmgs::Error),
+    Vmgs(#[source] VmgsBrokerError),
 }
 
-impl From<RpcError<vmgs::Error>> for VmgsClientError {
-    fn from(value: RpcError<vmgs::Error>) -> Self {
+impl From<RpcError> for VmgsClientError {
+    fn from(value: RpcError) -> Self {
+        match value {
+            RpcError::Channel(e) => VmgsClientError::BrokerOffline(RpcError::Channel(e)),
+        }
+    }
+}
+
+impl From<RpcError<VmgsBrokerError>> for VmgsClientError {
+    fn from(value: RpcError<VmgsBrokerError>) -> Self {
         match value {
             RpcError::Call(e) => VmgsClientError::Vmgs(e),
             RpcError::Channel(e) => VmgsClientError::BrokerOffline(RpcError::Channel(e)),
@@ -35,10 +45,10 @@ impl From<RpcError<vmgs::Error>> for VmgsClientError {
 }
 
 /// Client to interact with a backend-agnostic VMGS instance.
-#[derive(Clone, Inspect)]
+#[derive(Clone, Inspect, MeshPayload)]
 pub struct VmgsClient {
     #[inspect(flatten, send = "VmgsBrokerRpc::Inspect")]
-    pub(crate) control: mesh_channel::Sender<VmgsBrokerRpc>,
+    pub(crate) control: mesh::Sender<VmgsBrokerRpc>,
 }
 
 impl VmgsClient {
@@ -47,7 +57,7 @@ impl VmgsClient {
     pub async fn get_file_info(&self, file_id: FileId) -> Result<VmgsFileInfo, VmgsClientError> {
         let res = self
             .control
-            .call_failable(VmgsBrokerRpc::GetFileInfo, file_id)
+            .call_failable(VmgsBrokerRpc::GetFileInfo, file_id.into())
             .await?;
 
         Ok(res)
@@ -58,7 +68,7 @@ impl VmgsClient {
     pub async fn read_file(&self, file_id: FileId) -> Result<Vec<u8>, VmgsClientError> {
         let res = self
             .control
-            .call_failable(VmgsBrokerRpc::ReadFile, file_id)
+            .call_failable(VmgsBrokerRpc::ReadFile, file_id.into())
             .await?;
 
         Ok(res)
@@ -71,7 +81,7 @@ impl VmgsClient {
     #[instrument(skip_all, fields(file_id))]
     pub async fn write_file(&self, file_id: FileId, buf: Vec<u8>) -> Result<(), VmgsClientError> {
         self.control
-            .call_failable(VmgsBrokerRpc::WriteFile, (file_id, buf))
+            .call_failable(VmgsBrokerRpc::WriteFile, (file_id.into(), buf))
             .await?;
 
         Ok(())
@@ -88,7 +98,7 @@ impl VmgsClient {
         buf: Vec<u8>,
     ) -> Result<(), VmgsClientError> {
         self.control
-            .call_failable(VmgsBrokerRpc::WriteFileEncrypted, (file_id, buf))
+            .call_failable(VmgsBrokerRpc::WriteFileEncrypted, (file_id.into(), buf))
             .await?;
 
         Ok(())

vm/vmgs/vmgs_broker/src/lib.rs

@@ -20,7 +20,7 @@ use pal_async::task::Task;
 /// Given a fully-initialized VMGS instance, return a VMGS broker task +
 /// clonable VmgsClient
 pub fn spawn_vmgs_broker(spawner: impl Spawn, vmgs: vmgs::Vmgs) -> (VmgsClient, Task<()>) {
-    let (control_send, control_recv) = mesh_channel::mpsc_channel();
+    let (control_send, control_recv) = mesh::mpsc_channel();
 
     let process_loop_handle = spawner.spawn("vmgs-broker", async move {
         VmgsBrokerTask::new(vmgs).run(control_recv).await

vm/vmgs/vmgs_broker/src/non_volatile_store.rs

@@ -5,6 +5,7 @@
 
 use crate::VmgsClient;
 use crate::VmgsClientError;
+use crate::broker::VmgsBrokerError;
 use async_trait::async_trait;
 use thiserror::Error;
 use vmcore::non_volatile_store::NonVolatileStore;
@@ -67,7 +68,7 @@ impl NonVolatileStore for VmgsNonVolatileStore {
     async fn restore(&mut self) -> Result<Option<Vec<u8>>, NonVolatileStoreError> {
         match self.vmgs.read_file(self.file_id).await {
             Ok(buf) => Ok(Some(buf)),
-            Err(VmgsClientError::Vmgs(vmgs::Error::FileInfoNotAllocated)) => Ok(None),
+            Err(VmgsClientError::Vmgs(VmgsBrokerError::FileInfoNotAllocated)) => Ok(None),
             Err(e) => Err(NonVolatileStoreError::new(e)),
         }
     }