mini_dhcp/

lib.rs

use anyhow::Context;
use dhcproto::v4::{
    Decodable, Decoder, DhcpOption, Encodable, Encoder, Message, MessageType, Opcode, OptionCode,
};
use jiff::{ToSpan, Unit, Zoned};
use rand::Rng;
use serde::Serialize;
use sqlx::sqlite::{SqliteConnectOptions, SqlitePool};
use std::net::Ipv4Addr;
use std::str::FromStr;
use tokio::net::UdpSocket;
use tracing::{error, info, warn};
mod db;
pub mod info;

#[allow(dead_code)]
#[derive(Debug)]
struct Lease {
    ip: i64,
    client_id: Vec<u8>,
    leased: bool,
    expires_at: i64,
    network: i64,
    probation: bool,
}

#[derive(Debug, Serialize)]
pub struct Client {
    pub ip: Ipv4Addr,
    pub client_id: String,
}

fn check_ip_in_range(addr: Ipv4Addr) -> bool {
    let [a, b, c, d] = addr.octets();
    a == 192 && b == 168 && c == 1 && (100..200).contains(&d)
}

fn get_ip_in_range() -> Ipv4Addr {
    let octet = rand::thread_rng().gen_range(100..200);
    Ipv4Addr::new(192, 168, 1, octet)
}

async fn insert_lease(pool: &SqlitePool, ip: Ipv4Addr, client_id: &Vec<u8>) -> anyhow::Result<()> {
    let ip = u32::from(ip);

    let expire_at = Zoned::now()
        .round(Unit::Second)?
        .checked_add(1.hour())
        .with_context(|| "Failed to calculate lease expiry time".to_string())?
        .timestamp()
        .as_second();

    sqlx::query_file!(
        "./db/queries/insert-new-lease.sql",
        ip,
        client_id,
        expire_at,
        0
    )
    .execute(pool)
    .await?;
    Ok(())
}

async fn build_dhcp_offer_packet(
    leases: &SqlitePool,
    discover_message: &Message,
) -> anyhow::Result<Message> {
    let client_id = discover_message.chaddr().to_vec();

    // The client's current address as recorded in the client's current
    // binding
    // The client's current address as recorded in the client's current
    // binding, ELSE
    //
    // The client's previous address as recorded in the client's (now
    // expired or released) binding, if that address is in the server's
    // pool of available addresses and not already allocated, ELSE
    let mut suggested_address = match db::get_ip_from_client_id(leases, &client_id).await {
        Ok(address) => {
            // Check if this lease has expired
            match db::get_lease_by_ip(leases, &address).await {
                Ok(lease) => {
                    let current_time = Zoned::now().round(Unit::Second)?.timestamp().as_second();
                    if lease.expires_at < current_time {
                        // Lease has expired, renew it
                        let new_expiry = Zoned::now()
                            .round(Unit::Second)?
                            .checked_add(1.hour())
                            .with_context(|| "Failed to calculate lease expiry time".to_string())?
                            .timestamp()
                            .as_second();
                        db::update_lease_expiry(leases, address, new_expiry).await?;
                        info!(
                            "[OFFER] Client {:?} has expired lease for {:?}, renewed",
                            client_id, address
                        );
                    } else {
                        info!(
                            "[OFFER] Client {:?} already has IP assigned: {:?}",
                            client_id, address
                        );
                    }
                    Some(address)
                }
                Err(_) => {
                    warn!("[OFFER] Could not fetch lease details for {:?}", address);
                    Some(address)
                }
            }
        }
        _ => {
            warn!(
                "[OFFER] Client {:?} has no IP assigned in the database",
                client_id
            );
            None
        }
    };

    // The address requested in the 'Requested IP Address' option, if that
    // address is valid and not already allocated, ELSE
    if suggested_address.is_none() {
        let requested_ip_address = discover_message.opts().get(OptionCode::RequestedIpAddress);
        info!(
            "[OFFER] Client requested IP address: {:?}",
            requested_ip_address
        );
        match requested_ip_address {
            Some(DhcpOption::RequestedIpAddress(ip)) => {
                if !db::is_ip_assigned(leases, *ip).await? && check_ip_in_range(*ip) {
                    suggested_address = Some(*ip);
                }
            }
            _ => {
                warn!("[OFFER] No requested IP address")
            }
        };
    }

    if suggested_address.is_none() {
        let mut max_tries: u8 = 10;
        loop {
            let random_address = get_ip_in_range();

            if !db::is_ip_assigned(leases, random_address).await? {
                suggested_address = Some(random_address);
                // insert the lease into the database
                insert_lease(leases, random_address, &client_id).await?;
                break;
            }

            if max_tries == 0 {
                return Err(anyhow::anyhow!("Could not assign IP address"));
            } else {
                max_tries = max_tries.saturating_sub(1);
            }
        }
    }

    let suggested_address = match suggested_address {
        Some(address) => address,
        None => return Err(anyhow::anyhow!("Could not assign IP address")),
    };

    info!("[OFFER] creating offer with IP {}", suggested_address);

    let mut offer = Message::default();

    let reply_opcode = Opcode::BootReply;
    offer.set_opcode(reply_opcode);
    offer.set_xid(discover_message.xid());
    offer.set_yiaddr(suggested_address);
    offer.set_siaddr(Ipv4Addr::new(192, 168, 1, 69));
    offer.set_flags(discover_message.flags());
    offer.set_giaddr(discover_message.giaddr());
    offer.set_chaddr(discover_message.chaddr());

    offer
        .opts_mut()
        .insert(DhcpOption::MessageType(MessageType::Offer));
    offer.opts_mut().insert(DhcpOption::AddressLeaseTime(3600));
    offer
        .opts_mut()
        .insert(DhcpOption::ServerIdentifier(Ipv4Addr::new(192, 168, 1, 69)));
    offer
        .opts_mut()
        .insert(DhcpOption::SubnetMask(Ipv4Addr::new(255, 255, 255, 0)));
    offer
        .opts_mut()
        .insert(DhcpOption::BroadcastAddr(Ipv4Addr::new(255, 255, 255, 255)));
    offer
        .opts_mut()
        .insert(DhcpOption::Router(vec![Ipv4Addr::new(192, 168, 1, 69)]));

    Ok(offer)
}

/// Response type for DHCP REQUEST handling
#[derive(Debug)]
enum DhcpResponse {
    Ack(Message),
    Nak(Message),
}

/// Build a DHCP NAK packet according to RFC 2131 Table 3
fn build_dhcp_nack_packet(request_message: &Message, reason: &str) -> Message {
    info!("[NAK] Sending NACK: {}", reason);

    let mut nak = Message::default();
    nak.set_opcode(Opcode::BootReply);
    nak.set_xid(request_message.xid());
    nak.set_flags(request_message.flags());
    nak.set_chaddr(request_message.chaddr());

    // RFC 2131 Table 3: yiaddr and siaddr must be 0 for NACK
    nak.set_yiaddr(Ipv4Addr::new(0, 0, 0, 0));
    nak.set_siaddr(Ipv4Addr::new(0, 0, 0, 0));

    // Only include MessageType and ServerIdentifier options
    nak.opts_mut()
        .insert(DhcpOption::MessageType(MessageType::Nak));
    nak.opts_mut()
        .insert(DhcpOption::ServerIdentifier(Ipv4Addr::new(192, 168, 1, 69)));

    nak
}

async fn build_dhcp_ack_packet(
    leases: &SqlitePool,
    request_message: &Message,
) -> anyhow::Result<DhcpResponse> {
    let server_identifier_option = request_message.opts().get(OptionCode::ServerIdentifier);

    let requested_ip_option = request_message.opts().get(OptionCode::RequestedIpAddress);

    // `ciaddr` is the client’s current IP address (used in RENEWING/REBINDING)
    let ciaddr = request_message.ciaddr();
    let chaddr = request_message.chaddr().to_owned();

    let (is_selecting, is_init_reboot, is_renewing_rebinding) = {
        let have_server_id = server_identifier_option.is_some();
        let have_requested_ip = requested_ip_option.is_some();
        let ciaddr_is_zero = ciaddr == Ipv4Addr::new(0, 0, 0, 0);

        // SELECTING state
        let selecting = have_server_id && have_requested_ip && ciaddr_is_zero;

        // INIT-REBOOT
        let init_reboot = !have_server_id && have_requested_ip && ciaddr_is_zero;

        // RENEWING/REBINDING: ciaddr != 0, no 'requested IP address'
        let renewing_rebinding = ciaddr != Ipv4Addr::new(0, 0, 0, 0) && !have_requested_ip;

        (selecting, init_reboot, renewing_rebinding)
    };

    let ip_to_validate = if is_selecting || is_init_reboot {
        match requested_ip_option {
            Some(DhcpOption::RequestedIpAddress(ip)) => {
                info!("[ACK] Client requested IP address: {:?}", ip);
                ip
            }
            _ => {
                return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                    request_message,
                    "Client didn't provide requested IP address",
                )));
            }
        }
    } else if is_renewing_rebinding {
        info!("[ACK] using ciaddr {:?}", ciaddr);
        &ciaddr
    } else {
        return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
            request_message,
            "DHCPREQUEST does not match any known valid state",
        )));
    };

    // 4) Validate that the IP is on the correct subnet (RFC says to NAK if it's on the wrong net).
    //    Also check if you have a valid lease for this client in your DB, etc.

    // First check if IP is in valid range
    if !check_ip_in_range(*ip_to_validate) {
        return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
            request_message,
            "Requested IP address is outside valid range",
        )));
    }

    let lease = match db::get_lease_by_ip(leases, ip_to_validate).await {
        Ok(lease) => Some(lease),
        Err(sqlx::Error::RowNotFound) => {
            // No lease exists - for INIT-REBOOT state, create one if IP is available
            if is_init_reboot {
                // Check if IP is already assigned to someone else
                if db::is_ip_assigned(leases, *ip_to_validate).await? {
                    warn!("[ACK] IP {:?} is assigned to another client", ip_to_validate);
                    return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                        request_message,
                        "IP address is assigned to another client",
                    )));
                }

                // IP is available, create a lease for this client
                info!("[ACK] Creating new lease for {:?} on INIT-REBOOT request", ip_to_validate);
                insert_lease(leases, *ip_to_validate, &chaddr).await?;
                None // Will create ACK below
            } else {
                warn!("[ACK] NO RECORD FOUND ON DB for non-INIT-REBOOT state");
                return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                    request_message,
                    "No lease record found in database",
                )));
            }
        }
        Err(e) => {
            warn!("[ACK] Database error: {:?}", e);
            return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                request_message,
                "Database error",
            )));
        }
    };

    // Validate the lease if it exists (skip if we just created it)
    if let Some(lease) = &lease {
        if !lease.leased {
            return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                request_message,
                "IP address is not currently leased",
            )));
        }

        // Check if lease has expired
        let current_time = Zoned::now().round(Unit::Second)?.timestamp().as_second();
        if lease.expires_at < current_time {
            warn!("[ACK] Lease has expired: expires_at={}, current={}", lease.expires_at, current_time);
            return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                request_message,
                "Lease has expired",
            )));
        }

        // Validate that the client_id matches the lease
        if lease.client_id != chaddr {
            warn!("[ACK] Client ID mismatch: lease has {:?}, request has {:?}", lease.client_id, chaddr);
            return Ok(DhcpResponse::Nak(build_dhcp_nack_packet(
                request_message,
                "IP address is leased to a different client",
            )));
        }
    }

    let mut ack = Message::default();
    ack.set_opcode(Opcode::BootReply);
    ack.set_xid(request_message.xid());
    ack.set_flags(request_message.flags());
    ack.set_giaddr(request_message.giaddr());
    ack.set_chaddr(&chaddr);

    ack.set_yiaddr(*ip_to_validate);

    ack.opts_mut()
        .insert(DhcpOption::MessageType(MessageType::Ack));
    ack.opts_mut()
        .insert(DhcpOption::ServerIdentifier(Ipv4Addr::new(192, 168, 1, 69)));

    ack.opts_mut().insert(DhcpOption::AddressLeaseTime(3600));
    ack.opts_mut()
        .insert(DhcpOption::SubnetMask(Ipv4Addr::new(255, 255, 255, 0)));
    ack.opts_mut()
        .insert(DhcpOption::BroadcastAddr(Ipv4Addr::new(255, 255, 255, 255)));
    ack.opts_mut()
        .insert(DhcpOption::Router(vec![Ipv4Addr::new(192, 168, 1, 69)]));

    // Update lease expiry time in database
    let new_expiry = Zoned::now()
        .round(Unit::Second)?
        .checked_add(1.hour())
        .with_context(|| "Failed to calculate new lease expiry time".to_string())?
        .timestamp()
        .as_second();

    db::update_lease_expiry(leases, *ip_to_validate, new_expiry).await?;

    Ok(DhcpResponse::Ack(ack))
}

#[derive(Clone)]
pub struct MiniDHCPConfiguration {
    interface: String,
    leases: SqlitePool,
}

impl MiniDHCPConfiguration {
    pub async fn new(interface: String) -> anyhow::Result<Self> {
        let conn = SqliteConnectOptions::from_str("sqlite://dhcp.db")?.create_if_missing(true);

        let leases = SqlitePool::connect_with(conn).await?;

        sqlx::migrate!("./db/migrations").run(&leases).await?;

        Ok(Self { leases, interface })
    }
}

async fn handle_discover(
    config: &MiniDHCPConfiguration,
    decoded_message: &Message,
) -> anyhow::Result<Vec<u8>> {
    let transaction_id = decoded_message.xid();
    let client_address = decoded_message.chaddr();
    info!("[{:X}] DISCOVER {:?}", transaction_id, client_address);
    let offer = build_dhcp_offer_packet(&config.leases, decoded_message);

    match offer.await {
        Ok(offer) => {
            let offered_ip = offer.yiaddr();
            info!(
                "[{:X}] [OFFER]: client {:?} ip {:?}",
                transaction_id, client_address, offered_ip
            );

            let mut buf = Vec::new();
            let mut e = Encoder::new(&mut buf);
            offer.encode(&mut e)?;
            Ok(buf)
        }
        Err(e) => {
            anyhow::bail!("OFFER Error: {:?}", e)
        }
    }
}

async fn handle_request(
    config: &MiniDHCPConfiguration,
    decoded_message: &Message,
) -> anyhow::Result<Vec<u8>> {
    let options = decoded_message.opts();
    let transaction_id = decoded_message.xid();
    let client_address = decoded_message.chaddr();
    let server_identifier = options.get(OptionCode::ServerIdentifier);
    info!(
        "[{:X}] REQUEST from {:?} to {:?}",
        transaction_id, client_address, server_identifier
    );

    let response = build_dhcp_ack_packet(&config.leases, decoded_message).await?;

    match response {
        DhcpResponse::Ack(ack) => {
            let offered_ip = ack.yiaddr();
            info!(
                "[{:X}] [ACK]: {:?} {:?}",
                transaction_id, client_address, offered_ip
            );

            let mut buf = Vec::new();
            let mut e = Encoder::new(&mut buf);
            ack.encode(&mut e)?;
            Ok(buf)
        }
        DhcpResponse::Nak(nak) => {
            info!("[{:X}] [NAK]: {:?}", transaction_id, client_address);

            let mut buf = Vec::new();
            let mut e = Encoder::new(&mut buf);
            nak.encode(&mut e)?;
            Ok(buf)
        }
    }
}

pub async fn start(config: MiniDHCPConfiguration) -> anyhow::Result<()> {
    let address = "0.0.0.0:67";
    info!("Starting DHCP listener [{}] {}", config.interface, address);
    let socket = UdpSocket::bind(address).await?;
    socket.set_broadcast(true)?;
    socket.bind_device(Some(config.interface.as_bytes()))?;

    loop {
        // Receive a packet
        let mut read_buffer = vec![0u8; 1024];
        let (_len, addr) = socket.recv_from(&mut read_buffer).await?;
        info!("== Received packet from {:?} ==", addr);

        let decoded_message = Message::decode(&mut Decoder::new(&read_buffer))?;
        // https://datatracker.ietf.org/doc/html/rfc2131#page-13
        // The 'op' field of each DHCP message sent from a client to a server contains BOOTREQUEST.
        if decoded_message.opcode() != Opcode::BootRequest {
            error!("[ERROR] opcode is not BootRequest, ignoring message");
            continue;
        }

        let options = decoded_message.opts();

        if options.has_msg_type(MessageType::Discover) {
            let transaction_id = decoded_message.xid();
            let response = handle_discover(&config, &decoded_message).await;
            if let Ok(response) = response {
                info!("[{:X}] [OFFER] Sending...", transaction_id);
                if let Err(e) = socket
                    .send_to(&response, "255.255.255.255:68")
                    .await
                {
                    error!("[{:X}] [OFFER] Failed to send in socket: {:?}", transaction_id, e);
                }
            } else {
                error!("[ERROR] handling DISCOVER {:?}", response);
            }
            continue;
        }

        if options.has_msg_type(MessageType::Request) {
            let transaction_id = decoded_message.xid();
            let response = handle_request(&config, &decoded_message).await;
            if let Ok(response) = response {
                info!("[{:X}] [ACK/NAK] Sending...", transaction_id);
                if let Err(e) = socket
                    .send_to(&response, "255.255.255.255:68")
                    .await
                {
                    error!("[{:X}] [ACK/NAK] Failed to send in socket: {:?}", transaction_id, e);
                }
            } else {
                error!("[ERROR] handling REQUEST {:?}", response);
            }
            continue;
        }

        if options.has_msg_type(MessageType::Decline) {
            let transaction_id = decoded_message.xid();
            let requested_ip = decoded_message.opts().get(OptionCode::RequestedIpAddress);

            if let Some(DhcpOption::RequestedIpAddress(ip)) = requested_ip {
                info!("[{:X}] [DECLINE] Client declined IP {:?} (address conflict detected)", transaction_id, ip);
                if let Err(e) = db::mark_ip_declined(&config.leases, *ip).await {
                    error!("[{:X}] [DECLINE] Failed to mark IP as declined: {:?}", transaction_id, e);
                } else {
                    info!("[{:X}] [DECLINE] IP {:?} marked as unavailable", transaction_id, ip);
                }
            } else {
                warn!("[{:X}] [DECLINE] No requested IP in DECLINE message", transaction_id);
            }
            continue;
        }

        if options.has_msg_type(MessageType::Release) {
            let transaction_id = decoded_message.xid();
            let client_id = decoded_message.chaddr().to_vec();
            let ciaddr = decoded_message.ciaddr();

            info!("[{:X}] [RELEASE] Client releasing IP {:?}", transaction_id, ciaddr);
            if let Err(e) = db::release_lease(&config.leases, ciaddr, &client_id).await {
                error!("[{:X}] [RELEASE] Failed to release lease: {:?}", transaction_id, e);
            } else {
                info!("[{:X}] [RELEASE] Lease for {:?} released", transaction_id, ciaddr);
            }
            continue;
        }
        if options.has_msg_type(MessageType::Inform) {
            let transaction_id = decoded_message.xid();
            info!("[{:X}] [INFORM]", transaction_id);
            continue;
        }
    }
}