toe_beans/v4/message/
socket.rs

1use super::Deliverable;
2use crate::v4::Result;
3use inherface::get_interfaces;
4use log::{debug, info};
5use std::net::{IpAddr, Ipv4Addr, SocketAddr, SocketAddrV4, ToSocketAddrs, 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)
11pub struct Socket {
12    socket: UdpSocket,
13    broadcast: Ipv4Addr,
14}
15
16impl Socket {
17    /// Bind to an ip address/port and require that broadcast is enabled on the socket.
18    ///
19    /// Should be be called by both `Server::new` and `Client::new`, so this can be slower/panic
20    /// since it is not in the `listen_once` hot path.
21    pub fn new(address: SocketAddrV4, interface: Option<&String>) -> Self {
22        // hack to fix integration tests
23        let broadcast = if cfg!(feature = "integration") {
24            Ipv4Addr::LOCALHOST
25        } else {
26            Self::get_interface_broadcast(interface).unwrap_or(Ipv4Addr::BROADCAST)
27        };
28
29        let socket = UdpSocket::bind(address).expect("failed to bind to address");
30        info!("UDP socket bound on {}", address);
31
32        // set_broadcast sets the SO_BROADCAST option on the socket
33        // which is a way of ensuring that the application
34        // can't _accidentally_ spam many devices with a broadcast
35        // address without intentionally intending to do so.
36        socket
37            .set_broadcast(true)
38            .expect("Failed to enable broadcasting");
39
40        Self { socket, broadcast }
41    }
42
43    fn get_interface_broadcast(maybe_interface_name: Option<&String>) -> Option<Ipv4Addr> {
44        let interface = maybe_interface_name?;
45        let interfaces = get_interfaces().ok()?;
46        let maybe_interface = interfaces.get(interface);
47        let maybe_addr = maybe_interface?
48            .v4_addr
49            .iter()
50            .find(|address| address.broadcast.is_some());
51
52        debug!(
53            "found ipv4 broadcast address ({:?}) in list of addresses for interface",
54            maybe_addr
55        );
56
57        maybe_addr?.broadcast
58    }
59
60    /// Returns the ip address of the bound socket.
61    ///
62    /// Can panic as typically called outside of `listen_once` hot path.
63    pub fn get_ip(&self) -> Ipv4Addr {
64        match self.socket.local_addr().unwrap().ip() {
65            std::net::IpAddr::V4(ip) => ip,
66            std::net::IpAddr::V6(_) => todo!("ipv6 is not supported yet"),
67        }
68    }
69
70    /// Converts received bytes into a Deliverable and returns it and the source address.
71    pub fn receive<M: Deliverable<Output = M>>(&self) -> Result<(M, SocketAddr)> {
72        // If `buf` is too small to hold the message, it will be cut off.
73        let mut bytes = [0; 548];
74        let (_, src) = match self.socket.recv_from(&mut bytes) {
75            Ok(values) => values,
76            Err(_) => return Err("Failed to receive data"),
77        };
78
79        let decoded = match M::from_bytes(&bytes) {
80            Ok(message) => message,
81            Err(_) => return Err("Failed decoding message"),
82        };
83
84        debug!("Received dhcp message (from {}): {:?}", src, decoded);
85
86        Ok((decoded, src))
87    }
88
89    /// Converts a message to bytes and then sends it to the passed address.
90    pub fn unicast<A: ToSocketAddrs, M: Deliverable>(&self, message: &M, address: A) -> Result<()> {
91        let encoded = match message.to_bytes() {
92            Ok(bytes) => bytes,
93            Err(_) => return Err("Failed encoding message"),
94        };
95
96        let address = address.to_socket_addrs().unwrap().next().unwrap();
97        debug!("Sending dhcp message (to {:?}): {:?}", address, message);
98
99        match self.socket.send_to(&encoded, address) {
100            Ok(_) => Ok(()),
101            Err(_) => Err("Failed to send data"),
102        }
103    }
104
105    /// Send a message as bytes to many devices on the local network.
106    pub fn broadcast<M: Deliverable>(&self, message: &M, port: u16) -> Result<()> {
107        self.unicast(message, SocketAddr::new(IpAddr::V4(self.broadcast), port))
108    }
109}