[service-vm] Send shutdown request to exit rialto

And rialto will be in a loop of processing request till receiving
the shutdown request.

Bug: 299432766
Test: atest rialto_test
Change-Id: Ifd65adde74a370a91fc462930bb513a3e1fe65b9
diff --git a/libs/service_vm_comm/src/lib.rs b/libs/service_vm_comm/src/lib.rs
index 3b53b63..ca97ca1 100644
--- a/libs/service_vm_comm/src/lib.rs
+++ b/libs/service_vm_comm/src/lib.rs
@@ -22,5 +22,7 @@
 mod message;
 mod vsock;
 
-pub use message::{EcdsaP256KeyPair, GenerateCertificateRequestParams, Request, Response};
+pub use message::{
+    EcdsaP256KeyPair, GenerateCertificateRequestParams, Request, Response, ServiceVmRequest,
+};
 pub use vsock::VmType;
diff --git a/libs/service_vm_comm/src/message.rs b/libs/service_vm_comm/src/message.rs
index 0eddcfb..80956cb 100644
--- a/libs/service_vm_comm/src/message.rs
+++ b/libs/service_vm_comm/src/message.rs
@@ -21,7 +21,19 @@
 
 type MacedPublicKey = Vec<u8>;
 
-/// Represents a request to be sent to the service VM.
+/// The main request type to be sent to the service VM.
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub enum ServiceVmRequest {
+    /// A request to be processed by the service VM.
+    ///
+    /// Each request has a corresponding response item.
+    Process(Request),
+
+    /// Shuts down the service VM. No response is expected from it.
+    Shutdown,
+}
+
+/// Represents a process request to be sent to the service VM.
 ///
 /// Each request has a corresponding response item.
 #[derive(Clone, Debug, Serialize, Deserialize)]
diff --git a/rialto/src/communication.rs b/rialto/src/communication.rs
index ee4ecdb..50722f2 100644
--- a/rialto/src/communication.rs
+++ b/rialto/src/communication.rs
@@ -20,7 +20,7 @@
 use core::mem;
 use core::result;
 use log::info;
-use service_vm_comm::{Request, Response};
+use service_vm_comm::{Response, ServiceVmRequest};
 use tinyvec::ArrayVec;
 use virtio_drivers::{
     self,
@@ -84,7 +84,7 @@
         }
     }
 
-    pub fn read_request(&mut self) -> Result<Request> {
+    pub fn read_request(&mut self) -> Result<ServiceVmRequest> {
         Ok(ciborium::from_reader(self)?)
     }
 
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index d777b2d..4e91574 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -33,7 +33,7 @@
 use hyp::{get_mem_sharer, get_mmio_guard};
 use libfdt::FdtError;
 use log::{debug, error, info};
-use service_vm_comm::VmType;
+use service_vm_comm::{ServiceVmRequest, VmType};
 use virtio_drivers::{
     device::socket::{VsockAddr, VMADDR_CID_HOST},
     transport::{pci::bus::PciRoot, DeviceType, Transport},
@@ -140,9 +140,11 @@
     debug!("Found socket device: guest cid = {:?}", socket_device.guest_cid());
 
     let mut vsock_stream = VsockStream::new(socket_device, host_addr())?;
-    let response = requests::process_request(vsock_stream.read_request()?)?;
-    vsock_stream.write_response(&response)?;
-    vsock_stream.flush()?;
+    while let ServiceVmRequest::Process(req) = vsock_stream.read_request()? {
+        let response = requests::process_request(req)?;
+        vsock_stream.write_response(&response)?;
+        vsock_stream.flush()?;
+    }
     vsock_stream.shutdown()?;
 
     Ok(())
diff --git a/rialto/tests/test.rs b/rialto/tests/test.rs
index 20d00b5..fc6e20a 100644
--- a/rialto/tests/test.rs
+++ b/rialto/tests/test.rs
@@ -55,8 +55,8 @@
     let message = "abc".repeat(166);
     let request = Request::Reverse(message.as_bytes().to_vec());
 
-    let response = vm.process_request(&request)?;
-    info!("Received response '{response:?}' for the request '{request:?}'.");
+    let response = vm.process_request(request)?;
+    info!("Received response: {response:?}.");
 
     let expected_response: Vec<u8> = message.as_bytes().iter().rev().cloned().collect();
     assert_eq!(Response::Reverse(expected_response), response);
diff --git a/service_vm_manager/src/lib.rs b/service_vm_manager/src/lib.rs
index c27570c..a645202 100644
--- a/service_vm_manager/src/lib.rs
+++ b/service_vm_manager/src/lib.rs
@@ -24,16 +24,16 @@
     },
     binder::ParcelFileDescriptor,
 };
-use anyhow::{ensure, Context, Result};
+use anyhow::{anyhow, ensure, Context, Result};
 use log::{info, warn};
-use service_vm_comm::{Request, Response, VmType};
+use service_vm_comm::{Request, Response, ServiceVmRequest, VmType};
 use std::fs::{File, OpenOptions};
 use std::io::{self, BufRead, BufReader, BufWriter, Write};
 use std::os::unix::io::FromRawFd;
 use std::path::{Path, PathBuf};
 use std::thread;
 use std::time::Duration;
-use vmclient::VmInstance;
+use vmclient::{DeathReason, VmInstance};
 use vsock::{VsockListener, VsockStream, VMADDR_CID_HOST};
 
 const VIRT_DATA_DIR: &str = "/data/misc/apexdata/com.android.virt";
@@ -90,13 +90,13 @@
     }
 
     /// Processes the request in the service VM.
-    pub fn process_request(&mut self, request: &Request) -> Result<Response> {
-        self.write_request(request)?;
+    pub fn process_request(&mut self, request: Request) -> Result<Response> {
+        self.write_request(&ServiceVmRequest::Process(request))?;
         self.read_response()
     }
 
     /// Sends the request to the service VM.
-    fn write_request(&mut self, request: &Request) -> Result<()> {
+    fn write_request(&mut self, request: &ServiceVmRequest) -> Result<()> {
         let mut buffer = BufWriter::with_capacity(WRITE_BUFFER_CAPACITY, &mut self.vsock_stream);
         ciborium::into_writer(request, &mut buffer)?;
         buffer.flush().context("Failed to flush the buffer")?;
@@ -111,14 +111,22 @@
         info!("Received response from the service VM.");
         Ok(response)
     }
+
+    /// Shuts down the service VM.
+    fn shutdown(&mut self) -> Result<DeathReason> {
+        self.write_request(&ServiceVmRequest::Shutdown)?;
+        self.vm
+            .wait_for_death_with_timeout(Duration::from_secs(10))
+            .ok_or_else(|| anyhow!("Timed out to exit the service VM"))
+    }
 }
 
 impl Drop for ServiceVm {
     fn drop(&mut self) {
         // Wait till the service VM finishes releasing all the resources.
-        match self.vm.wait_for_death_with_timeout(Duration::from_secs(10)) {
-            Some(e) => info!("Exit the service VM: {e:?}"),
-            None => warn!("Timed out waiting for service VM exit"),
+        match self.shutdown() {
+            Ok(reason) => info!("Exit the service VM successfully: {reason:?}"),
+            Err(e) => warn!("Service VM shutdown request failed '{e:?}', killing it."),
         }
     }
 }
diff --git a/virtualizationservice/src/rkpvm.rs b/virtualizationservice/src/rkpvm.rs
index 2c9230b..dbadd60 100644
--- a/virtualizationservice/src/rkpvm.rs
+++ b/virtualizationservice/src/rkpvm.rs
@@ -26,7 +26,7 @@
     // TODO(b/271275206): Send the correct request type with client VM's
     // information to be attested.
     let request = Request::Reverse(csr.to_vec());
-    match vm.process_request(&request).context("Failed to process request")? {
+    match vm.process_request(request).context("Failed to process request")? {
         Response::Reverse(cert) => Ok(cert),
         _ => bail!("Incorrect response type"),
     }