From 88388539571062034dc3182fc94f7346d037b39f Mon Sep 17 00:00:00 2001 From: kieran Date: Wed, 5 Mar 2025 10:03:29 +0000 Subject: [PATCH] feat: delete ip resources --- lnvps_db/src/lib.rs | 2 +- src/api/routes.rs | 1 + src/mocks.rs | 54 ++++++++----- src/provisioner/lnvps.rs | 160 +++++++++++++++++++++++++++------------ src/router/mikrotik.rs | 57 ++++++++------ src/router/mod.rs | 25 +++--- 6 files changed, 200 insertions(+), 99 deletions(-) diff --git a/lnvps_db/src/lib.rs b/lnvps_db/src/lib.rs index 2a26f09..162b59e 100644 --- a/lnvps_db/src/lib.rs +++ b/lnvps_db/src/lib.rs @@ -104,7 +104,7 @@ pub trait LNVpsDb: Sync + Send { /// List VM ip assignments async fn insert_vm_ip_assignment(&self, ip_assignment: &VmIpAssignment) -> Result; - /// Update VM ip assignments + /// Update VM ip assignments (arp/dns refs) async fn update_vm_ip_assignment(&self, ip_assignment: &VmIpAssignment) -> Result<()>; /// List VM ip assignments diff --git a/src/api/routes.rs b/src/api/routes.rs index e6ee000..8940788 100644 --- a/src/api/routes.rs +++ b/src/api/routes.rs @@ -232,6 +232,7 @@ async fn v1_patch_vm( for mut ip in ips.iter_mut() { ip.dns_reverse = Some(ptr.to_string()); provisioner.update_reverse_ip_dns(&mut ip).await?; + db.update_vm_ip_assignment(&ip).await?; } } diff --git a/src/mocks.rs b/src/mocks.rs index 560f807..fe04193 100644 --- a/src/mocks.rs +++ b/src/mocks.rs @@ -5,7 +5,7 @@ use crate::lightning::{AddInvoiceRequest, AddInvoiceResult, InvoiceUpdate, Light use crate::router::{ArpEntry, Router}; use crate::settings::NetworkPolicy; use crate::status::{VmRunningState, VmState}; -use anyhow::{anyhow, bail}; +use anyhow::{anyhow, bail, ensure}; use chrono::{DateTime, Utc}; use fedimint_tonic_lnd::tonic::codegen::tokio_stream::Stream; use lnvps_db::{ @@ -14,7 +14,6 @@ use lnvps_db::{ VmIpAssignment, VmOsImage, VmPayment, VmTemplate, }; use std::collections::HashMap; -use std::net::IpAddr; use std::pin::Pin; use std::sync::{Arc, LazyLock}; use tokio::sync::Mutex; @@ -510,24 +509,15 @@ impl Router for MockRouter { Ok(arp.values().cloned().collect()) } - async fn add_arp_entry( - &self, - ip: IpAddr, - mac: &str, - interface: &str, - comment: Option<&str>, - ) -> anyhow::Result { + async fn add_arp_entry(&self, entry: &ArpEntry) -> anyhow::Result { let mut arp = self.arp.lock().await; - if arp.iter().any(|(k, v)| v.address == ip.to_string()) { + if arp.iter().any(|(k, v)| v.address == entry.address) { bail!("Address is already in use"); } let max_id = *arp.keys().max().unwrap_or(&0); let e = ArpEntry { - id: (max_id + 1).to_string(), - address: ip.to_string(), - mac_address: mac.to_string(), - interface: Some(interface.to_string()), - comment: comment.map(|s| s.to_string()), + id: Some((max_id + 1).to_string()), + ..entry.clone() }; arp.insert(max_id + 1, e.clone()); Ok(e) @@ -538,6 +528,18 @@ impl Router for MockRouter { arp.remove(&id.parse::()?); Ok(()) } + + async fn update_arp_entry(&self, entry: &ArpEntry) -> anyhow::Result { + ensure!(entry.id.is_some(), "id is missing"); + let mut arp = self.arp.lock().await; + if let Some(mut a) = arp.get_mut(&entry.id.as_ref().unwrap().parse::()?) { + a.mac_address = entry.mac_address.clone(); + a.address = entry.address.clone(); + a.interface = entry.interface.clone(); + a.comment = entry.comment.clone(); + } + Ok(entry.clone()) + } } #[derive(Clone, Debug, Default)] @@ -728,10 +730,26 @@ impl DnsServer for MockDnsServer { } async fn delete_record(&self, record: &BasicRecord) -> anyhow::Result<()> { - todo!() + let mut table = match record.kind { + RecordType::PTR => self.reverse.lock().await, + _ => self.forward.lock().await, + }; + ensure!(record.id.is_some(), "Id is missing"); + table.remove(record.id.as_ref().unwrap()); + Ok(()) } - async fn update_record(&self, name: &BasicRecord) -> anyhow::Result { - todo!() + async fn update_record(&self, record: &BasicRecord) -> anyhow::Result { + let mut table = match record.kind { + RecordType::PTR => self.reverse.lock().await, + _ => self.forward.lock().await, + }; + ensure!(record.id.is_some(), "Id is missing"); + if let Some(mut r) = table.get_mut(record.id.as_ref().unwrap()) { + r.name = record.name.clone(); + r.value = record.value.clone(); + r.kind = record.kind.to_string(); + } + Ok(record.clone()) } } diff --git a/src/provisioner/lnvps.rs b/src/provisioner/lnvps.rs index 7b04735..4f1140e 100644 --- a/src/provisioner/lnvps.rs +++ b/src/provisioner/lnvps.rs @@ -3,9 +3,9 @@ use crate::exchange::{ExchangeRateService, Ticker}; use crate::host::{get_host_client, FullVmInfo}; use crate::lightning::{AddInvoiceRequest, LightningNode}; use crate::provisioner::{NetworkProvisioner, ProvisionerMethod}; -use crate::router::Router; +use crate::router::{ArpEntry, Router}; use crate::settings::{NetworkAccessPolicy, NetworkPolicy, ProvisionerConfig, Settings}; -use anyhow::{bail, Context, Result}; +use anyhow::{bail, ensure, Context, Result}; use chrono::{Days, Months, Utc}; use futures::future::join_all; use lnvps_db::{IpRange, LNVpsDb, Vm, VmCostPlanIntervalType, VmIpAssignment, VmPayment}; @@ -55,24 +55,78 @@ impl LNVpsProvisioner { } } - pub async fn delete_ip_assignment(&self, vm: &Vm) -> Result<()> { - // Delete access policy - if let NetworkAccessPolicy::StaticArp { .. } = &self.network_policy.access { + /// Create or Update access policy for a given ip assignment, does not save to database! + pub async fn update_access_policy(&self, assignment: &mut VmIpAssignment) -> Result<()> { + // apply network policy + if let NetworkAccessPolicy::StaticArp { interface } = &self.network_policy.access { if let Some(r) = self.router.as_ref() { - let ent = r.list_arp_entry().await?; - if let Some(ent) = ent - .iter() - .find(|e| e.mac_address.eq_ignore_ascii_case(&vm.mac_address)) - { - r.remove_arp_entry(&ent.id).await?; + let vm = self.db.get_vm(assignment.vm_id).await?; + let entry = ArpEntry::new(&vm, &assignment, Some(interface.clone()))?; + let arp = if let Some(_id) = &assignment.arp_ref { + r.update_arp_entry(&entry).await? } else { - warn!("ARP entry not found, skipping") - } + r.add_arp_entry(&entry).await? + }; + ensure!(arp.id.is_some(), "ARP id was empty"); + assignment.arp_ref = arp.id; + } else { + bail!("No router found to apply static arp entry!") } } Ok(()) } + /// Remove an access policy for a given ip assignment, does not save to database! + pub async fn remove_access_policy(&self, assignment: &mut VmIpAssignment) -> Result<()> { + // Delete access policy + if let NetworkAccessPolicy::StaticArp { .. } = &self.network_policy.access { + if let Some(r) = self.router.as_ref() { + let id = if let Some(id) = &assignment.arp_ref { + Some(id.clone()) + } else { + warn!("ARP REF not found, using arp list"); + + let ent = r.list_arp_entry().await?; + if let Some(ent) = ent.iter().find(|e| e.address == assignment.ip) { + ent.id.clone() + } else { + warn!("ARP entry not found, skipping"); + None + } + }; + + if let Some(id) = id { + if let Err(e) = r.remove_arp_entry(&id).await { + warn!("Failed to remove arp entry, skipping: {}", e); + } + } + assignment.arp_ref = None; + } + } + Ok(()) + } + + /// Delete DNS on the dns server, does not save to database! + pub async fn remove_ip_dns(&self, assignment: &mut VmIpAssignment) -> Result<()> { + // Delete forward/reverse dns + if let Some(dns) = &self.dns { + if let Some(_r) = &assignment.dns_reverse_ref { + let rev = BasicRecord::reverse(assignment)?; + dns.delete_record(&rev).await?; + assignment.dns_reverse_ref = None; + assignment.dns_reverse = None; + } + if let Some(_r) = &assignment.dns_forward_ref { + let rev = BasicRecord::forward(assignment)?; + dns.delete_record(&rev).await?; + assignment.dns_forward_ref = None; + assignment.dns_forward = None; + } + } + Ok(()) + } + + /// Update DNS on the dns server, does not save to database! pub async fn update_forward_ip_dns(&self, assignment: &mut VmIpAssignment) -> Result<()> { if let Some(dns) = &self.dns { let fwd = BasicRecord::forward(assignment)?; @@ -83,15 +137,11 @@ impl LNVpsProvisioner { }; assignment.dns_forward = Some(ret_fwd.name); assignment.dns_forward_ref = Some(ret_fwd.id.context("Record id is missing")?); - - // save to db - if assignment.id != 0 { - self.db.update_vm_ip_assignment(assignment).await?; - } } Ok(()) } + /// Update DNS on the dns server, does not save to database! pub async fn update_reverse_ip_dns(&self, assignment: &mut VmIpAssignment) -> Result<()> { if let Some(dns) = &self.dns { let ret_rev = if assignment.dns_reverse_ref.is_some() { @@ -103,32 +153,30 @@ impl LNVpsProvisioner { }; assignment.dns_reverse = Some(ret_rev.value); assignment.dns_reverse_ref = Some(ret_rev.id.context("Record id is missing")?); - - // save to db - if assignment.id != 0 { - self.db.update_vm_ip_assignment(assignment).await?; - } } Ok(()) } - async fn save_ip_assignment(&self, vm: &Vm, assignment: &mut VmIpAssignment) -> Result<()> { - let ip = IpAddr::from_str(&assignment.ip)?; - - // apply network policy - if let NetworkAccessPolicy::StaticArp { interface } = &self.network_policy.access { - if let Some(r) = self.router.as_ref() { - r.add_arp_entry( - ip.clone(), - &vm.mac_address, - interface, - Some(&format!("VM{}", vm.id)), - ) - .await?; - } else { - bail!("No router found to apply static arp entry!") - } + /// Delete all ip assignments for a given vm + pub async fn delete_ip_assignments(&self, vm_id: u64) -> Result<()> { + let ips = self.db.list_vm_ip_assignments(vm_id).await?; + for mut ip in ips { + // remove access policy + self.remove_access_policy(&mut ip).await?; + // remove dns + self.remove_ip_dns(&mut ip).await?; + // save arp/dns changes + self.db.update_vm_ip_assignment(&ip).await?; } + // mark as deleted + self.db.delete_vm_ip_assignment(vm_id).await?; + + Ok(()) + } + + async fn save_ip_assignment(&self, assignment: &mut VmIpAssignment) -> Result<()> { + // apply access policy + self.update_access_policy(assignment).await?; // Add DNS records self.update_forward_ip_dns(assignment).await?; @@ -164,7 +212,7 @@ impl LNVpsProvisioner { dns_reverse_ref: None, }; - self.save_ip_assignment(&vm, &mut assignment).await?; + self.save_ip_assignment(&mut assignment).await?; Ok(vec![assignment]) } @@ -324,14 +372,11 @@ impl LNVpsProvisioner { /// Delete a VM and its associated resources pub async fn delete_vm(&self, vm_id: u64) -> Result<()> { - let vm = self.db.get_vm(vm_id).await?; - // host client currently doesn't support delete (proxmox) // VM should already be stopped by [Worker] - self.delete_ip_assignment(&vm).await?; - self.db.delete_vm_ip_assignment(vm.id).await?; - self.db.delete_vm(vm.id).await?; + self.delete_ip_assignments(vm_id).await?; + self.db.delete_vm(vm_id).await?; Ok(()) } @@ -351,7 +396,7 @@ impl LNVpsProvisioner { mod tests { use super::*; use crate::exchange::DefaultRateCache; - use crate::mocks::{MockDb, MockNode}; + use crate::mocks::{MockDb, MockDnsServer, MockNode, MockRouter}; use crate::settings::{DnsServerConfig, LightningConfig, QemuConfig, RouterConfig}; use lnvps_db::UserSshKey; @@ -403,7 +448,8 @@ mod tests { let db = Arc::new(MockDb::default()); let node = Arc::new(MockNode::default()); let rates = Arc::new(DefaultRateCache::default()); - let router = settings.get_router().expect("router").unwrap(); + let router = MockRouter::new(settings.network_policy.clone()); + let dns = MockDnsServer::new(); let provisioner = LNVpsProvisioner::new(settings, db.clone(), node.clone(), rates.clone()); let pubkey: [u8; 32] = random(); @@ -448,6 +494,26 @@ mod tests { assert!(!ip.ip.ends_with("/8")); assert!(!ip.ip.ends_with("/24")); + // now expire + provisioner.delete_vm(vm.id).await?; + + // test arp/dns is removed + let arp = router.list_arp_entry().await?; + assert!(arp.is_empty()); + assert_eq!(dns.forward.lock().await.len(), 0); + assert_eq!(dns.reverse.lock().await.len(), 0); + + // ensure IPS are deleted + let ips = db.ip_assignments.lock().await; + let ip = ips.values().next().unwrap(); + assert!(ip.arp_ref.is_none()); + assert!(ip.dns_forward.is_none()); + assert!(ip.dns_reverse.is_none()); + assert!(ip.dns_reverse_ref.is_none()); + assert!(ip.dns_forward_ref.is_none()); + assert!(ip.deleted); + println!("{:?}", ip); + Ok(()) } } diff --git a/src/router/mikrotik.rs b/src/router/mikrotik.rs index 3315946..89a8df2 100644 --- a/src/router/mikrotik.rs +++ b/src/router/mikrotik.rs @@ -1,6 +1,6 @@ use crate::json_api::JsonApi; use crate::router::{ArpEntry, Router}; -use anyhow::Result; +use anyhow::{ensure, Result}; use base64::engine::general_purpose::STANDARD; use base64::Engine; use log::debug; @@ -32,27 +32,9 @@ impl Router for MikrotikRouter { Ok(rsp.into_iter().map(|e| e.into()).collect()) } - async fn add_arp_entry( - &self, - ip: IpAddr, - mac: &str, - arp_interface: &str, - comment: Option<&str>, - ) -> Result { - let rsp: MikrotikArpEntry = self - .api - .req( - Method::PUT, - "/rest/ip/arp", - MikrotikArpEntry { - id: None, - address: ip.to_string(), - mac_address: Some(mac.to_string()), - interface: arp_interface.to_string(), - comment: comment.map(|c| c.to_string()), - }, - ) - .await?; + async fn add_arp_entry(&self, entry: &ArpEntry) -> Result { + let req: MikrotikArpEntry = entry.clone().into(); + let rsp: MikrotikArpEntry = self.api.req(Method::PUT, "/rest/ip/arp", req).await?; debug!("{:?}", rsp); Ok(rsp.into()) } @@ -60,11 +42,26 @@ impl Router for MikrotikRouter { async fn remove_arp_entry(&self, id: &str) -> Result<()> { let rsp: MikrotikArpEntry = self .api - .req(Method::DELETE, &format!("/rest/ip/arp/{id}"), ()) + .req(Method::DELETE, &format!("/rest/ip/arp/{}", id), ()) .await?; debug!("{:?}", rsp); Ok(()) } + + async fn update_arp_entry(&self, entry: &ArpEntry) -> Result { + ensure!(entry.id.is_some(), "Cannot update an arp entry without ID"); + let req: MikrotikArpEntry = entry.clone().into(); + let rsp: MikrotikArpEntry = self + .api + .req( + Method::PATCH, + &format!("/rest/ip/arp/{}", entry.id.as_ref().unwrap()), + req, + ) + .await?; + debug!("{:?}", rsp); + Ok(rsp.into()) + } } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -84,7 +81,7 @@ pub struct MikrotikArpEntry { impl Into for MikrotikArpEntry { fn into(self) -> ArpEntry { ArpEntry { - id: self.id.unwrap(), + id: self.id, address: self.address, mac_address: self.mac_address.unwrap(), interface: Some(self.interface), @@ -92,3 +89,15 @@ impl Into for MikrotikArpEntry { } } } + +impl Into for ArpEntry { + fn into(self) -> MikrotikArpEntry { + MikrotikArpEntry { + id: self.id, + address: self.address, + mac_address: Some(self.mac_address), + interface: self.interface.unwrap(), + comment: self.comment, + } + } +} diff --git a/src/router/mod.rs b/src/router/mod.rs index c685969..c69b526 100644 --- a/src/router/mod.rs +++ b/src/router/mod.rs @@ -1,6 +1,6 @@ use anyhow::Result; +use lnvps_db::{Vm, VmIpAssignment}; use rocket::async_trait; -use std::net::IpAddr; /// Router defines a network device used to access the hosts /// @@ -12,25 +12,32 @@ use std::net::IpAddr; #[async_trait] pub trait Router: Send + Sync { async fn list_arp_entry(&self) -> Result>; - async fn add_arp_entry( - &self, - ip: IpAddr, - mac: &str, - interface: &str, - comment: Option<&str>, - ) -> Result; + async fn add_arp_entry(&self, entry: &ArpEntry) -> Result; async fn remove_arp_entry(&self, id: &str) -> Result<()>; + async fn update_arp_entry(&self, entry: &ArpEntry) -> Result; } #[derive(Debug, Clone)] pub struct ArpEntry { - pub id: String, + pub id: Option, pub address: String, pub mac_address: String, pub interface: Option, pub comment: Option, } +impl ArpEntry { + pub fn new(vm: &Vm, ip: &VmIpAssignment, interface: Option) -> Result { + Ok(Self { + id: ip.arp_ref.clone(), + address: ip.ip.clone(), + mac_address: vm.mac_address.clone(), + interface, + comment: Some(format!("VM{}", vm.id)), + }) + } +} + #[cfg(feature = "mikrotik")] mod mikrotik; #[cfg(feature = "mikrotik")]