From aaa9226df69f7316ccb5a4c1f80a14d464f80436 Mon Sep 17 00:00:00 2001
From: ghe0 <gheorghe@gheo.tech>
Date: Sat, 25 Jan 2025 16:12:46 +0200
Subject: [PATCH] refactored the vm update process

---
 snp.proto   |  7 +++++++
 src/data.rs | 46 +++++++++++++++++++++++++++++++++++++++++++---
 src/grpc.rs | 11 +++++++++++
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/snp.proto b/snp.proto
index 4addaf3..1d6a2f7 100644
--- a/snp.proto
+++ b/snp.proto
@@ -175,6 +175,12 @@ message NodeListResp {
   uint64 price = 8;
 }
 
+message ExtendVmReq {
+  string uuid = 1;
+  string admin_pubkey = 2;
+  uint64 locked_nano = 3;
+}
+
 service BrainCli {
   rpc GetAirdrop (Pubkey) returns (Empty);
   rpc GetBalance (Pubkey) returns (AccountBalance);
@@ -184,4 +190,5 @@ service BrainCli {
   rpc GetOneNode (NodeFilters) returns (NodeListResp);
   rpc DeleteVm (DeleteVmReq) returns (Empty);
   rpc UpdateVm (UpdateVmReq) returns (UpdateVmResp);
+  rpc ExtendVm (ExtendVmReq) returns (Empty);
 }
diff --git a/src/data.rs b/src/data.rs
index 4736b60..b2d0ce3 100644
--- a/src/data.rs
+++ b/src/data.rs
@@ -16,6 +16,8 @@ pub enum Error {
     InsufficientFunds,
     #[error("I have no idea how this happened. Please report this bug.")]
     ImpossibleError,
+    #[error("Could not find contract {0}")]
+    ContractNotFound(String),
 }
 
 #[derive(Clone)]
@@ -258,6 +260,38 @@ impl BrainData {
         }
     }
 
+    pub fn extend_contract_time(
+        &self,
+        uuid: &str,
+        account: &str,
+        nanotokens: u64,
+    ) -> Result<(), Error> {
+        if nanotokens > 100_000_000_000_000 {
+            return Err(Error::TxTooBig);
+        }
+        let mut account = match self.accounts.get_mut(account) {
+            Some(account) => account,
+            None => return Err(Error::InsufficientFunds),
+        };
+        match self
+            .contracts
+            .write()
+            .unwrap()
+            .iter_mut()
+            .find(|c| c.uuid == uuid)
+        {
+            Some(contract) => {
+                if account.balance + contract.locked_nano < nanotokens {
+                    return Err(Error::InsufficientFunds);
+                }
+                account.balance = account.balance + contract.locked_nano - nanotokens;
+                contract.locked_nano = nanotokens;
+                Ok(())
+            }
+            None => Err(Error::ContractNotFound(uuid.to_string())),
+        }
+    }
+
     pub fn submit_node_resources(&self, res: grpc::NodeResources) {
         let mut nodes = self.nodes.write().unwrap();
         for n in nodes.iter_mut() {
@@ -406,9 +440,15 @@ impl BrainData {
         let mut contracts = self.contracts.write().unwrap();
         match contracts.iter_mut().find(|c| c.uuid == update_vm_resp.uuid) {
             Some(contract) => {
-                contract.disk_size_gb = update_vm_req.0.disk_size_gb;
-                contract.vcpus = update_vm_req.0.vcpus;
-                contract.memory_mb = update_vm_req.0.memory_mb;
+                if update_vm_req.0.vcpus != 0 {
+                    contract.vcpus = update_vm_req.0.vcpus;
+                }
+                if update_vm_req.0.memory_mb != 0 {
+                    contract.memory_mb = update_vm_req.0.memory_mb;
+                }
+                if update_vm_req.0.disk_size_gb != 0 {
+                    contract.disk_size_gb = update_vm_req.0.disk_size_gb;
+                }
                 if !update_vm_req.0.kernel_sha.is_empty() {
                     debug!(
                         "Updating kernel sha for {} to {}",
diff --git a/src/grpc.rs b/src/grpc.rs
index de08ad0..746f9e4 100644
--- a/src/grpc.rs
+++ b/src/grpc.rs
@@ -171,6 +171,17 @@ impl BrainCli for BrainCliMock {
         }
     }
 
+    async fn extend_vm(&self, req: Request<ExtendVmReq>) -> Result<Response<Empty>, Status> {
+        let req = req.into_inner();
+        match self
+            .data
+            .extend_contract_time(&req.uuid, &req.admin_pubkey, req.locked_nano)
+        {
+            Ok(()) => Ok(Response::new(Empty {})),
+            Err(e) => Err(Status::unknown(format!("Could not extend contract: {e}"))),
+        }
+    }
+
     type ListContractsStream = Pin<Box<dyn Stream<Item = Result<Contract, Status>> + Send>>;
     async fn list_contracts(
         &self,
-- 
2.43.0