use anyhow::Result; use krata::{ events::EventStream, v1::{ common::Guest, control::{ control_service_client::ControlServiceClient, watch_events_reply::Event, ListGuestsRequest, }, }, }; use log::warn; use smoltcp::wire::{EthernetAddress, Ipv4Cidr, Ipv6Cidr}; use std::{collections::HashMap, str::FromStr, time::Duration}; use tokio::{select, sync::broadcast::Receiver, time::sleep}; use tonic::transport::Channel; use uuid::Uuid; pub struct AutoNetworkWatcher { control: ControlServiceClient, pub events: EventStream, known: HashMap, } #[derive(Debug, Clone)] pub struct NetworkSide { pub ipv4: Ipv4Cidr, pub ipv6: Ipv6Cidr, pub mac: EthernetAddress, } #[derive(Debug, Clone)] pub struct NetworkMetadata { pub domid: u32, pub uuid: Uuid, pub guest: NetworkSide, pub gateway: NetworkSide, } impl NetworkMetadata { pub fn interface(&self) -> String { format!("vif{}.20", self.domid) } } #[derive(Debug, Clone)] pub struct AutoNetworkChangeset { pub added: Vec, pub removed: Vec, } impl AutoNetworkWatcher { pub async fn new(control: ControlServiceClient) -> Result { let client = control.clone(); Ok(AutoNetworkWatcher { control, events: EventStream::open(client).await?, known: HashMap::new(), }) } pub async fn read(&mut self) -> Result> { let mut all_guests: HashMap = HashMap::new(); for guest in self .control .list_guests(ListGuestsRequest {}) .await? .into_inner() .guests { let Ok(uuid) = Uuid::from_str(&guest.id) else { continue; }; all_guests.insert(uuid, guest); } let mut networks: Vec = Vec::new(); for (uuid, guest) in &all_guests { let Some(ref state) = guest.state else { continue; }; if state.domid == u32::MAX { continue; } let Some(ref network) = state.network else { continue; }; let Ok(guest_ipv4_cidr) = Ipv4Cidr::from_str(&network.guest_ipv4) else { continue; }; let Ok(guest_ipv6_cidr) = Ipv6Cidr::from_str(&network.guest_ipv6) else { continue; }; let Ok(guest_mac) = EthernetAddress::from_str(&network.guest_mac) else { continue; }; let Ok(gateway_ipv4_cidr) = Ipv4Cidr::from_str(&network.gateway_ipv4) else { continue; }; let Ok(gateway_ipv6_cidr) = Ipv6Cidr::from_str(&network.gateway_ipv6) else { continue; }; let Ok(gateway_mac) = EthernetAddress::from_str(&network.gateway_mac) else { continue; }; networks.push(NetworkMetadata { domid: state.domid, uuid: *uuid, guest: NetworkSide { ipv4: guest_ipv4_cidr, ipv6: guest_ipv6_cidr, mac: guest_mac, }, gateway: NetworkSide { ipv4: gateway_ipv4_cidr, ipv6: gateway_ipv6_cidr, mac: gateway_mac, }, }); } Ok(networks) } pub async fn read_changes(&mut self) -> Result { let mut seen: Vec = Vec::new(); let mut added: Vec = Vec::new(); let mut removed: Vec = Vec::new(); let networks = match self.read().await { Ok(networks) => networks, Err(error) => { warn!("failed to read network changes: {}", error); return Ok(AutoNetworkChangeset { added, removed }); } }; for network in networks { seen.push(network.uuid); if self.known.contains_key(&network.uuid) { continue; } let _ = self.known.insert(network.uuid, network.clone()); added.push(network); } let mut gone: Vec = Vec::new(); for uuid in self.known.keys() { if seen.contains(uuid) { continue; } gone.push(*uuid); } for uuid in &gone { let Some(network) = self.known.remove(uuid) else { continue; }; removed.push(network); } Ok(AutoNetworkChangeset { added, removed }) } pub async fn wait(&mut self, receiver: &mut Receiver) -> Result<()> { loop { select! { x = receiver.recv() => match x { Ok(Event::GuestChanged(_)) => { break; }, Ok(_) => {}, Err(error) => { warn!("failed to receive event: {}", error); } }, _ = sleep(Duration::from_secs(10)) => { break; } }; } Ok(()) } pub fn mark_unknown(&mut self, uuid: Uuid) -> Result { Ok(self.known.remove(&uuid).is_some()) } }