toe_beans/v4/message/
socket.rs

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