toe-beans 0.11.0

mod config;
mod lease_time;

pub use config::*;
pub use lease_time::*;

use crate::v4::error::Result;
use ip_network::Ipv4Network;
use jiff::Zoned;
use log::{debug, error, info, warn};
use mac_address::MacAddress;
use std::fmt::Display;
use std::fs::{File, OpenOptions};
use std::io::{BufRead, BufReader, Seek, SeekFrom, Write};
use std::str::FromStr;
use std::{
    collections::{HashMap, HashSet},
    net::Ipv4Addr,
};

/// How much space is used by each lease in the leases file.
const PAGE_SIZE: usize = 200;

#[derive(Clone, Debug)]
enum LeaseStatus {
    Offered,
    Acked,
}

impl Display for LeaseStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            LeaseStatus::Offered => write!(f, "Offered"),
            LeaseStatus::Acked => write!(f, "Acked"),
        }
    }
}

impl From<&str> for LeaseStatus {
    fn from(value: &str) -> Self {
        match value {
            "Offered" => Self::Offered,
            "Acked" => Self::Acked,
            _ => panic!("Unknown value used for LeaseStatus"),
        }
    }
}

#[derive(Clone, Debug)]
struct Lease {
    /// What line in the leases file this lease occupies.
    /// Where 0 indicates it has not been assigned a line yet.
    line: u32,
    /// What ip address was assigned
    ip: Ipv4Addr,
    status: LeaseStatus,
    is_static: bool,
    /// A UTC time with time zone for when the lease was assigned/extended.
    when: Zoned,
}

impl Lease {
    #[inline]
    fn new(ip: Ipv4Addr, status: LeaseStatus, is_static: bool) -> Self {
        Self {
            line: 0,
            ip,
            status,
            is_static,
            when: Zoned::now(),
        }
    }

    /// Get back the entry that maps the mac address to the lease
    /// The passed string must exactly match the expected format.
    fn from_string(string: String, line: u32) -> Result<(MacAddress, Lease)> {
        let mut parts = string.split(',');
        // TODO handle unwraps
        let owner = MacAddress::from_str(parts.next().unwrap()).unwrap();
        let ip = Ipv4Addr::from_str(parts.next().unwrap()).unwrap();
        let status = LeaseStatus::from(parts.next().unwrap());
        let is_static = parts.next().unwrap().parse().unwrap();

        // The last value might have trailing spaces
        let when = Zoned::from_str(parts.next().unwrap().trim_end()).unwrap();

        let lease = Self {
            line,
            ip,
            status,
            is_static,
            when,
        };
        Ok((owner, lease))
    }

    /// Resets the time that this was considered leased.
    #[inline]
    fn extend(&mut self) {
        self.when = Zoned::now();
    }

    /// If the elapsed time since the assignment of the lease exceeds
    /// the allowed lease time then the lease is expired.
    fn is_expired(&self, lease_time: u32) -> bool {
        let elapsed = self.when.duration_until(&Zoned::now());
        elapsed.as_secs() >= lease_time as i64
    }

    fn to_string(&self, owner: MacAddress) -> String {
        // Some parts of this string are variable, but let's assume 100 for now.
        let mut string = String::with_capacity(PAGE_SIZE);

        string.push_str(&owner.to_string());
        string.push(',');
        string.push_str(&self.ip.to_string()); // variable
        string.push(',');
        string.push_str(&self.status.to_string()); // variable
        string.push(',');
        string.push_str(&self.is_static.to_string()); // variable
        string.push(',');
        string.push_str(&self.when.to_string()); // variable

        string
    }
}

/// Stores the leases and other information used across leasing functions in memory.
///
/// The leases are also written to a file on disk to preserve state in the event that something, such as a power outage, causes the server to stop.
/// This file is attempted to be read, parsed, and validated during `Leases` creation during `Server` start. After which it is only written to in order to sync it with its in-memory state.
/// Therefore, changing the file when the server is running will have no effect and may impact syncing to the file correctly.
#[derive(Debug)]
pub struct Leases {
    available: HashSet<Ipv4Addr>,
    leased: HashMap<MacAddress, Lease>,
    /// Passed to Leases from LeaseConfig.lease_time, so Leases can check for lease expiration.
    lease_time: u32,
    /// Keeps track of the number of lines written to the leases file.
    lines: u32,
    /// All configuration for `Leases`.
    config: LeaseConfig,
    /// A handle to the leases file.
    /// Will always be Some after `::new()` if `LeasesConfig.use_leases_file` is true.
    /// This is kept open for the duration of a `Leases` which is usually the lifetime of the `Server`.
    /// An exclusive file lock is attempted to reduce the chances that another process (or even another toe-beans server) will write to the same file. If you want to run multiple toe-beans servers then change the `LeasesConfig.leases_file_path` so that each writes to a different location.
    file: Option<File>,
}

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

    /// Offers only valid for about 5 minutes (300 seconds)
    /// at which point the offer's IP address may be reused.
    ///
    /// "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."
    const OFFER_TIMEOUT: u32 = 300;

    /// Once you pass a `LeaseConfig` to `Leases`, `Leases` privately owns it,
    /// but you can use this fn to borrow an immutable reference.
    #[inline]
    pub fn get_config(&self) -> &LeaseConfig {
        &self.config
    }

    fn empty(config: &LeaseConfig) -> Self {
        let size = config.network_cidr.hosts().len();
        Self {
            available: HashSet::with_capacity(size),
            leased: HashMap::with_capacity(size),
            lease_time: config.lease_time.as_server(),
            config: LeaseConfig::default(),
            lines: 0,
            file: None,
        }
    }

    /// Depending on whether the passed config enables `use_leases_file`:
    /// Restore and validate the leases file or create a new, empty `Leases`.
    pub fn new(config: LeaseConfig) -> Self {
        let file_path = config.leases_file_path.join(Self::FILE_NAME);

        let mut leases = if config.use_leases_file {
            let mut leases = Self::read_leases(&config).unwrap_or_else(|_| Self::empty(&config));

            let file = OpenOptions::new()
                .write(true)
                .create(true)
                .truncate(false)
                .open(file_path)
                .expect("Problem opening leases file for writing");
            file.lock().expect("Failed to acquire lock on leases file");
            leases.file = Some(file);

            leases
        } else {
            Self::empty(&config)
        };

        leases.config = config;
        leases.fill_available();

        leases
    }

    fn read_leases(config: &LeaseConfig) -> Result<Leases> {
        info!("Trying to read, parse, and validate leases file");

        let full_leases_file_path = config.leases_file_path.join(Self::FILE_NAME);
        let open_result = File::open(full_leases_file_path);
        let leased: HashMap<MacAddress, Lease> = match open_result {
            Ok(file) => {
                let mut line_num = 0;
                BufReader::new(file)
                    .lines()
                    .map(|line| {
                        let lease_string = line.expect("Problem reading leases file");
                        line_num += 1;
                        Lease::from_string(lease_string, line_num)
                            .expect("Problem parsing leases file")
                    })
                    .collect()
            }
            Err(message) => {
                error!("Problem opening leases file: {}", message);
                return Err("Problem opening leases file");
            }
        };

        let mut leases = Leases::empty(config);
        leases.leased = leased;

        leases.validate(&config.network_cidr)?;

        info!("Restored and validated leases file");

        Ok(leases)
    }

    /// Checks that both `static_leases` and `leased`:
    /// 1. Only contain ip addresses in the configured network range
    /// 2. Don't contain the broadcast or network address (which are unassignable)
    #[inline]
    fn validate(&self, network: &Ipv4Network) -> Result<()> {
        let broadcast_address = network.broadcast_address();
        let network_address = network.network_address();
        let static_lease_ips = self.config.static_leases.values();
        let leased_ips = self.leased.values().map(|lease| &lease.ip);

        static_lease_ips.chain(leased_ips).try_for_each(|ip| {
            if !network.contains(*ip) {
                return Err("The configured network range does not include the restored ip. Did the network range change?");
            }

            if ip == &broadcast_address {
                return Err("The network's broadcast address can't be leased");
            }

            if ip == &network_address {
                return Err("The network's network address can't be leased");
            }

            Ok(())
        })?;

        Ok(())
    }

    /// Only call this _after_ `leased` and `static_leases` are filled.
    fn fill_available(&mut self) {
        let network = self.config.network_cidr;
        let hosts = network.hosts();

        self.available.extend(hosts);

        // Dont assign the broadcast or network addresses
        self.available.remove(&network.broadcast_address());
        self.available.remove(&network.network_address());

        // Dont assign the already assigned addresses
        self.leased.values().for_each(|lease| {
            self.available.remove(&lease.ip);
        });

        // Dont assign the statically reserved addresses
        self.config.static_leases.values().for_each(|ip| {
            self.available.remove(ip);
        });
    }

    /// Writes leases to persistent storage as toe-beans-leases.toml
    fn commit(&mut self, mut lease: Lease, owner: MacAddress) -> Result<()> {
        if !self.config.use_leases_file {
            return Ok(());
        }

        // if line has not be assigned, assign next one:
        if lease.line == 0 {
            lease.line = self.lines + 1
        }

        // Convert string to a page of bytes for writing.
        // If string is less than the page size then remaining characters are spaces.
        let mut file_content: [u8; PAGE_SIZE] = [32; PAGE_SIZE];
        lease
            .to_string(owner)
            .as_bytes()
            .iter()
            .enumerate()
            .for_each(|(i, byte)| file_content[i] = *byte);
        // The last character in a page is a line break.
        file_content[PAGE_SIZE - 1] = 10;

        if self
            .file
            .as_mut()
            .expect("This should always be Some here")
            .seek(SeekFrom::Start((PAGE_SIZE * lease.line as usize) as u64))
            .is_err()
        {
            return Err("Problem writing to leases file");
        };

        if self
            .file
            .as_mut()
            .expect("This should always be Some here")
            .write_all(&file_content)
            .is_err()
        {
            return Err("Problem writing to leases file");
        };

        self.lines += 1;

        Ok(())
    }

    /// 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 available,
    ///   - _and_ there is not a static lease
    /// - Otherwise any available address (unless none are available).
    pub fn offer(&mut self, owner: MacAddress, requested_ip: Option<Ipv4Addr>) -> Result<Ipv4Addr> {
        if let Some(lease) = self.leased.get_mut(&owner) {
            match lease.status {
                LeaseStatus::Offered => {
                    warn!(
                        "{} will be offered an IP address that it has already been offered",
                        owner
                    );
                    lease.extend();
                    return Ok(lease.ip);
                }
                LeaseStatus::Acked => {
                    if !lease.is_expired(self.lease_time) {
                        warn!(
                            "{} will be offered its non-expired, acked IP address again",
                            owner
                        );
                        lease.status = LeaseStatus::Offered;
                        lease.extend();
                        return Ok(lease.ip);
                    }
                    // else if it is expired then continue to offer new address below...
                }
            }
        }

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

            if not_available {
                // we validated that available and leased were in range in `restore`
                debug!("Requested IP Address is not available (or maybe not in network range)");
            } else if has_static_lease {
                debug!("Requested IP Address ignored because owner has static lease");
            } else {
                self.available.remove(&requested_ip);
                self.leased
                    .insert(owner, Lease::new(requested_ip, LeaseStatus::Offered, false));
                return Ok(requested_ip);
            }
        }

        // else give an available ip address below...
        let lease = self.get_lease(&owner, LeaseStatus::Offered)?;
        let ip = lease.ip;
        self.leased.insert(owner, lease);
        Ok(ip)
    }

    /// Takes a chaddr's offered ip address and marks it as reserved
    /// then commits it to persistent storage and returns the address.
    pub fn ack(&mut self, owner: MacAddress) -> Result<Ipv4Addr> {
        let maybe_lease = self.leased.get_mut(&owner);
        let lease = match maybe_lease {
            Some(lease) => {
                match lease.status {
                    LeaseStatus::Offered => {
                        lease.status = LeaseStatus::Acked;
                    }
                    LeaseStatus::Acked => {
                        warn!(
                            "{} will be leased an IP address it was already leased",
                            owner
                        );
                    }
                };

                lease.extend();
                lease.to_owned()
            }
            None => {
                // nothing was offered, but this might be a rapid commit
                let lease = self.get_lease(&owner, LeaseStatus::Acked)?;
                self.leased.insert(owner, lease.clone());
                lease
            }
        };

        let ip = lease.ip;
        self.commit(lease, owner)?;

        Ok(ip)
    }

    /// Resets the time of an offered or acked lease.
    /// Used by the server in the lease renew/rebind process.
    pub fn extend(&mut self, owner: MacAddress) -> Result<()> {
        match self.leased.get_mut(&owner) {
            Some(lease) => {
                lease.extend();
                Ok(())
            }
            None => Err("No IP address lease found with that owner"),
        }
    }

    /// Makes an ip address that was offered or acked available again.
    pub fn release(&mut self, owner: MacAddress) -> Result<()> {
        let maybe_leased = self.leased.remove(&owner);

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

    /// Returns either a lease with a static address or one with the next available address
    fn get_lease(&mut self, owner: &MacAddress, status: LeaseStatus) -> Result<Lease> {
        let lease = match self.config.static_leases.get(owner) {
            Some(ip) => Lease::new(*ip, status, true),
            None => Lease::new(self.get_ip()?, status, false),
        };
        Ok(lease)
    }

    /// Returns 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 get_ip(&mut self) -> Result<Ipv4Addr> {
        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());
        }

        let maybe_expired_lease = self
            .leased
            .clone() // TODO slow
            .into_iter() // TODO slow
            .find(|(_owner, lease)| {
                let expiration = match lease.status {
                    LeaseStatus::Offered => Self::OFFER_TIMEOUT,
                    LeaseStatus::Acked => self.lease_time,
                };
                lease.is_expired(expiration)
            });

        if let Some(expired_lease) = maybe_expired_lease {
            debug!("Reusing expired lease's IP address");
            // UNWRAP is okay because we've found this element above
            return Ok(self.leased.remove(&expired_lease.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.leased.get(&owner);

        match maybe_leased {
            Some(lease) => {
                if let LeaseStatus::Offered = lease.status {
                    return Err("Lease offered but not previously acked");
                }

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

                if lease.is_expired(self.lease_time) {
                    return Err("Lease has expired");
                }

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