zlayer_overlay/
allocator.rs

1//! IP address allocation for overlay networks
2//!
3//! Manages allocation and tracking of overlay IP addresses within a CIDR range.
4
5use crate::error::{OverlayError, Result};
6use ipnet::Ipv4Net;
7use serde::{Deserialize, Serialize};
8use std::collections::HashSet;
9use std::net::Ipv4Addr;
10use std::path::Path;
11
12/// IP allocator for overlay network addresses
13///
14/// Tracks allocated IP addresses and provides next-available allocation
15/// from a configured CIDR range.
16#[derive(Debug, Clone)]
17pub struct IpAllocator {
18    /// Network CIDR range
19    network: Ipv4Net,
20    /// Set of allocated IP addresses
21    allocated: HashSet<Ipv4Addr>,
22}
23
24/// Persistent state for IP allocator
25#[derive(Debug, Clone, Serialize, Deserialize)]
26pub struct IpAllocatorState {
27    /// CIDR string
28    pub cidr: String,
29    /// List of allocated IPs
30    pub allocated: Vec<Ipv4Addr>,
31}
32
33impl IpAllocator {
34    /// Create a new IP allocator for the given CIDR range
35    ///
36    /// # Arguments
37    /// * `cidr` - Network CIDR notation (e.g., "10.200.0.0/16")
38    ///
39    /// # Errors
40    ///
41    /// Returns `OverlayError::InvalidCidr` if the CIDR string cannot be parsed.
42    ///
43    /// # Example
44    /// ```
45    /// use zlayer_overlay::allocator::IpAllocator;
46    ///
47    /// let allocator = IpAllocator::new("10.200.0.0/16").unwrap();
48    /// ```
49    pub fn new(cidr: &str) -> Result<Self> {
50        let network: Ipv4Net = cidr
51            .parse()
52            .map_err(|e| OverlayError::InvalidCidr(format!("{cidr}: {e}")))?;
53
54        Ok(Self {
55            network,
56            allocated: HashSet::new(),
57        })
58    }
59
60    /// Create an allocator from persisted state
61    ///
62    /// # Errors
63    ///
64    /// Returns an error if the CIDR is invalid or any IP is out of range.
65    pub fn from_state(state: IpAllocatorState) -> Result<Self> {
66        let mut allocator = Self::new(&state.cidr)?;
67        for ip in state.allocated {
68            allocator.mark_allocated(ip)?;
69        }
70        Ok(allocator)
71    }
72
73    /// Get the current state for persistence
74    #[must_use]
75    pub fn to_state(&self) -> IpAllocatorState {
76        IpAllocatorState {
77            cidr: self.network.to_string(),
78            allocated: self.allocated.iter().copied().collect(),
79        }
80    }
81
82    /// Load allocator state from a file
83    ///
84    /// # Errors
85    ///
86    /// Returns an error if the file cannot be read or the state is invalid.
87    pub async fn load(path: &Path) -> Result<Self> {
88        let contents = tokio::fs::read_to_string(path).await?;
89        let state: IpAllocatorState = serde_json::from_str(&contents)?;
90        Self::from_state(state)
91    }
92
93    /// Save allocator state to a file
94    ///
95    /// # Errors
96    ///
97    /// Returns an error if the file cannot be written or serialization fails.
98    pub async fn save(&self, path: &Path) -> Result<()> {
99        let state = self.to_state();
100        let contents = serde_json::to_string_pretty(&state)?;
101        tokio::fs::write(path, contents).await?;
102        Ok(())
103    }
104
105    /// Allocate the next available IP address
106    ///
107    /// Returns `None` if all addresses in the CIDR range are allocated.
108    ///
109    /// # Example
110    /// ```
111    /// use zlayer_overlay::allocator::IpAllocator;
112    ///
113    /// let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
114    /// let ip = allocator.allocate().unwrap();
115    /// assert_eq!(ip.to_string(), "10.200.0.1");
116    /// ```
117    pub fn allocate(&mut self) -> Option<Ipv4Addr> {
118        // Skip network address and broadcast address
119        for ip in self.network.hosts() {
120            if !self.allocated.contains(&ip) {
121                self.allocated.insert(ip);
122                return Some(ip);
123            }
124        }
125        None
126    }
127
128    /// Allocate a specific IP address
129    ///
130    /// # Errors
131    ///
132    /// Returns an error if the IP is already allocated or not in the CIDR range.
133    pub fn allocate_specific(&mut self, ip: Ipv4Addr) -> Result<()> {
134        if !self.network.contains(&ip) {
135            return Err(OverlayError::IpNotInRange(ip, self.network.to_string()));
136        }
137
138        if self.allocated.contains(&ip) {
139            return Err(OverlayError::IpAlreadyAllocated(ip));
140        }
141
142        self.allocated.insert(ip);
143        Ok(())
144    }
145
146    /// Allocate the first usable IP in the range (typically for the leader)
147    ///
148    /// # Example
149    /// ```
150    /// use zlayer_overlay::allocator::IpAllocator;
151    ///
152    /// let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
153    /// let ip = allocator.allocate_first().unwrap();
154    /// assert_eq!(ip.to_string(), "10.200.0.1");
155    /// ```
156    ///
157    /// # Errors
158    ///
159    /// Returns an error if no IPs are available or the first IP is already allocated.
160    pub fn allocate_first(&mut self) -> Result<Ipv4Addr> {
161        let first_ip = self
162            .network
163            .hosts()
164            .next()
165            .ok_or(OverlayError::NoAvailableIps)?;
166
167        if self.allocated.contains(&first_ip) {
168            return Err(OverlayError::IpAlreadyAllocated(first_ip));
169        }
170
171        self.allocated.insert(first_ip);
172        Ok(first_ip)
173    }
174
175    /// Mark an IP address as allocated (for restoring state)
176    ///
177    /// # Errors
178    ///
179    /// Returns an error if the IP is not in the CIDR range.
180    pub fn mark_allocated(&mut self, ip: Ipv4Addr) -> Result<()> {
181        if !self.network.contains(&ip) {
182            return Err(OverlayError::IpNotInRange(ip, self.network.to_string()));
183        }
184        self.allocated.insert(ip);
185        Ok(())
186    }
187
188    /// Release an IP address back to the pool
189    ///
190    /// Returns `true` if the IP was released, `false` if it wasn't allocated.
191    pub fn release(&mut self, ip: Ipv4Addr) -> bool {
192        self.allocated.remove(&ip)
193    }
194
195    /// Check if an IP address is allocated
196    #[must_use]
197    pub fn is_allocated(&self, ip: Ipv4Addr) -> bool {
198        self.allocated.contains(&ip)
199    }
200
201    /// Check if an IP address is within the CIDR range
202    #[must_use]
203    pub fn contains(&self, ip: Ipv4Addr) -> bool {
204        self.network.contains(&ip)
205    }
206
207    /// Get the number of allocated addresses
208    #[must_use]
209    pub fn allocated_count(&self) -> usize {
210        self.allocated.len()
211    }
212
213    /// Get the total number of usable addresses in the range
214    #[must_use]
215    #[allow(clippy::cast_possible_truncation)]
216    pub fn total_hosts(&self) -> u32 {
217        self.network.hosts().count() as u32
218    }
219
220    /// Get the number of available addresses
221    #[must_use]
222    #[allow(clippy::cast_possible_truncation)]
223    pub fn available_count(&self) -> u32 {
224        self.total_hosts()
225            .saturating_sub(self.allocated.len() as u32)
226    }
227
228    /// Get the CIDR string
229    #[must_use]
230    pub fn cidr(&self) -> String {
231        self.network.to_string()
232    }
233
234    /// Get the network address
235    #[must_use]
236    pub fn network_addr(&self) -> Ipv4Addr {
237        self.network.network()
238    }
239
240    /// Get the broadcast address
241    #[must_use]
242    pub fn broadcast_addr(&self) -> Ipv4Addr {
243        self.network.broadcast()
244    }
245
246    /// Get the prefix length
247    #[must_use]
248    pub fn prefix_len(&self) -> u8 {
249        self.network.prefix_len()
250    }
251
252    /// Get all allocated IPs
253    #[must_use]
254    pub fn allocated_ips(&self) -> Vec<Ipv4Addr> {
255        self.allocated.iter().copied().collect()
256    }
257}
258
259/// Helper function to get the first usable IP from a CIDR
260///
261/// # Errors
262///
263/// Returns an error if the CIDR is invalid or has no usable hosts.
264pub fn first_ip_from_cidr(cidr: &str) -> Result<Ipv4Addr> {
265    let network: Ipv4Net = cidr
266        .parse()
267        .map_err(|e| OverlayError::InvalidCidr(format!("{cidr}: {e}")))?;
268
269    network.hosts().next().ok_or(OverlayError::NoAvailableIps)
270}
271
272#[cfg(test)]
273mod tests {
274    use super::*;
275
276    #[test]
277    fn test_allocator_new() {
278        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
279        assert_eq!(allocator.cidr(), "10.200.0.0/24");
280        assert_eq!(allocator.allocated_count(), 0);
281    }
282
283    #[test]
284    fn test_allocator_invalid_cidr() {
285        let result = IpAllocator::new("invalid");
286        assert!(result.is_err());
287    }
288
289    #[test]
290    fn test_allocate_sequential() {
291        let mut allocator = IpAllocator::new("10.200.0.0/30").unwrap();
292
293        // /30 has 2 usable hosts (excluding network and broadcast)
294        let ip1 = allocator.allocate().unwrap();
295        let ip2 = allocator.allocate().unwrap();
296
297        assert_eq!(ip1.to_string(), "10.200.0.1");
298        assert_eq!(ip2.to_string(), "10.200.0.2");
299
300        // Should be exhausted
301        assert!(allocator.allocate().is_none());
302    }
303
304    #[test]
305    fn test_allocate_first() {
306        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
307
308        let first = allocator.allocate_first().unwrap();
309        assert_eq!(first.to_string(), "10.200.0.1");
310
311        // Can't allocate first again
312        assert!(allocator.allocate_first().is_err());
313    }
314
315    #[test]
316    fn test_allocate_specific() {
317        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
318
319        let specific_ip: Ipv4Addr = "10.200.0.50".parse().unwrap();
320        allocator.allocate_specific(specific_ip).unwrap();
321
322        assert!(allocator.is_allocated(specific_ip));
323
324        // Can't allocate same IP again
325        assert!(allocator.allocate_specific(specific_ip).is_err());
326    }
327
328    #[test]
329    fn test_allocate_specific_out_of_range() {
330        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
331
332        let out_of_range: Ipv4Addr = "192.168.1.1".parse().unwrap();
333        assert!(allocator.allocate_specific(out_of_range).is_err());
334    }
335
336    #[test]
337    fn test_release() {
338        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
339
340        let ip = allocator.allocate().unwrap();
341        assert!(allocator.is_allocated(ip));
342
343        assert!(allocator.release(ip));
344        assert!(!allocator.is_allocated(ip));
345
346        // Can allocate same IP again
347        let ip2 = allocator.allocate().unwrap();
348        assert_eq!(ip, ip2);
349    }
350
351    #[test]
352    fn test_mark_allocated() {
353        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
354
355        let ip: Ipv4Addr = "10.200.0.100".parse().unwrap();
356        allocator.mark_allocated(ip).unwrap();
357
358        assert!(allocator.is_allocated(ip));
359    }
360
361    #[test]
362    fn test_contains() {
363        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
364
365        assert!(allocator.contains("10.200.0.50".parse().unwrap()));
366        assert!(!allocator.contains("10.201.0.50".parse().unwrap()));
367    }
368
369    #[test]
370    fn test_total_hosts() {
371        // /24 has 254 usable hosts
372        let allocator = IpAllocator::new("10.200.0.0/24").unwrap();
373        assert_eq!(allocator.total_hosts(), 254);
374
375        // /30 has 2 usable hosts
376        let allocator = IpAllocator::new("10.200.0.0/30").unwrap();
377        assert_eq!(allocator.total_hosts(), 2);
378    }
379
380    #[test]
381    fn test_available_count() {
382        let mut allocator = IpAllocator::new("10.200.0.0/30").unwrap();
383
384        assert_eq!(allocator.available_count(), 2);
385
386        allocator.allocate();
387        assert_eq!(allocator.available_count(), 1);
388
389        allocator.allocate();
390        assert_eq!(allocator.available_count(), 0);
391    }
392
393    #[test]
394    fn test_state_roundtrip() {
395        let mut allocator = IpAllocator::new("10.200.0.0/24").unwrap();
396        allocator.allocate();
397        allocator.allocate();
398
399        let state = allocator.to_state();
400        let restored = IpAllocator::from_state(state).unwrap();
401
402        assert_eq!(allocator.cidr(), restored.cidr());
403        assert_eq!(allocator.allocated_count(), restored.allocated_count());
404    }
405
406    #[test]
407    fn test_first_ip_from_cidr() {
408        let ip = first_ip_from_cidr("10.200.0.0/24").unwrap();
409        assert_eq!(ip.to_string(), "10.200.0.1");
410    }
411}
zlayer_overlay/allocator.rs

zlayer_overlay/
allocator.rs
