toe_beans/v4/message/
socket.rs

1use super::{Decodable, Encodable};
2use crate::v4::{MAX_MESSAGE_SIZE, Result, UndecodedMessage};
3use inherface::get_interfaces;
4use log::{debug, info};
5use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4, UdpSocket};
6
7/// A connection to a UdpSocket that understands how to send/receive Deliverable.
8/// Meant to be used by both Client and Server.
9///
10/// Currently wraps [UdpSocket::bind](https://doc.rust-lang.org/std/net/struct.UdpSocket.html#method.bind)
11#[derive(Debug)]
12pub struct Socket {
13    socket: UdpSocket,
14    broadcast: Ipv4Addr,
15}
16
17/// Bytes, received from the socket, that can be decoded into a Message.
18pub type MessageBuffer = [u8; MAX_MESSAGE_SIZE];
19
20impl Socket {
21    /// A `MessageBuffer` with all 0 bytes.
22    pub const EMPTY_MESSAGE_BUFFER: MessageBuffer = [0; MAX_MESSAGE_SIZE];
23
24    /// Bind to an ip address/port and require that broadcast is enabled on the socket.
25    ///
26    /// Should be be called by both `Server::new` and `Client::new`, so this can be slower/panic
27    /// since it is not in the `listen_once` hot path.
28    pub fn new(address: SocketAddrV4, interface: Option<&String>) -> Self {
29        let broadcast = Self::get_interface_broadcast(interface).unwrap_or(Ipv4Addr::BROADCAST);
30
31        let socket = UdpSocket::bind(address).expect("failed to bind to address");
32        info!("UDP socket bound on {}", address);
33
34        // set_broadcast sets the SO_BROADCAST option on the socket
35        // which is a way of ensuring that the application
36        // can't _accidentally_ spam many devices with a broadcast
37        // address without intentionally intending to do so.
38        socket
39            .set_broadcast(true)
40            .expect("Failed to enable broadcasting");
41
42        Self { socket, broadcast }
43    }
44
45    fn get_interface_broadcast(maybe_interface_name: Option<&String>) -> Option<Ipv4Addr> {
46        let interface = maybe_interface_name?;
47        let interfaces = get_interfaces().ok()?;
48        let maybe_interface = interfaces.get(interface);
49        let maybe_addr = maybe_interface?
50            .v4_addr
51            .iter()
52            .find(|address| address.broadcast.is_some());
53        let maybe_broadcast = maybe_addr?.broadcast;
54
55        debug!(
56            "Found ipv4 broadcast address ({:?}) in list of interface addresses",
57            maybe_broadcast
58        );
59
60        maybe_broadcast
61    }
62
63    /// Returns the ip address of the bound socket.
64    ///
65    /// Can panic as typically called outside of `listen_once` hot path.
66    pub fn get_ip(&self) -> Ipv4Addr {
67        match self.socket.local_addr().unwrap().ip() {
68            std::net::IpAddr::V4(ip) => ip,
69            std::net::IpAddr::V6(_) => todo!("ipv6 is not supported yet"),
70        }
71    }
72
73    /// Decodes received bytes into a "message" type that implements `DecodeMessage`
74    /// and returns it and the source address.
75    pub fn receive<D: Decodable<Output = D>>(&self) -> Result<(D, SocketAddr)> {
76        // Any bytes over MAX_MESSAGE_SIZE will be discarded.
77        let mut buffer = Self::EMPTY_MESSAGE_BUFFER;
78        let (_, src) = match self.socket.recv_from(&mut buffer) {
79            Ok(values) => values,
80            Err(_) => return Err("Failed to receive data"),
81        };
82
83        let decoded = D::from_bytes(&UndecodedMessage::new(buffer));
84
85        debug!("Received dhcp message (from {}): {:?}", src, decoded);
86
87        Ok((decoded, src))
88    }
89
90    /// Returns received bytes directly without being decoded into a `Deliverable`,
91    /// which allows you to partially decode them yourself later.
92    pub fn receive_raw(&self) -> Result<(UndecodedMessage, SocketAddr)> {
93        // Any bytes over MAX_MESSAGE_SIZE will be discarded.
94        let mut buffer = Self::EMPTY_MESSAGE_BUFFER;
95        let (_, src) = match self.socket.recv_from(&mut buffer) {
96            Ok(values) => values,
97            Err(_) => return Err("Failed to receive data"),
98        };
99
100        debug!("Received dhcp message from {}", src);
101
102        Ok((UndecodedMessage::new(buffer), src))
103    }
104
105    /// Fills an empty undecoded message with passed bytes.
106    /// Used to mock receiving a specific `UndecodedMessage` in tests, etc.
107    pub fn receive_mock(partial_message: &[u8]) -> UndecodedMessage {
108        let mut buffer = Self::EMPTY_MESSAGE_BUFFER;
109        partial_message
110            .iter()
111            .enumerate()
112            .for_each(|(i, byte)| buffer[i] = *byte);
113        UndecodedMessage::new(buffer)
114    }
115
116    /// Converts a message to bytes and then sends it to the passed address.
117    pub fn unicast<E: Encodable>(&self, encodable: &E, address: SocketAddrV4) -> Result<()> {
118        let encoded = encodable.to_bytes();
119
120        let ip = address.ip();
121        debug!("Sending dhcp message (to {:?}): {:?}", ip, encodable);
122
123        match self.socket.send_to(&encoded, address) {
124            Ok(_) => Ok(()),
125            Err(_) => Err("Failed to send data"),
126        }
127    }
128
129    /// Send a message as bytes to many devices on the local network.
130    pub fn broadcast<E: Encodable>(&self, encodable: &E, port: u16) -> Result<()> {
131        self.unicast(encodable, SocketAddrV4::new(self.broadcast, port))
132    }
133}
toe_beans/v4/message/socket.rs

toe_beans/v4/message/
socket.rs
