toe_beans/v4/server/

ip_pool.rs

use super::Config;
use crate::v4::error::Result;
use ipnetwork::Ipv4Network;
use log::{debug, info, warn};
use mac_address::MacAddress;
use serde::{Deserialize, Serialize};
use std::fs::read_to_string;
use std::io::Write;
use std::os::unix::fs::OpenOptionsExt;
use std::path::PathBuf;
use std::time::{Duration, SystemTime};
use std::{
    collections::{HashMap, HashSet},
    fs::OpenOptions,
    net::Ipv4Addr,
};

#[derive(Deserialize, Serialize, Clone)]
struct Lease {
    ip: Ipv4Addr,
    // TODO this time measurement is not monotonic.
    when: SystemTime,
}

/// Used to manage the dynamic collection of IP addresses that the server leases.
// TODO Use a CIDR Trie for this?
#[derive(Deserialize, Serialize)]
pub struct IpPool {
    #[serde(skip)]
    available: HashSet<Ipv4Addr>,
    #[serde(skip)]
    offered: HashMap<MacAddress, Lease>,
    acked: HashMap<MacAddress, Lease>,
    /// The server will always assign these ip adddresses to these mac addresses.
    /// All static leases must be within the network_cidr range.
    static_leases: HashMap<MacAddress, Ipv4Addr>,
    /// Passed to IpPool from Config.network_cidr.
    network: Ipv4Network,
    /// Passed to IpPool from Config.lease_time, so IpPool check for lease expiration.
    lease_time: u32,
    #[serde(skip)]
    /// Passed to IpPool from Config.path because they share a single directory.
    file_path: PathBuf,
}

impl IpPool {
    /// The file name that the IpPool is read from and written to.
    pub const FILE_NAME: &'static str = "toe-beans-leases.toml";

    /// Create a pool with the capacity of the configured network's size.
    pub fn new(config: &Config) -> Self {
        let network = config.network_cidr;
        let size = network.size() as usize;
        let mut available = HashSet::with_capacity(size);
        available.extend(network.iter());
        let offered = HashMap::with_capacity(size);
        let acked = HashMap::with_capacity(size);
        let static_leases = HashMap::new();

        Self {
            file_path: config.path.clone(),
            available,
            offered,
            acked,
            static_leases,
            network,
            lease_time: config.lease_time.as_server(),
        }
    }

    /// Restores the state of an IpPool from a toe-beans-leases.toml
    /// as long as the IpPool's network (range/capacity) has not changed.
    /// Static leases are verified to be within network range.
    ///
    /// This does not restore dhcp offers.
    pub fn restore(config: &Config) -> Result<Self> {
        let network = config.network_cidr;
        let read_result = read_to_string(config.path.join(Self::FILE_NAME));

        let mut ip_pool: IpPool = match read_result {
            Ok(toml_string) => {
                let toml_result = toml::from_str(&toml_string);
                match toml_result {
                    Ok(ip_pool) => ip_pool,
                    Err(_) => return Err("Problem parsing leases file"),
                }
            }
            Err(_) => {
                return Err("Problem reading leases file");
            }
        };

        if network != ip_pool.network {
            return Err(
                "The restored and configured IP address pool's network range/capacity do not match",
            );
        }

        ip_pool.static_leases.values().try_for_each(|static_ip| {
            if !network.contains(*static_ip) {
                return Err("The configured network range does not include the static ip");
            }

            Ok(())
        })?;

        ip_pool.available.extend(network.iter());

        // TODO remove restored "acked" from "available"

        info!("Restored IP address leases");
        Ok(ip_pool)
    }

    /// Use the toe-beans-leases.toml to restore past leases
    /// or return a new IpPool if that was not successful.
    pub fn restore_or_new(config: &Config) -> Self {
        IpPool::restore(config).unwrap_or_else(|_| IpPool::new(config))
    }

    /// Writes leases to persistent storage as toe-beans-leases.toml
    fn commit(&self) -> Result<()> {
        debug!("Writing {}", Self::FILE_NAME);

        let file_content = match toml::to_string_pretty(&self) {
            Ok(content) => content,
            Err(_) => return Err("Failed to generate toml data"),
        };

        let open_result = OpenOptions::new()
            .read(false)
            .write(true)
            .create(true)
            .truncate(true)
            .mode(0o644) // ensure that file permissions are set before creating file
            .open(self.file_path.join(Self::FILE_NAME));

        let mut file = match open_result {
            Ok(file) => file,
            Err(_) => return Err("Failed to open file for writing"),
        };

        match file.write_all(file_content.as_bytes()) {
            Ok(_) => Ok(()),
            Err(_) => Err("Failed to write to file"),
        }
    }

    /// Takes an available IP address, marks it as offered, and returns it.
    ///
    /// Returns one of:
    /// - The previously offered address, if one has been offered.
    /// - The requested address:
    ///   - If one was requested,
    ///   - _and_ it is in the pool's range,
    ///   - _and_ it is available
    /// - Otherwise any available address (unless none are available).
    pub fn offer(&mut self, owner: MacAddress, requested_ip: Option<Ipv4Addr>) -> Result<Ipv4Addr> {
        if self.acked.contains_key(&owner) && !self.is_expired(&owner)? {
            return Err("This device already has already been leased a non-expired IP address");
        }

        if let Some(offered) = self.offered.get(&owner) {
            warn!(
                "{} will be offered an IP address that it has already been offered",
                owner
            );
            return Ok(offered.ip);
        }

        if let Some(requested_ip) = requested_ip {
            let ip_in_pool = self.network.contains(requested_ip);
            let not_available = !self.available.contains(&requested_ip);
            let has_static_lease = self.static_leases.contains_key(&owner);

            if !ip_in_pool {
                debug!("Requested IP Address is not in network range");
            } else if not_available {
                debug!("Requested IP Address is not available");
            } else if has_static_lease {
                debug!("Requested IP Address ignored because owner has static lease");
            } else {
                self.available.remove(&requested_ip);
                self.offered.insert(
                    owner,
                    Lease {
                        ip: requested_ip,
                        when: SystemTime::now(),
                    },
                );
                return Ok(requested_ip);
            }
        }

        // else give an available ip address below...
        let ip = self.find_available_ip(&owner)?;
        self.offered.insert(
            owner,
            Lease {
                ip,
                when: SystemTime::now(),
            },
        );
        Ok(ip)
    }

    /// Takes a chaddr's offered ip address and marks it as reserved
    /// then commits it to persistent storage and returns the address.
    ///
    /// Currently only one IP is allowed per chaddr.
    pub fn ack(&mut self, owner: MacAddress) -> Result<Ipv4Addr> {
        if self.acked.contains_key(&owner) && !self.is_expired(&owner)? {
            return Err("This device already has already been leased a non-expired IP address");
        }

        let maybe_offered_ip = self.offered.remove(&owner);
        let ip = match maybe_offered_ip {
            Some(lease) => lease.ip,
            None => {
                // nothing was offered, but this might be a rapid commit
                self.find_available_ip(&owner)?
            }
        };

        self.acked.insert(
            owner,
            Lease {
                ip,
                when: SystemTime::now(),
            },
        );

        self.commit()?;

        Ok(ip)
    }

    /// Resets the time of an acked lease. Used in the lease renew/rebind process.
    pub fn extend(&mut self, owner: MacAddress, ip: Ipv4Addr) -> Result<()> {
        match self.acked.insert(
            owner,
            Lease {
                ip,
                when: SystemTime::now(),
            },
        ) {
            Some(_) => Ok(()),
            None => Err("No IP address lease found with that owner"),
        }
    }

    /// Makes an ip address that was reserved during an offer available again.
    pub fn unoffer(&mut self, owner: MacAddress) -> Result<()> {
        let maybe_offered = self.offered.remove(&owner);

        match maybe_offered {
            Some(lease) => {
                self.available.insert(lease.ip);
                Ok(())
            }
            None => Err("No IP address offer found to unoffer"),
        }
    }

    /// Makes an ip address that was acked available again.
    // TODO require IP address to release as well?
    pub fn release(&mut self, owner: MacAddress) -> Result<()> {
        let maybe_leased = self.acked.remove(&owner);

        match maybe_leased {
            Some(lease) => {
                self.available.insert(lease.ip);
                Ok(())
            }
            None => Err("No IP address lease found with that owner"),
        }
    }

    /// Returns a static ip address or gets the next available address.
    /// If there are no more available addresses, then it will:
    /// 1. Search for offered addresses older than 1 minute
    /// 2. Search for acked addresses that have expired.
    fn find_available_ip(&mut self, owner: &MacAddress) -> Result<Ipv4Addr> {
        if let Some(ip) = self.static_leases.get(owner) {
            return match self.available.take(ip) {
                Some(ip) => Ok(ip),
                None => Err("A static lease was found but is not available"),
            };
        }

        if let Some(any) = self.available.iter().next().cloned() {
            debug!("Chose available IP address {any}");
            // UNWRAP is okay because we've found this element above
            return Ok(self.available.take(&any).unwrap());
        }

        // "Because the servers have not committed any network address assignments"
        // "on the basis of a DHCPOFFER, server are free to reuse offered network addresses in response to subsequent requests."
        // "Servers SHOULD NOT reuse offered addresses and may use an implementation-specific timeout mechanism to decide when to reuse an offered address."
        //
        // TODO collect all expired offers into a Vec and add them back to available except for the last which should be returned instead.
        let maybe_expired_offer = self
            .offered
            .clone() // TODO slow
            .into_iter() // TODO slow
            .find(|offer| {
                let elapsed_result = offer.1.when.elapsed();
                match elapsed_result {
                    Ok(elapsed) => elapsed > Duration::from_secs(300),
                    Err(_) => false,
                }
            });

        if let Some(expired_offer) = maybe_expired_offer {
            debug!("Found an expired offer's IP address");
            // UNWRAP is okay because we've found this element above
            return Ok(self.offered.remove(&expired_offer.0).unwrap().ip);
        }

        // TODO collect all expired acks into a Vec and add them back to available except for the last which should be returned instead.
        let maybe_expired_ack = self
            .acked
            .clone() // TODO slow
            .into_iter() // TODO slow
            .find(|ack| {
                let elapsed_result = ack.1.when.elapsed();
                match elapsed_result {
                    Ok(elapsed) => elapsed > Duration::from_secs(300),
                    Err(_) => false,
                }
            });

        if let Some(expired_ack) = maybe_expired_ack {
            debug!("Found an expired ack's IP address");
            // UNWRAP is okay because we've found this element above
            return Ok(self.acked.remove(&expired_ack.0).unwrap().ip);
        }

        Err("No more IP addresses available")
    }

    /// Checks whether the passed IP address is available in the pool.
    pub fn is_available(&self, ip: &Ipv4Addr) -> bool {
        self.available.contains(ip)
    }

    /// Checks if the passed IP address matches the committed, leased IP address,
    /// and that the lease is not expired.
    pub fn verify_lease(&self, owner: MacAddress, ip: &Ipv4Addr) -> Result<()> {
        let maybe_leased = self.acked.get(&owner);

        match maybe_leased {
            Some(lease) => {
                if &lease.ip != ip {
                    return Err("Client's notion of ip address is wrong");
                }

                if self.is_expired(&owner)? {
                    return Err("Lease has expired");
                }

                Ok(())
            }
            None => Err("No IP address lease found with that owner"),
        }
    }

    /// Looks up a lease, if one is found, and checks if it has expired.
    pub fn is_expired(&self, owner: &MacAddress) -> Result<bool> {
        let maybe_leased = self.acked.get(owner);

        match maybe_leased {
            Some(lease) => match lease.when.elapsed() {
                Ok(elapsed) => Ok(elapsed >= Duration::from_secs(self.lease_time as u64)),
                Err(_) => Err("Problem getting elapsed system time"),
            },
            None => Err("No IP address lease found with that owner"),
        }
    }
}