Shutdown Feature (#41)

* Feat: add initramfs implementation for vmm (#34) * feat(vmm): implemented automatic generation of rootfs with initramfs Signed-off-by: Muriel Paraire <[email protected]> * feat: image generation based off language Signed-off-by: Muriel Paraire <[email protected]> * feat(vmm): implemented automatic generation of rootfs with initramfs Signed-off-by: Muriel Paraire <[email protected]> * fix(vmm): fix logging & language order Signed-off-by: Muriel Paraire <[email protected]> * feat(vmm): one image per language Signed-off-by: Muriel Paraire <[email protected]> * feat(vmm): implemented initramfs Signed-off-by: Muriel Paraire <[email protected]> * fix(vmm): code cleanup Signed-off-by: Muriel Paraire <[email protected]> * fix(vmm): code cleanup Signed-off-by: Muriel Paraire <[email protected]> * fix(vmm): code cleanup Signed-off-by: Muriel Paraire <[email protected]> * fix: rust export for cargo agent and increase MMIO_GAP_END Signed-off-by: Mauran <[email protected]> * chore: lint Signed-off-by: Mauran <[email protected]> * fix: add back tracing Signed-off-by: Mauran <[email protected]> --------- Signed-off-by: Muriel Paraire <[email protected]> Signed-off-by: Mauran <[email protected]> Co-authored-by: Mauran <[email protected]> Signed-off-by: Font Vincent <[email protected]> * feat: shutdown vm Signed-off-by: Font Vincent <[email protected]> * docs: add readme (#38) * docs: add readme Signed-off-by: Mauran <[email protected]> --------- Signed-off-by: Mauran <[email protected]> Signed-off-by: Matéo Fernandez <[email protected]> Co-authored-by: Matéo Fernandez <[email protected]> Co-authored-by: Alexandre Sollier <[email protected]> Signed-off-by: Font Vincent <[email protected]> * remove(vmm): useless crate Signed-off-by: Font Vincent <[email protected]> * fix: lint Signed-off-by: Font Vincent <[email protected]> * fix: fmt Signed-off-by: Font Vincent <[email protected]> * fix: lint Signed-off-by: Font Vincent <[email protected]> --------- Signed-off-by: Muriel Paraire <[email protected]> Signed-off-by: Mauran <[email protected]> Signed-off-by: Font Vincent <[email protected]> Signed-off-by: Matéo Fernandez <[email protected]> Co-authored-by: Muriel Paraire <[email protected]> Co-authored-by: Mauran <[email protected]> Co-authored-by: Thomas Mauran <[email protected]> Co-authored-by: Matéo Fernandez <[email protected]> Co-authored-by: Alexandre Sollier <[email protected]>
virt-do · May 3, 2024 · 8ea1568 · 8ea1568
1 parent 1e6e7c7
commit 8ea1568
Show file tree

Hide file tree

Showing 9 changed files with 142 additions and 9 deletions.
diff --git a/proto/vmm.proto b/proto/vmm.proto
@@ -32,6 +32,7 @@ message ExecuteResponse {
 }
 
 service VmmService {
+  rpc Shutdown (ShutdownVmRequest) returns (ShutdownVmResponse) {};
   rpc Run (RunVmmRequest) returns (stream ExecuteResponse) {};
 }
 
@@ -45,3 +46,10 @@ message RunVmmRequest {
 
 message RunVmmResponse {
 }
+
+message ShutdownVmRequest {
+}
+
+message ShutdownVmResponse {
+  bool success = 1;
+}
diff --git a/src/api/src/client.rs b/src/api/src/client.rs
@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use tonic::{transport::Channel, Streaming};
 use vmmorchestrator::vmm_service_client::VmmServiceClient;
 
@@ -27,4 +29,17 @@ impl VmmClient {
 
         Ok(response_stream)
     }
+
+    pub async fn shutdown_vm(
+        &mut self,
+        request: vmmorchestrator::ShutdownVmRequest,
+    ) -> Result<vmmorchestrator::ShutdownVmResponse, tonic::Status> {
+        let mut request = tonic::Request::new(request);
+        request.set_timeout(Duration::from_secs(5));
+        let response = self.client.shutdown(request).await?.into_inner();
+
+        println!("shutdown response: {:?}", response);
+
+        Ok(response)
+    }
 }
diff --git a/src/api/src/service.rs b/src/api/src/service.rs
@@ -1,8 +1,8 @@
 use crate::client::{
-    vmmorchestrator::{ExecuteResponse, RunVmmRequest},
+    vmmorchestrator::{ExecuteResponse, RunVmmRequest, ShutdownVmRequest, ShutdownVmResponse},
     VmmClient,
 };
-use actix_web::{post, web, HttpResponse, Responder};
+use actix_web::{post, web, HttpRequest, HttpResponse, Responder};
 use actix_web_lab::sse;
 use async_stream::stream;
 use serde::Serialize;
@@ -65,7 +65,37 @@ impl From<ExecuteResponse> for ExecuteJsonResponse {
 }
 
 #[post("/shutdown")]
-pub async fn shutdown(req_body: String) -> impl Responder {
-    // TODO: Get the id from the body and shutdown the vm
-    HttpResponse::Ok().body(req_body)
+pub async fn shutdown(request: HttpRequest) -> impl Responder {
+    let req = request;
+
+    let mut client = VmmClient::new().await.unwrap();
+
+    println!("Request: {:?}", req);
+
+    let shutdown_request = ShutdownVmRequest {};
+    let response_result = client.shutdown_vm(shutdown_request).await;
+
+    match response_result {
+        Ok(response) => {
+            let json_response: ShutdownJsonResponse = response.into();
+            HttpResponse::Ok().body(serde_json::to_string(&json_response).unwrap())
+        }
+        Err(_) => {
+            let json_response: ShutdownJsonResponse = ShutdownJsonResponse { success: false };
+            HttpResponse::Ok().body(serde_json::to_string(&json_response).unwrap())
+        }
+    }
+}
+
+#[derive(Debug, Serialize)]
+pub struct ShutdownJsonResponse {
+    pub success: bool,
+}
+
+impl From<ShutdownVmResponse> for ShutdownJsonResponse {
+    fn from(value: ShutdownVmResponse) -> Self {
+        Self {
+            success: value.success,
+        }
+    }
 }
diff --git a/src/cli/src/args.rs b/src/cli/src/args.rs
@@ -14,4 +14,5 @@ pub enum Commands {
         #[arg(short, long)]
         config_path: PathBuf,
     },
+    Shutdown {},
 }
diff --git a/src/cli/src/main.rs b/src/cli/src/main.rs
@@ -30,6 +30,19 @@ async fn main() -> io::Result<()> {
                 Err(e) => eprintln!("Error while making the request: {}", e),
             }
         }
+        Commands::Shutdown {} => {
+            let response = CloudletClient::shutdown().await;
+            match response {
+                Ok(bool) => {
+                    if bool {
+                        println!("Shutdown Request successful !")
+                    } else {
+                        println!("Shutdown Request Failed")
+                    }
+                }
+                Err(()) => println!("Cannot send shutdown Request"),
+            }
+        }
     }
 
     Ok(())

diff --git a/src/cli/src/services.rs b/src/cli/src/services.rs
@@ -1,7 +1,9 @@
 use crate::utils::ConfigFileHandler;
 use reqwest::Client;
 use serde::Deserialize;
-use shared_models::{BuildConfig, CloudletDtoRequest, Language, ServerConfig};
+use shared_models::{
+    BuildConfig, CloudletDtoRequest, CloudletShutdownResponse, Language, ServerConfig,
+};
 use std::error::Error;
 
 #[derive(Deserialize, Debug)]
@@ -50,4 +52,17 @@ impl CloudletClient {
         println!("Response: {:?}", res.text().await?);
         Ok(())
     }
+
+    pub async fn shutdown() -> Result<bool, ()> {
+        let client = Client::new();
+        let response = client.post("http://127.0.0.1:3000/shutdown").send().await;
+
+        let shutdown_response: CloudletShutdownResponse = response
+            .unwrap()
+            .json::<CloudletShutdownResponse>()
+            .await
+            .unwrap();
+
+        Ok(shutdown_response.success)
+    }
 }
diff --git a/src/shared-models/src/lib.rs b/src/shared-models/src/lib.rs
@@ -39,6 +39,11 @@ pub struct CloudletDtoRequest {
     pub build: BuildConfig,
 }
 
+#[derive(Debug, Deserialize)]
+pub struct CloudletShutdownResponse {
+    pub success: bool,
+}
+
 #[derive(Serialize, Deserialize, Debug)]
 pub struct ServerConfig {
     pub address: String,

diff --git a/src/vmm/src/grpc/client.rs b/src/vmm/src/grpc/client.rs
@@ -1,6 +1,7 @@
-use self::agent::{workload_runner_client::WorkloadRunnerClient, ExecuteRequest};
+use self::agent::{workload_runner_client::WorkloadRunnerClient, ExecuteRequest, SignalRequest};
+use super::server::vmmorchestrator::{ShutdownVmRequest, ShutdownVmResponse};
 use log::error;
-use std::{net::Ipv4Addr, time::Duration};
+use std::{error::Error, net::Ipv4Addr, time::Duration};
 use tonic::{transport::Channel, Streaming};
 
 pub mod agent {
@@ -37,4 +38,23 @@ impl WorkloadClient {
 
         Ok(response_stream)
     }
+
+    pub async fn shutdown(
+        &mut self,
+        _request: ShutdownVmRequest,
+    ) -> Result<ShutdownVmResponse, tonic::Status> {
+        const BROKEN_PIPE_ERROR: &str = "stream closed because of a broken pipe";
+
+        let signal_request = SignalRequest::default();
+        let response = self.client.signal(signal_request).await;
+
+        if let Err(status) = response {
+            let error = status.source().unwrap().source().unwrap().source().unwrap();
+            if error.to_string().as_str().eq(BROKEN_PIPE_ERROR) {
+                return Ok(ShutdownVmResponse { success: true });
+            }
+        }
+
+        Ok(ShutdownVmResponse { success: false })
+    }
 }
diff --git a/src/vmm/src/grpc/server.rs b/src/vmm/src/grpc/server.rs
@@ -1,5 +1,6 @@
 use self::vmmorchestrator::{
-    vmm_service_server::VmmService as VmmServiceTrait, Language, RunVmmRequest,
+    vmm_service_server::VmmService as VmmServiceTrait, Language, RunVmmRequest, ShutdownVmRequest,
+    ShutdownVmResponse,
 };
 use crate::grpc::client::agent::ExecuteRequest;
 use crate::VmmErrors;
@@ -154,6 +155,31 @@ impl VmmServiceTrait for VmmService {
     type RunStream =
         ReceiverStream<std::result::Result<vmmorchestrator::ExecuteResponse, tonic::Status>>;
 
+    async fn shutdown(&self, request: Request<ShutdownVmRequest>) -> Result<ShutdownVmResponse> {
+        const GUEST_IP: Ipv4Addr = Ipv4Addr::new(172, 29, 0, 2);
+
+        let grpc_client = tokio::spawn(async move {
+            // Wait 2 seconds
+            tokio::time::sleep(Duration::from_secs(2)).await;
+            println!("Connecting to Agent service");
+
+            WorkloadClient::new(GUEST_IP, 50051).await
+        })
+        .await
+        .unwrap();
+
+        if let Ok(mut client) = grpc_client {
+            info!("Attempting to shutdown the VM...");
+
+            let response = client.shutdown(request.into_inner()).await.unwrap();
+
+            return Ok(Response::new(response));
+        } else if let Err(e) = grpc_client {
+            error!("ERROR {:?}", e);
+        }
+        return Err(Status::internal("Failed to shutdown the VM"));
+    }
+
     async fn run(&self, request: Request<RunVmmRequest>) -> Result<Self::RunStream> {
         let (tx, rx) = tokio::sync::mpsc::channel(4);