VM updates working

.env

@@ -3,3 +3,5 @@ DB_USER = "root"
 DB_PASS = "root"
 DB_NAMESPACE = "brain"
 DB_NAME = "migration"
+CERT_PATH = "./tmp/brain-crt.pem"
+CERT_KEY_PATH = "./tmp/brain-key.pem"

interim_tables.surql

@@ -60,6 +60,8 @@ DEFINE FIELD dtrfs_sha ON TABLE update_vm_req TYPE string;
 DEFINE FIELD dtrfs_url ON TABLE update_vm_req TYPE string;
 DEFINE FIELD kernel_sha ON TABLE update_vm_req TYPE string;
 DEFINE FIELD kernel_url ON TABLE update_vm_req TYPE string;
+DEFINE FIELD created_at ON TABLE update_vm_req TYPE datetime;
+DEFINE FIELD error ON TABLE update_vm_req TYPE string;
 
 DEFINE TABLE deleted_vm TYPE RELATION FROM account TO vm_node SCHEMAFULL;
 DEFINE FIELD hostname ON TABLE deleted_vm TYPE string;

src/bin/brain.rs

@@ -12,14 +12,16 @@ use tonic::transport::{Identity, Server, ServerTlsConfig};
 
 #[tokio::main]
 async fn main() {
+    if dotenv::from_filename("/etc/detee/brain/config.ini").is_err() {
         dotenv().ok();
+    }
     env_logger::builder().filter_level(log::LevelFilter::Debug).init();
 
-    let db_url = std::env::var("DB_URL").expect("DB_URL not set in .env");
+    let db_url = std::env::var("DB_URL").expect("the environment variable DB_URL is not set");
-    let db_user = std::env::var("DB_USER").expect("DB_USER not set in .env");
+    let db_user = std::env::var("DB_USER").expect("the environment variable DB_USER is not set");
-    let db_pass = std::env::var("DB_PASS").expect("DB_PASS not set in .env");
+    let db_pass = std::env::var("DB_PASS").expect("the environment variable DB_PASS is not set");
-    let db_ns = std::env::var("DB_NAMESPACE").expect("DB_NAMESPACE not set in .env");
+    let db_ns = std::env::var("DB_NAMESPACE").expect("the env variable DB_NAMESPACE is not set");
-    let db_name = std::env::var("DB_NAME").expect("DB_NAME not set in .env");
+    let db_name = std::env::var("DB_NAME").expect("the environment variable DB_NAME is not set");
 
     let db = db::db_connection(&db_url, &db_user, &db_pass, &db_ns, &db_name).await.unwrap();
     let db_arc = Arc::new(db);
@@ -30,8 +32,10 @@ async fn main() {
     let snp_cli_server = BrainVmCliServer::new(VmCliServer::new(db_arc.clone()));
     let general_service_server = BrainGeneralCliServer::new(GeneralCliServer::new(db_arc.clone()));
 
-    let cert = std::fs::read_to_string(CERT_PATH).unwrap();
+    let cert_path = std::env::var("CERT_PATH").unwrap_or(CERT_PATH.to_string());
-    let key = std::fs::read_to_string(CERT_KEY_PATH).unwrap();
+    let key_path = std::env::var("CERT_KEY_PATH").unwrap_or(CERT_KEY_PATH.to_string());
+    let cert = std::fs::read_to_string(cert_path).unwrap();
+    let key = std::fs::read_to_string(key_path).unwrap();
 
     let identity = Identity::from_pem(cert, key);
 

src/constants.rs

@@ -1,6 +1,7 @@
 pub const BRAIN_GRPC_ADDR: &str = "0.0.0.0:31337";
-pub const CERT_PATH: &str = "./tmp/brain-crt.pem";
+pub const CERT_PATH: &str = "/etc/detee/brain/brain-crt.pem";
-pub const CERT_KEY_PATH: &str = "./tmp/brain-key.pem";
+pub const CERT_KEY_PATH: &str = "/etc/detee/brain/brain-key.pem";
+pub const CONFIG_PATH: &str = "/etc/detee/brain/config.ini";
 
 pub const DB_SCHEMA_FILE: &str = "interim_tables.surql";
 
@@ -15,6 +16,7 @@ pub const OLD_BRAIN_DATA_PATH: &str = "./saved_data.yaml";
 pub const ACCOUNT: &str = "account";
 pub const VM_NODE: &str = "vm_node";
 pub const ACTIVE_VM: &str = "active_vm";
+pub const VM_UPDATE_EVENT: &str = "vm_update_event";
 pub const NEW_VM_REQ: &str = "new_vm_req";
 pub const UPDATE_VM_REQ: &str = "update_vm_req";
 pub const DELETED_VM: &str = "deleted_vm";

src/db/mod.rs

@@ -29,6 +29,7 @@ pub mod prelude {
     pub use super::general::*;
     pub use super::vm::*;
     pub use super::*;
+    pub use detee_shared::snp::pb::vm_proto::{MeasurementArgs, VmNodeFilters};
 }
 
 pub async fn db_connection(
@@ -78,12 +79,12 @@ pub async fn upsert_record<SomeRecord: Serialize + 'static>(
 ) -> Result<(), Error> {
     #[derive(Deserialize)]
     struct Wrapper {}
-    let _: Option<Wrapper> = db.create((table, id)).content(my_record).await?;
+    let _: Option<Wrapper> = db.upsert((table, id)).content(my_record).await?;
     Ok(())
 }
 
 pub async fn live_vmnode_msgs<
-    T: Into<vm::VmDaemonMsg> + std::marker::Unpin + for<'de> Deserialize<'de>,
+    T: std::fmt::Debug + Into<vm::VmDaemonMsg> + std::marker::Unpin + for<'de> Deserialize<'de>,
 >(
     db: &Surreal<Client>,
     node: &str,
@@ -104,12 +105,13 @@ pub async fn live_vmnode_msgs<
     while let Some(result) = live_stream.next().await {
         match result {
             Ok(notification) => {
+                log::debug!("Got notification for node {node}: {notification:?}");
                 if notification.action == surrealdb::Action::Create {
                     tx.send(notification.data.into()).await?
                 }
             }
             Err(e) => {
-                log::error!("listen_for_{table_name} DB stream failed for {node}: {e}");
+                log::error!("live_vmnode_msgs for {table_name} DB stream failed for {node}: {e}");
                 return Err(Error::from(e));
             }
         }

src/db/vm.rs

@@ -2,8 +2,10 @@ use std::str::FromStr;
 use std::time::Duration;
 
 use super::Error;
-use crate::constants::{ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, VM_NODE};
+use crate::constants::{
-use crate::db::general::Report;
+    ACCOUNT, ACTIVE_VM, DELETED_VM, NEW_VM_REQ, UPDATE_VM_REQ, VM_NODE, VM_UPDATE_EVENT,
+};
+use crate::db::{MeasurementArgs, Report, VmNodeFilters};
 use crate::old_brain;
 use serde::{Deserialize, Serialize};
 use surrealdb::engine::remote::ws::Client;
@@ -80,7 +82,7 @@ impl VmNodeWithReports {
     // https://en.wikipedia.org/wiki/Dependency_inversion_principle
     pub async fn find_by_filters(
         db: &Surreal<Client>,
-        filters: detee_shared::snp::pb::vm_proto::VmNodeFilters,
+        filters: VmNodeFilters,
     ) -> Result<Vec<Self>, Error> {
         let mut query = format!(
             "select *, <-report.* as reports from {VM_NODE} where
@@ -193,35 +195,54 @@ impl NewVmReq {
 }
 
 /// first string is the vm_id
-pub enum NewVmResp {
+pub enum WrappedMeasurement {
-    // TODO: find a more elegant way to do this than importing gRPC in the DB module
+    Args(String, MeasurementArgs),
-    // https://en.wikipedia.org/wiki/Dependency_inversion_principle
-    Args(String, detee_shared::snp::pb::vm_proto::MeasurementArgs),
     Error(String, String),
 }
 
-impl NewVmResp {
+impl WrappedMeasurement {
-    pub async fn listen(db: &Surreal<Client>, vm_id: &str) -> Result<NewVmResp, Error> {
+    /// table must be NEW_VM_REQ or UPDATE_VM_REQ
+    /// it will however be enforced if you send anything else
+    pub async fn listen(
+        db: &Surreal<Client>,
+        vm_id: &str,
+        table: &str,
+    ) -> Result<WrappedMeasurement, Error> {
+        let table = match table {
+            UPDATE_VM_REQ => UPDATE_VM_REQ,
+            _ => NEW_VM_REQ,
+        };
+        #[derive(Deserialize)]
+        struct ErrorMessage {
+            error: String,
+        }
         let mut resp = db
-            .query(format!("live select * from {NEW_VM_REQ} where id = {NEW_VM_REQ}:{vm_id};"))
+            .query(format!("live select error from {table} where id = {NEW_VM_REQ}:{vm_id};"))
             .query(format!(
                 "live select * from measurement_args where id = measurement_args:{vm_id};"
             ))
             .await?;
-        let mut new_vm_stream = resp.stream::<Notification<NewVmReq>>(0)?;
+        let mut error_stream = resp.stream::<Notification<ErrorMessage>>(0)?;
-        let mut args_stream =
+        let mut args_stream = resp.stream::<Notification<MeasurementArgs>>(1)?;
-            resp.stream::<Notification<detee_shared::snp::pb::vm_proto::MeasurementArgs>>(1)?;
+
+        let args: Option<MeasurementArgs> = db.delete(("measurement_args", vm_id)).await?;
+        if let Some(args) = args {
+            return Ok(Self::Args(vm_id.to_string(), args));
+        }
 
         tokio::time::timeout(Duration::from_secs(10), async {
             loop {
                 tokio::select! {
-                    new_vm_req_notif = new_vm_stream.next() => {
+                    error_notification = error_stream.next() => {
-                        log::debug!("Got stream 1...");
+                        if let Some(err_notif) = error_notification {
-                        if let Some(new_vm_req_notif) = new_vm_req_notif {
+                            match err_notif {
-                            match new_vm_req_notif {
+                                Ok(err_notif) => {
-                                Ok(new_vm_req_notif) => {
+                                    if err_notif.action == surrealdb::Action::Update
-                                    if new_vm_req_notif.action == surrealdb::Action::Update && !new_vm_req_notif.data.error.is_empty() {
+                                    && !err_notif.data.error.is_empty() {
-                                        return Ok::<NewVmResp, Error>(Self::Error(vm_id.to_string(), new_vm_req_notif.data.error));
+                                        return Ok::<WrappedMeasurement, Error>(
+                                            Self::Error(vm_id.to_string(),
+                                            err_notif.data.error)
+                                        );
                                     };
                                 },
                                 Err(e) => return Err(e.into()),
@@ -233,6 +254,7 @@ impl NewVmResp {
                             match args_notif {
                                 Ok(args_notif) => {
                                     if args_notif.action == surrealdb::Action::Create {
+                                        let _: Option<MeasurementArgs> = db.delete(("measurement_args", vm_id)).await?;
                                         return Ok(Self::Args(vm_id.to_string(), args_notif.data));
                                     };
                                 },
@@ -242,11 +264,12 @@ impl NewVmResp {
                     }
                 }
             }
-        }).await?
+        })
+        .await?
     }
 }
 
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct ActiveVm {
     pub id: RecordId,
     #[serde(rename = "in")]
@@ -269,10 +292,32 @@ pub struct ActiveVm {
 }
 
 impl ActiveVm {
+    /// total hardware units of this VM
+    fn total_units(&self) -> u64 {
+        // TODO: Optimize this based on price of hardware.
+        // I tried, but this can be done better.
+        // Storage cost should also be based on tier
+        (self.vcpus as u64 * 10)
+            + ((self.memory_mb + 256) as u64 / 200)
+            + (self.disk_size_gb as u64 / 10)
+            + (!self.public_ipv4.is_empty() as u64 * 10)
+    }
+
+    /// Returns price per minute in nanoLP
+    pub fn price_per_minute(&self) -> u64 {
+        self.total_units() * self.price_per_unit
+    }
+
+    pub async fn get_by_uuid(db: &Surreal<Client>, uuid: &str) -> Result<Option<Self>, Error> {
+        let contract: Option<Self> =
+            db.query(format!("select * from {ACTIVE_VM}:{uuid};")).await?.take(0)?;
+        Ok(contract)
+    }
+
     pub async fn activate(
         db: &Surreal<Client>,
         id: &str,
-        args: detee_shared::vm_proto::MeasurementArgs,
+        args: MeasurementArgs,
     ) -> Result<(), Error> {
         let new_vm_req = match NewVmReq::get(db, id).await? {
             Some(r) => r,
@@ -298,7 +343,7 @@ impl ActiveVm {
         let mut guest_ports = vec![22];
         guest_ports.append(&mut args.exposed_ports.clone());
         let mut i = 0;
-        while i < new_vm_req.extra_ports.len() && i < guest_ports.len() {
+        while i < args.exposed_ports.len() && i < guest_ports.len() {
             mapped_ports.push((args.exposed_ports[i], guest_ports[i]));
             i += 1;
         }
@@ -327,6 +372,55 @@ impl ActiveVm {
         NewVmReq::delete(db, id).await?;
         Ok(())
     }
+
+    pub async fn update(db: &Surreal<Client>, id: &str) -> Result<(), Error> {
+        let update_vm_req = match UpdateVmReq::get(db, id).await? {
+            Some(r) => r,
+            None => return Ok(()),
+        };
+
+        let mut active_vm = match Self::get_by_uuid(db, id).await? {
+            Some(vm) => vm,
+            None => return Ok(()),
+        };
+
+        if update_vm_req.vcpus > 0 {
+            active_vm.vcpus = update_vm_req.vcpus;
+        }
+        if update_vm_req.memory_mb > 0 {
+            active_vm.memory_mb = update_vm_req.memory_mb;
+        }
+        if update_vm_req.disk_size_gb > 0 {
+            active_vm.disk_size_gb = update_vm_req.disk_size_gb;
+        }
+        if !update_vm_req.dtrfs_sha.is_empty() && !update_vm_req.kernel_sha.is_empty() {
+            active_vm.dtrfs_sha = update_vm_req.dtrfs_sha;
+            active_vm.kernel_sha = update_vm_req.kernel_sha;
+        }
+
+        let _: Option<ActiveVm> = db.update(active_vm.id.clone()).content(active_vm).await?;
+        UpdateVmReq::delete(db, id).await?;
+        Ok(())
+    }
+
+    pub async fn change_hostname(
+        db: &Surreal<Client>,
+        id: &str,
+        new_hostname: &str,
+    ) -> Result<bool, Error> {
+        let contract: Option<Self> = db
+            .query(format!(
+                "UPDATE {ACTIVE_VM}:{id} SET hostname = '{new_hostname}' RETURN BEFORE;"
+            ))
+            .await?
+            .take(0)?;
+        if let Some(contract) = contract {
+            if contract.hostname != new_hostname {
+                return Ok(true);
+            }
+        }
+        Ok(false)
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -344,8 +438,94 @@ pub struct UpdateVmReq {
     pub kernel_sha: String,
     pub kernel_url: String,
     pub created_at: Datetime,
-    pub price_per_unit: u64,
+    pub error: String,
-    pub locked_nano: u64,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct UpdateVmEvent {
+    pub vm_id: RecordId,
+    #[serde(rename = "in")]
+    pub admin: RecordId,
+    #[serde(rename = "out")]
+    pub vm_node: RecordId,
+    pub disk_size_gb: u32,
+    pub vcpus: u32,
+    pub memory_mb: u32,
+    pub dtrfs_url: String,
+    pub dtrfs_sha: String,
+    pub kernel_sha: String,
+    pub kernel_url: String,
+    pub executed_at: Datetime,
+}
+
+impl From<UpdateVmReq> for UpdateVmEvent {
+    fn from(update_vm_req: UpdateVmReq) -> Self {
+        Self {
+            vm_id: RecordId::from((VM_UPDATE_EVENT, update_vm_req.id.key().to_string())),
+            admin: update_vm_req.admin,
+            vm_node: update_vm_req.vm_node,
+            disk_size_gb: update_vm_req.disk_size_gb,
+            vcpus: update_vm_req.vcpus,
+            memory_mb: update_vm_req.memory_mb,
+            dtrfs_url: update_vm_req.dtrfs_url,
+            dtrfs_sha: update_vm_req.dtrfs_sha,
+            kernel_sha: update_vm_req.kernel_sha,
+            kernel_url: update_vm_req.kernel_url,
+            executed_at: Datetime::default(),
+        }
+    }
+}
+
+impl UpdateVmReq {
+    pub async fn get(db: &Surreal<Client>, id: &str) -> Result<Option<Self>, Error> {
+        let update_vm_req: Option<Self> = db.select((UPDATE_VM_REQ, id)).await?;
+        Ok(update_vm_req)
+    }
+
+    pub async fn delete(db: &Surreal<Client>, id: &str) -> Result<(), Error> {
+        let update_vm_req: Option<Self> = db.delete((UPDATE_VM_REQ, id)).await?;
+        if let Some(update_vm_req) = update_vm_req {
+            let update_vm_event: UpdateVmEvent = update_vm_req.into();
+            let _: Option<UpdateVmEvent> =
+                db.create(VM_UPDATE_EVENT).content(update_vm_event).await?;
+        }
+        Ok(())
+    }
+
+    /// returns None if VM does not exist
+    /// returns Some(false) if hw update is not needed
+    /// returns Some(true) if hw update is needed and got submitted
+    /// returns error if something happened with the DB
+    pub async fn request_hw_update(mut self, db: &Surreal<Client>) -> Result<Option<bool>, Error> {
+        let contract = ActiveVm::get_by_uuid(db, &self.id.key().to_string()).await?;
+
+        if contract.is_none() {
+            return Ok(None);
+        }
+        let contract = contract.unwrap();
+        // this is needed cause TryFrom does not support await
+        self.vm_node = contract.vm_node;
+
+        if !((self.vcpus != 0 && contract.vcpus != self.vcpus)
+            || (self.memory_mb != 0 && contract.memory_mb != self.memory_mb)
+            || (!self.dtrfs_sha.is_empty() && contract.dtrfs_sha != self.dtrfs_sha)
+            || (self.disk_size_gb != 0 && contract.disk_size_gb != self.disk_size_gb))
+        {
+            return Ok(Some(false));
+        }
+
+        let _: Vec<Self> = db.insert(UPDATE_VM_REQ).relation(self).await?;
+        Ok(Some(true))
+    }
+
+    pub async fn submit_error(db: &Surreal<Client>, id: &str, error: String) -> Result<(), Error> {
+        #[derive(Serialize)]
+        struct UpdateVmError {
+            error: String,
+        }
+        let _: Option<Self> = db.update((UPDATE_VM_REQ, id)).merge(UpdateVmError { error }).await?;
+        Ok(())
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -369,6 +549,28 @@ pub struct DeletedVm {
     pub price_per_unit: u64,
 }
 
+impl From<ActiveVm> for DeletedVm {
+    fn from(active_vm: ActiveVm) -> Self {
+        Self {
+            id: RecordId::from((DELETED_VM, active_vm.id.key().to_string())),
+            admin: active_vm.admin,
+            vm_node: active_vm.vm_node,
+            hostname: active_vm.hostname,
+            mapped_ports: active_vm.mapped_ports,
+            public_ipv4: active_vm.public_ipv4,
+            public_ipv6: active_vm.public_ipv6,
+            disk_size_gb: active_vm.disk_size_gb,
+            vcpus: active_vm.vcpus,
+            memory_mb: active_vm.memory_mb,
+            dtrfs_sha: active_vm.dtrfs_sha,
+            kernel_sha: active_vm.kernel_sha,
+            created_at: active_vm.created_at,
+            deleted_at: Datetime::default(),
+            price_per_unit: active_vm.price_per_unit,
+        }
+    }
+}
+
 impl DeletedVm {
     pub async fn get_by_uuid(db: &Surreal<Client>, uuid: &str) -> Result<Option<Self>, Error> {
         let contract: Option<Self> =
@@ -431,24 +633,6 @@ impl DeletedVm {
     }
 }
 
-impl ActiveVm {
-    /// total hardware units of this VM
-    fn total_units(&self) -> u64 {
-        // TODO: Optimize this based on price of hardware.
-        // I tried, but this can be done better.
-        // Storage cost should also be based on tier
-        (self.vcpus as u64 * 10)
-            + ((self.memory_mb + 256) as u64 / 200)
-            + (self.disk_size_gb as u64 / 10)
-            + (!self.public_ipv4.is_empty() as u64 * 10)
-    }
-
-    /// Returns price per minute in nanoLP
-    pub fn price_per_minute(&self) -> u64 {
-        self.total_units() * self.price_per_unit
-    }
-}
-
 #[derive(Debug, Serialize, Deserialize)]
 pub struct ActiveVmWithNode {
     pub id: RecordId,

src/grpc/types.rs

@@ -61,15 +61,45 @@ impl From<db::NewVmReq> for NewVmReq {
     }
 }
 
-impl From<db::NewVmResp> for NewVmResp {
+impl From<db::WrappedMeasurement> for NewVmResp {
-    fn from(resp: db::NewVmResp) -> Self {
+    fn from(resp: db::WrappedMeasurement) -> Self {
         match resp {
-            // TODO: This will require a small architecture change to pass MeasurementArgs from
+            db::WrappedMeasurement::Args(uuid, args) => {
-            // Daemon to CLI
+                Self { uuid, error: String::new(), args: Some(args) }
-            db::NewVmResp::Args(uuid, args) => {
-                NewVmResp { uuid, error: String::new(), args: Some(args) }
             }
-            db::NewVmResp::Error(uuid, error) => NewVmResp { uuid, error, args: None },
+            db::WrappedMeasurement::Error(uuid, error) => NewVmResp { uuid, error, args: None },
+        }
+    }
+}
+
+// TODO: NewVmResp is identical to UpdateVmResp so we can actually remove it from proto
+impl From<db::WrappedMeasurement> for UpdateVmResp {
+    fn from(resp: db::WrappedMeasurement) -> Self {
+        match resp {
+            db::WrappedMeasurement::Args(uuid, args) => {
+                Self { uuid, error: String::new(), args: Some(args) }
+            }
+            db::WrappedMeasurement::Error(uuid, error) => Self { uuid, error, args: None },
+        }
+    }
+}
+
+impl From<UpdateVmReq> for db::UpdateVmReq {
+    fn from(new_vm_req: UpdateVmReq) -> Self {
+        Self {
+            id: RecordId::from((NEW_VM_REQ, new_vm_req.uuid)),
+            admin: RecordId::from((ACCOUNT, new_vm_req.admin_pubkey)),
+            // vm_node gets modified later, and only if the db::UpdateVmReq is required
+            vm_node: RecordId::from((VM_NODE, String::new())),
+            disk_size_gb: new_vm_req.disk_size_gb,
+            vcpus: new_vm_req.vcpus,
+            memory_mb: new_vm_req.memory_mb,
+            kernel_url: new_vm_req.kernel_url,
+            kernel_sha: new_vm_req.kernel_sha,
+            dtrfs_url: new_vm_req.dtrfs_url,
+            dtrfs_sha: new_vm_req.dtrfs_sha,
+            created_at: surrealdb::sql::Datetime::default(),
+            error: String::new(),
         }
     }
 }

src/grpc/vm.rs

@@ -1,5 +1,5 @@
 #![allow(dead_code)]
-use crate::constants::{ACCOUNT, VM_NODE};
+use crate::constants::{ACCOUNT, NEW_VM_REQ, UPDATE_VM_REQ, VM_NODE};
 use crate::db::prelude as db;
 use crate::grpc::{check_sig_from_parts, check_sig_from_req};
 use detee_shared::common_proto::Empty;
@@ -148,8 +148,6 @@ impl BrainVmDaemon for VmDaemonServer {
             match daemon_message {
                 Ok(msg) => match msg.msg {
                     Some(vm_daemon_message::Msg::NewVmResp(new_vm_resp)) => {
-                        // TODO: move new_vm_req to active_vm
-                        // also handle failure properly
                         if !new_vm_resp.error.is_empty() {
                             db::NewVmReq::submit_error(
                                 &self.db,
@@ -170,9 +168,24 @@ impl BrainVmDaemon for VmDaemonServer {
                             }
                         }
                     }
-                    Some(vm_daemon_message::Msg::UpdateVmResp(_update_vm_resp)) => {
+                    Some(vm_daemon_message::Msg::UpdateVmResp(update_vm_resp)) => {
-                        todo!();
+                        if !update_vm_resp.error.is_empty() {
-                        // self.data.submit_updatevm_resp(update_vm_resp).await;
+                            db::UpdateVmReq::submit_error(
+                                &self.db,
+                                &update_vm_resp.uuid,
+                                update_vm_resp.error,
+                            )
+                            .await?;
+                        } else {
+                            db::upsert_record(
+                                &self.db,
+                                "measurement_args",
+                                &update_vm_resp.uuid,
+                                update_vm_resp.args.clone(),
+                            )
+                            .await?;
+                            db::ActiveVm::update(&self.db, &update_vm_resp.uuid).await?;
+                        }
                     }
                     Some(vm_daemon_message::Msg::VmNodeResources(node_resources)) => {
                         let node_resources: db::VmNodeResources = node_resources.into();
@@ -207,28 +220,31 @@ impl BrainVmCli for VmCliServer {
     async fn new_vm(&self, req: Request<NewVmReq>) -> Result<Response<NewVmResp>, Status> {
         let req = check_sig_from_req(req)?;
         info!("New VM requested via CLI: {req:?}");
-        if db::general::Account::is_banned_by_node(&self.db, &req.admin_pubkey, &req.node_pubkey).await? {
+        if db::general::Account::is_banned_by_node(&self.db, &req.admin_pubkey, &req.node_pubkey)
+            .await?
+        {
             return Err(Status::permission_denied("This operator banned you. What did you do?"));
         }
 
-        let new_vm_req: db::NewVmReq = req.into();
+        let db_req: db::NewVmReq = req.into();
-        let id = new_vm_req.id.key().to_string();
+        let id = db_req.id.key().to_string();
-
-        let db = self.db.clone();
 
         let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
+        let db = self.db.clone();
         tokio::spawn(async move {
-            let _ = oneshot_tx.send(db::NewVmResp::listen(&db, &id).await);
+            let _ = oneshot_tx.send(db::WrappedMeasurement::listen(&db, &id, NEW_VM_REQ).await);
         });
-        new_vm_req.submit(&self.db).await?;
+        db_req.submit(&self.db).await?;
 
         match oneshot_rx.await {
-            Ok(Err(db::Error::TimeOut(_))) => Err(Status::deadline_exceeded("Request failed due to timeout. Please try again later or contact the DeTEE devs team.")),
+            Ok(Err(db::Error::TimeOut(_))) => Err(Status::deadline_exceeded(
+                "Network timeout. Please try again later or contact the DeTEE devs team.",
+            )),
             Ok(new_vm_resp) => Ok(Response::new(new_vm_resp?.into())),
             Err(e) => {
-                log::error!("Something weird happened. Reached error {e:?}");
+                log::error!("Something weird happened during CLI NewVmReq. Reached error {e:?}");
                 Err(Status::unknown(
-                    "Request failed due to unknown error. Please try again or contact the DeTEE devs team.",
+                    "Unknown error. Please try again or contact the DeTEE devs team.",
                 ))
             }
         }
@@ -237,18 +253,62 @@ impl BrainVmCli for VmCliServer {
     async fn update_vm(&self, req: Request<UpdateVmReq>) -> Result<Response<UpdateVmResp>, Status> {
         let req = check_sig_from_req(req)?;
         info!("Update VM requested via CLI: {req:?}");
-        todo!();
+
-        // let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
+        let db_req: db::UpdateVmReq = req.clone().into();
-        // self.data.submit_updatevm_req(req, oneshot_tx).await;
+        let id = db_req.id.key().to_string();
-        // match oneshot_rx.await {
+
-        //     Ok(response) => {
+        let mut hostname_changed = false;
-        //         info!("Sending UpdateVMResp: {response:?}");
+        if !req.hostname.is_empty() {
-        //         Ok(Response::new(response))
+            hostname_changed =
-        //     }
+                db::ActiveVm::change_hostname(&self.db, &req.uuid, &req.hostname).await?;
-        //     Err(e) => Err(Status::unknown(format!(
+        }
-        //         "Update VM request failed due to error: {e}"
+
-        //     ))),
+        let hw_change_submitted = db_req.request_hw_update(&self.db).await?;
-        // }
+        if hw_change_submitted.is_none() {
+            return Ok(Response::new(UpdateVmResp {
+                uuid: req.uuid.clone(),
+                error: "VM Contract does not exist.".to_string(),
+                args: None,
+            }));
+        }
+        let hw_change_needed = hw_change_submitted.unwrap();
+
+        if !hostname_changed && !hw_change_needed {
+            return Ok(Response::new(UpdateVmResp {
+                uuid: req.uuid.clone(),
+                error: "No modification required".to_string(),
+                args: None,
+            }));
+        }
+
+        // if only the hostname got changed, return a confirmation
+        if !hw_change_needed {
+            return Ok(Response::new(UpdateVmResp {
+                uuid: req.uuid.clone(),
+                error: String::new(),
+                args: None,
+            }));
+        }
+
+        // if HW changes got requested, wait for the new args
+        let (oneshot_tx, oneshot_rx) = tokio::sync::oneshot::channel();
+        let db = self.db.clone();
+        tokio::spawn(async move {
+            let _ = oneshot_tx.send(db::WrappedMeasurement::listen(&db, &id, UPDATE_VM_REQ).await);
+        });
+
+        match oneshot_rx.await {
+            Ok(Err(db::Error::TimeOut(_))) => Err(Status::deadline_exceeded(
+                "Network timeout. Please try again later or contact the DeTEE devs team.",
+            )),
+            Ok(new_vm_resp) => Ok(Response::new(new_vm_resp?.into())),
+            Err(e) => {
+                log::error!("Something weird happened during CLI VM Update. Reached error {e:?}");
+                Err(Status::unknown(
+                    "Unknown error. Please try again or contact the DeTEE devs team.",
+                ))
+            }
+        }
     }
 
     async fn extend_vm(&self, req: Request<ExtendVmReq>) -> Result<Response<Empty>, Status> {

tests/common/prepare_test_env.rs

@@ -5,10 +5,8 @@ use dotenv::dotenv;
 use hyper_util::rt::TokioIo;
 use std::net::SocketAddr;
 use std::sync::Arc;
-use surreal_brain::grpc::{
+use surreal_brain::grpc::general::GeneralCliServer;
-    general::GeneralCliServer,
+use surreal_brain::grpc::vm::{VmCliServer, VmDaemonServer};
-    vm::{VmCliServer, VmDaemonServer},
-};
 use surrealdb::engine::remote::ws::Client;
 use surrealdb::Surreal;
 use tokio::io::DuplexStream;