toe-beans 0.10.0

use super::{Decodable, Encodable};
use crate::v4::{MAX_MESSAGE_SIZE, Result, UndecodedMessage};
use inherface::get_interfaces;
use log::{debug, info};
use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};

/// A connection to a UdpSocket that understands how to send/receive Deliverable.
/// Meant to be used by both Client and Server.
///
/// Currently wraps [UdpSocket::bind](https://doc.rust-lang.org/std/net/struct.UdpSocket.html#method.bind)
#[derive(Debug)]
pub struct Socket {
    socket: UdpSocket,
    broadcast: Ipv4Addr,
}

/// Bytes, received from the socket, that can be decoded into a Message.
pub type MessageBuffer = [u8; MAX_MESSAGE_SIZE];

impl Socket {
    /// A `MessageBuffer` with all 0 bytes.
    pub const EMPTY_MESSAGE_BUFFER: MessageBuffer = [0; MAX_MESSAGE_SIZE];

    /// Bind to an ip address/port and require that broadcast is enabled on the socket.
    ///
    /// Should be be called by both `Server::new` and `Client::new`, so this can be slower/panic
    /// since it is not in the `listen_once` hot path.
    pub fn new(address: SocketAddrV4, interface: Option<&String>) -> Self {
        let broadcast = Self::get_interface_broadcast(interface).unwrap_or(Ipv4Addr::BROADCAST);

        let socket = UdpSocket::bind(address).expect("failed to bind to address");
        info!("UDP socket bound on {}", address);

        // set_broadcast sets the SO_BROADCAST option on the socket
        // which is a way of ensuring that the application
        // can't _accidentally_ spam many devices with a broadcast
        // address without intentionally intending to do so.
        socket
            .set_broadcast(true)
            .expect("Failed to enable broadcasting");

        Self { socket, broadcast }
    }

    /// Tries to get the directed broadcast address for the interface with the given name.
    fn get_interface_broadcast(maybe_interface_name: Option<&String>) -> Option<Ipv4Addr> {
        let interface_name = maybe_interface_name?;
        let interfaces = get_interfaces().ok()?;
        let maybe_interface = interfaces.get(interface_name);
        let maybe_addr = maybe_interface?
            .v4_addr
            .iter()
            .find(|address| address.broadcast.is_some());
        let maybe_broadcast = maybe_addr?.broadcast;

        debug!(
            "Found ipv4 broadcast address ({:?}) in list of interface addresses",
            maybe_broadcast
        );

        maybe_broadcast
    }

    /// Returns the ip address of the bound socket.
    ///
    /// Can panic as typically called outside of `listen_once` hot path.
    pub fn get_ip(&self) -> Ipv4Addr {
        match self.socket.local_addr().unwrap().ip() {
            std::net::IpAddr::V4(ip) => ip,
            std::net::IpAddr::V6(_) => todo!("ipv6 is not supported yet"),
        }
    }

    /// Decodes received bytes into a "message" type that implements `DecodeMessage`
    /// and returns it and the source address.
    pub fn receive<D: Decodable<Output = D>>(&self) -> Result<(D, SocketAddr)> {
        // Any bytes over MAX_MESSAGE_SIZE will be discarded.
        let mut buffer = Self::EMPTY_MESSAGE_BUFFER;
        let (_, src) = match self.socket.recv_from(&mut buffer) {
            Ok(values) => values,
            Err(_) => return Err("Failed to receive data"),
        };

        let decoded = D::from_bytes(&UndecodedMessage::new(buffer));

        debug!("Received dhcp message (from {}): {:?}", src, decoded);

        Ok((decoded, src))
    }

    /// Returns received bytes directly without being decoded into a `Deliverable`,
    /// which allows you to partially decode them yourself later.
    pub fn receive_raw(&self) -> Result<(UndecodedMessage, SocketAddr)> {
        // Any bytes over MAX_MESSAGE_SIZE will be discarded.
        let mut buffer = Self::EMPTY_MESSAGE_BUFFER;
        let (_, src) = match self.socket.recv_from(&mut buffer) {
            Ok(values) => values,
            Err(_) => return Err("Failed to receive data"),
        };

        debug!("Received dhcp message from {}", src);

        Ok((UndecodedMessage::new(buffer), src))
    }

    /// Fills an empty undecoded message with passed bytes.
    /// Used to mock receiving a specific `UndecodedMessage` in tests, etc.
    pub fn receive_mock(partial_message: &[u8]) -> UndecodedMessage {
        let mut buffer = Self::EMPTY_MESSAGE_BUFFER;
        partial_message
            .iter()
            .enumerate()
            .for_each(|(i, byte)| buffer[i] = *byte);
        UndecodedMessage::new(buffer)
    }

    /// Converts a message to bytes and then sends it to the passed address.
    pub fn unicast<E: Encodable>(&self, encodable: &E, address: SocketAddrV4) -> Result<()> {
        let encoded = encodable.to_bytes();

        let ip = address.ip();
        debug!("Sending dhcp message (to {:?}): {:?}", ip, encodable);

        match self.socket.send_to(&encoded, address) {
            Ok(_) => Ok(()),
            Err(_) => Err("Failed to send data"),
        }
    }

    /// Send a message as bytes to many devices on the local network.
    pub fn broadcast<E: Encodable>(&self, encodable: &E, port: u16) -> Result<()> {
        self.unicast(encodable, SocketAddrV4::new(self.broadcast, port))
    }
}