ghe0/brain-mock
1
0
Fork 1
Code Issues Pull Requests 2 Packages Projects Releases Wiki Activity

update functionality for the brain #2

Merged
ghe0 merged 11 commits from new-updater into main 2025-01-09 22:29:56 +00:00
Conversation 1 Commits 11 Files Changed 3 +101 -55
1 changed files with 101 additions and 55 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit 5755ece192 - Show all commits

130
src/data.rs
View File

@ -196,29 +196,47 @@ impl BrainData {
return;
}
};
let args = match new_vm_resp.args.clone() {
Some(args) => args,
match new_vm_resp.args {
Some(ref mut args) => {
if args.dtrfs_api_endpoint.starts_with(':') {
match self.find_nodes_by_pubkey(&new_vm_req.0.node_pubkey) {
Some(node) => {
args.dtrfs_api_endpoint =
format!("{}{}", node.ip, args.dtrfs_api_endpoint);
}
None => {
log::error!("NewVmResp does not contain MeasurementArgs for {}", new_vm_resp.uuid);
log::error!("Node not found for pubkey {}", new_vm_req.0.node_pubkey);
new_vm_resp.error = "Node not found.".to_string();
return;
}
}
}
ghe0 marked this conversation as resolved Outdated
ghe0 commented 2025-01-04 23:03:54 +00:00
Outdated
Review
Copy Link

If the daemon does not provide Measurement Args, the CLI can't execute the attestation. I believe we should not accept grpc::NewVmResp that does not contain Args. This should be replaced with a NewVmResp that contains an error with the message "SNP Node did not provide measurement args"

Also, this means the contract is broken so we need to discuss the cases.

If the daemon does not provide Measurement Args, the CLI can't execute the attestation. I believe we should not accept `grpc::NewVmResp` that does not contain Args. This should be replaced with a NewVmResp that contains an error with the message "SNP Node did not provide measurement args" Also, this means the contract is broken so we need to discuss the cases.
}
None => {
log::error!(
"NewVmResp does not contain MeasurementArgs for {}",
new_vm_resp.uuid
);
new_vm_resp.error = "Daemon did not return measurement args.".to_string();
return;
}
};
if let Err(e) = new_vm_req.1.send(new_vm_resp.clone()) {
}
if let Err(_) = new_vm_req.1.send(new_vm_resp.clone()) {
log::error!(
"CLI RX for {} dropped before receiving confirmation {:?}. Error is: {:?}",
"CLI RX for {} dropped before receiving confirmation {:?}.",
&new_vm_req.0.admin_pubkey,
new_vm_resp,
e
);
}
if new_vm_resp.error != "" {
return;
}
let mut public_ipv4 = String::new();
let mut public_ipv6 = String::new();
for ip in args.ips {
let args = new_vm_resp.args.as_ref().unwrap();
for ip in args.ips.iter() {
if let Ok(ipv4_addr) = std::net::Ipv4Addr::from_str(&ip.address) {
if !ipv4_addr.is_private() && !ipv4_addr.is_link_local() {
public_ipv4 = ipv4_addr.to_string();
@ -232,7 +250,7 @@ impl BrainData {
let contract = Contract {
uuid: new_vm_resp.uuid,
exposed_ports: args.exposed_ports,
exposed_ports: args.exposed_ports.clone(),
public_ipv4,
public_ipv6,
created_at: Utc::now().to_rfc3339(),
@ -248,60 +266,88 @@ impl BrainData {
};
info!("Created new contract: {contract:?}");
self.contracts.write().unwrap().push(contract);
let args = new_vm_resp.args.as_mut().unwrap();
if args.dtrfs_api_endpoint.starts_with(':') {
if let Some(node) = self.find_nodes_by_pubkey(&new_vm_req.0.node_pubkey) {
args.dtrfs_api_endpoint = format!("{}{}", node.ip, args.dtrfs_api_endpoint);
}
}
}
pub async fn submit_updatevm_resp(&self, mut resp: grpc::UpdateVmResp) {
let updatevmreq = match self.tmp_updatevm_reqs.remove(&resp.uuid) {
pub async fn submit_updatevm_resp(&self, mut update_vm_resp: grpc::UpdateVmResp) {
let update_vm_req = match self.tmp_updatevm_reqs.remove(&update_vm_resp.uuid) {
Some((_, r)) => r,
None => {
log::error!(
"Received confirmation for ghost UpdateVMRequest {}",
resp.uuid
update_vm_resp.uuid
);
resp.error = "Received confirmation for ghost UpdateVMRequest.".to_string();
update_vm_resp.error =
"Received confirmation for ghost UpdateVMRequest.".to_string();
return;
}
};
if let None = resp.args {
log::error!("NewVmResp does not contain MeasurementArgs for {}", resp.uuid);
resp.error = "Daemon did not return measurement args.".to_string();
return;
};
if let Err(e) = updatevmreq.1.send(resp.clone()) {
log::error!("CLI RX dropped before receiving UpdateVMResp {resp:?}. Error is: {e:?}");
}
if resp.error != "" {
if update_vm_resp.error != "" {
return;
}
let mut contracts = self.contracts.write().unwrap();
if let Some(contract) = contracts.iter_mut().find(|c| c.uuid == resp.uuid) {
let args = resp.args.as_mut().unwrap();
match contracts.iter_mut().find(|c| c.uuid == update_vm_resp.uuid) {
Some(contract) => {
match update_vm_resp.args {
Some(ref mut args) => {
if args.dtrfs_api_endpoint.starts_with(':') {
if let Some(node) = self.find_nodes_by_pubkey(&contract.node_pubkey) {
args.dtrfs_api_endpoint = format!("{}{}", node.ip, args.dtrfs_api_endpoint);
match self.find_nodes_by_pubkey(&contract.node_pubkey) {
Some(node) => {
args.dtrfs_api_endpoint =
format!("{}{}", node.ip, args.dtrfs_api_endpoint);
}
None => {
log::error!(
"Node not found for pubkey {}",
contract.node_pubkey
);
update_vm_resp.error = "Node not found.".to_string();
return;
}
}
}
}
None => {
log::error!(
"NewVmResp does not contain MeasurementArgs for {}",
update_vm_resp.uuid
);
update_vm_resp.error =
"Daemon did not return measurement args.".to_string();
return;
}
}
contract.disk_size_gb = updatevmreq.0.disk_size_gb;
contract.vcpus = updatevmreq.0.vcpus;
contract.memory_mb = updatevmreq.0.memory_mb;
if !updatevmreq.0.kernel_sha.is_empty() {
info!("Updating kernel sha for {} to {}", contract.uuid, updatevmreq.0.kernel_sha);
contract.kernel_sha = updatevmreq.0.kernel_sha;
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.kernel_sha.is_empty() {
info!(
"Updating kernel sha for {} to {}",
contract.uuid, update_vm_req.0.kernel_sha
);
contract.kernel_sha = update_vm_req.0.kernel_sha;
}
if !updatevmreq.0.dtrfs_sha.is_empty() {
info!("Updating dtrfs sha for {} to {}", contract.uuid, updatevmreq.0.dtrfs_sha);
contract.dtrfs_sha = updatevmreq.0.dtrfs_sha;
if !update_vm_req.0.dtrfs_sha.is_empty() {
info!(
"Updating dtrfs sha for {} to {}",
contract.uuid, update_vm_req.0.dtrfs_sha
);
contract.dtrfs_sha = update_vm_req.0.dtrfs_sha;
}
contract.updated_at = Utc::now().to_rfc3339();
}
None => {
log::error!("Contract not found for {}.", update_vm_req.0.uuid);
update_vm_resp.error = "Contract not found.".to_string();
}
}
if let Err(_) = update_vm_req.1.send(update_vm_resp.clone()) {
log::error!(
"CLI RX dropped before receiving UpdateVMResp {:?}.",
update_vm_resp
);
}
}
pub async fn submit_newvmrequest(