kratad/network/
assignment.rs

1use advmac::MacAddr6;
2use anyhow::{anyhow, Result};
3use ipnetwork::{Ipv4Network, Ipv6Network};
4use std::{
5    collections::HashMap,
6    net::{Ipv4Addr, Ipv6Addr},
7    sync::Arc,
8};
9use tokio::sync::RwLock;
10use uuid::Uuid;
11
12use crate::db::network::{NetworkReservation, NetworkReservationStore};
13
14#[derive(Default, Clone)]
15pub struct NetworkAssignmentState {
16    pub ipv4: HashMap<Ipv4Addr, NetworkReservation>,
17    pub ipv6: HashMap<Ipv6Addr, NetworkReservation>,
18}
19
20#[derive(Clone)]
21pub struct NetworkAssignment {
22    ipv4_network: Ipv4Network,
23    ipv6_network: Ipv6Network,
24    gateway_ipv4: Ipv4Addr,
25    gateway_ipv6: Ipv6Addr,
26    gateway_mac: MacAddr6,
27    store: NetworkReservationStore,
28    state: Arc<RwLock<NetworkAssignmentState>>,
29}
30
31impl NetworkAssignment {
32    pub async fn new(
33        host_uuid: Uuid,
34        ipv4_network: Ipv4Network,
35        ipv6_network: Ipv6Network,
36        store: NetworkReservationStore,
37    ) -> Result<Self> {
38        let mut state = NetworkAssignment::fetch_current_state(&store).await?;
39        let gateway_reservation = if let Some(reservation) = store.read(Uuid::nil()).await? {
40            reservation
41        } else {
42            NetworkAssignment::allocate(
43                &mut state,
44                &store,
45                Uuid::nil(),
46                ipv4_network,
47                ipv6_network,
48                None,
49                None,
50                None,
51            )
52            .await?
53        };
54
55        if store.read(host_uuid).await?.is_none() {
56            let _ = NetworkAssignment::allocate(
57                &mut state,
58                &store,
59                host_uuid,
60                ipv4_network,
61                ipv6_network,
62                Some(gateway_reservation.gateway_ipv4),
63                Some(gateway_reservation.gateway_ipv6),
64                Some(gateway_reservation.gateway_mac),
65            )
66            .await?;
67        }
68
69        let assignment = NetworkAssignment {
70            ipv4_network,
71            ipv6_network,
72            gateway_ipv4: gateway_reservation.ipv4,
73            gateway_ipv6: gateway_reservation.ipv6,
74            gateway_mac: gateway_reservation.mac,
75            store,
76            state: Arc::new(RwLock::new(state)),
77        };
78        Ok(assignment)
79    }
80
81    async fn fetch_current_state(
82        store: &NetworkReservationStore,
83    ) -> Result<NetworkAssignmentState> {
84        let reservations = store.list().await?;
85        let mut state = NetworkAssignmentState::default();
86        for reservation in reservations.values() {
87            state.ipv4.insert(reservation.ipv4, reservation.clone());
88            state.ipv6.insert(reservation.ipv6, reservation.clone());
89        }
90        Ok(state)
91    }
92
93    #[allow(clippy::too_many_arguments)]
94    async fn allocate(
95        state: &mut NetworkAssignmentState,
96        store: &NetworkReservationStore,
97        uuid: Uuid,
98        ipv4_network: Ipv4Network,
99        ipv6_network: Ipv6Network,
100        gateway_ipv4: Option<Ipv4Addr>,
101        gateway_ipv6: Option<Ipv6Addr>,
102        gateway_mac: Option<MacAddr6>,
103    ) -> Result<NetworkReservation> {
104        let found_ipv4: Option<Ipv4Addr> = ipv4_network
105            .iter()
106            .filter(|ip| {
107                ip.is_private() && !(ip.is_loopback() || ip.is_multicast() || ip.is_broadcast())
108            })
109            .filter(|ip| {
110                let last = ip.octets()[3];
111                // filter for IPs ending in .1 to .250 because .250+ can have special meaning
112                (1..250).contains(&last)
113            })
114            .find(|ip| !state.ipv4.contains_key(ip));
115
116        let found_ipv6: Option<Ipv6Addr> = ipv6_network
117            .iter()
118            .filter(|ip| !ip.is_loopback() && !ip.is_multicast())
119            .filter(|ip| {
120                let last = ip.octets()[15];
121                last > 0
122            })
123            .find(|ip| !state.ipv6.contains_key(ip));
124
125        let Some(ipv4) = found_ipv4 else {
126            return Err(anyhow!(
127                "unable to allocate ipv4 address, assigned network is exhausted"
128            ));
129        };
130
131        let Some(ipv6) = found_ipv6 else {
132            return Err(anyhow!(
133                "unable to allocate ipv6 address, assigned network is exhausted"
134            ));
135        };
136
137        let mut mac = MacAddr6::random();
138        mac.set_local(true);
139        mac.set_multicast(false);
140
141        let reservation = NetworkReservation {
142            uuid: uuid.to_string(),
143            ipv4,
144            ipv6,
145            mac,
146            ipv4_prefix: ipv4_network.prefix(),
147            ipv6_prefix: ipv6_network.prefix(),
148            gateway_ipv4: gateway_ipv4.unwrap_or(ipv4),
149            gateway_ipv6: gateway_ipv6.unwrap_or(ipv6),
150            gateway_mac: gateway_mac.unwrap_or(mac),
151        };
152        state.ipv4.insert(ipv4, reservation.clone());
153        state.ipv6.insert(ipv6, reservation.clone());
154        store.update(uuid, reservation.clone()).await?;
155        Ok(reservation)
156    }
157
158    pub async fn assign(&self, uuid: Uuid) -> Result<NetworkReservation> {
159        let mut state = self.state.write().await;
160        let reservation = NetworkAssignment::allocate(
161            &mut state,
162            &self.store,
163            uuid,
164            self.ipv4_network,
165            self.ipv6_network,
166            Some(self.gateway_ipv4),
167            Some(self.gateway_ipv6),
168            Some(self.gateway_mac),
169        )
170        .await?;
171        Ok(reservation)
172    }
173
174    pub async fn recall(&self, uuid: Uuid) -> Result<()> {
175        let mut state = self.state.write().await;
176        self.store.remove(uuid).await?;
177        state
178            .ipv4
179            .retain(|_, reservation| reservation.uuid != uuid.to_string());
180        state
181            .ipv6
182            .retain(|_, reservation| reservation.uuid != uuid.to_string());
183        Ok(())
184    }
185
186    pub async fn retrieve(&self, uuid: Uuid) -> Result<Option<NetworkReservation>> {
187        self.store.read(uuid).await
188    }
189
190    pub async fn reload(&self) -> Result<()> {
191        let mut state = self.state.write().await;
192        let intermediate = NetworkAssignment::fetch_current_state(&self.store).await?;
193        *state = intermediate;
194        Ok(())
195    }
196
197    pub async fn read(&self) -> Result<NetworkAssignmentState> {
198        Ok(self.state.read().await.clone())
199    }
200
201    pub async fn read_reservations(&self) -> Result<HashMap<Uuid, NetworkReservation>> {
202        self.store.list().await
203    }
204}
kratad/network/assignment.rs

kratad/network/
assignment.rs
