reliable shutdown

src/netif.rs

@@ -12,8 +12,10 @@ use crate::mac::Mac;
 pub struct LocalNetwork {
     // 192.168.0.7
     pub source_ip_hint: Ipv4Addr,
-    // 192.168.0.255 = source_ip_hint | ~netmask
+    // 192.168.0.255 = source_ip_hint | ~subnet_mask
     pub subnet_broadcast: Ipv4Addr,
+    pub subnet_mask: Ipv4Addr,
+    pub subnet_addr: Ipv4Addr,
 }
 
 pub fn get_local_networks() -> WResult<Vec<LocalNetwork>> {
@@ -33,6 +35,8 @@ pub fn get_local_networks() -> WResult<Vec<LocalNetwork>> {
             let local_network = LocalNetwork {
                 source_ip_hint: i.ip,
                 subnet_broadcast: i.broadcast,
+                subnet_mask: i.mask,
+                subnet_addr: i.address,
             };
             networks.push(local_network);
         }
@@ -41,7 +45,7 @@ pub fn get_local_networks() -> WResult<Vec<LocalNetwork>> {
     Ok(networks)
 }
 
-pub fn find_ip_with_mac(mac: Mac) -> WResult<Ipv4Addr> {
+pub fn find_ip_with_mac(mac: &Mac) -> WResult<Ipv4Addr> {
     let Some(ip) = find_ip_from_arp(&mac, "/proc/net/arp") else {
         return Err(WError::MacNotInNetwork(mac.to_string()));
     };
@@ -150,11 +154,12 @@ fn get_network_interfaces() -> WResult<Vec<NetworkInterface>> {
                 let sa_ip = ifa.ifa_addr as *const SockAddrIn;
                 let sa_netmask = ifa.ifa_netmask as *const SockAddrIn;
                 let ip = read_sa_ip(sa_ip);
-                let netmask = read_sa_ip(sa_netmask);
+                let mask = read_sa_ip(sa_netmask);
-                let broadcast = calculate_broadcast_address(ip, netmask);
+                let broadcast = calculate_broadcast_address(ip, mask);
                 let network_address = NetworkAddress {
                     source_ip_hint: ip,
                     broadcast,
+                    mask,
                 };
 
                 network_addresses.insert(name, network_address);
@@ -178,10 +183,13 @@ fn get_network_interfaces() -> WResult<Vec<NetworkInterface>> {
     // Assume all link_addresses have unique names for now
     for (n, m) in link_addresses {
         if let Some(n_addr) = network_addresses.get(&n) {
+            let network_address = calculate_network_address(n_addr);
             let network_interface = NetworkInterface::create(
                 n,
                 n_addr.source_ip_hint,
                 n_addr.broadcast,
+                network_address,
+                n_addr.mask,
                 m.mac,
             );
             network_interfaces.push(network_interface);
@@ -238,6 +246,7 @@ unsafe extern "C" {
 struct NetworkAddress {
     source_ip_hint: Ipv4Addr,
     broadcast: Ipv4Addr,
+    mask: Ipv4Addr,
 }
 
 struct LinkAddress {
@@ -245,9 +254,16 @@ struct LinkAddress {
 }
 
 struct NetworkInterface {
+    // enp42s0, wlan0, wlo1, eth0
     name: String,
+    // IP of this device - 192.168.0.8
     ip: Ipv4Addr,
+    // Broadcast address of the network - 192.168.0.255
     broadcast: Ipv4Addr,
+    // Network address - 192.168.0.0
+    address: Ipv4Addr,
+    // Netmask for the network - 255.255.255.0
+    mask: Ipv4Addr,
     _mac: [u8; 6],
 }
 
@@ -256,29 +272,31 @@ impl NetworkInterface {
         name: String,
         ip: Ipv4Addr,
         broadcast: Ipv4Addr,
+        address: Ipv4Addr,
+        mask: Ipv4Addr,
         mac: [u8; 6],
     ) -> Self {
         Self {
             name,
             ip,
             broadcast,
+            address,
+            mask,
             _mac: mac,
         }
     }
 }
 
-fn calculate_broadcast_address(ip: Ipv4Addr, netmask: Ipv4Addr) -> Ipv4Addr {
+fn calculate_network_address(n_addr: &NetworkAddress) -> Ipv4Addr {
-    let ip_octets = ip.octets();
+    let network_address = n_addr.source_ip_hint.to_bits() & n_addr.mask.to_bits();
-    let netmask_octets = netmask.octets();
+    Ipv4Addr::from_bits(network_address)
-
-    let mut broadcast_octets = [0u8; 4];
-    let mut broadcast_octet_index = 0;
-    for (i, m) in ip_octets.iter().zip(netmask_octets) {
-        broadcast_octets[broadcast_octet_index] = i | 255 - m;
-        broadcast_octet_index += 1;
 }
 
-    Ipv4Addr::from_octets(broadcast_octets)
+fn calculate_broadcast_address(ip: Ipv4Addr, netmask: Ipv4Addr) -> Ipv4Addr {
+    let neg_netmask = u32::MAX - netmask.to_bits();
+    let ip = ip.to_bits();
+    let broadcast = ip | neg_netmask;
+    Ipv4Addr::from_bits(broadcast)
 }
 
 fn read_sa_ip(ifa_addr: *const SockAddrIn) -> Ipv4Addr {
@@ -299,7 +317,7 @@ mod tests {
     use std::{net::Ipv4Addr, str::FromStr};
 
     use super::find_ip_from_arp;
-    use crate::mac::Mac;
+    use crate::{mac::Mac, netif::NetworkAddress};
 
     #[test]
     fn arp_file_parse() {
@@ -328,6 +346,14 @@ mod tests {
         validate_broadcast_address_calculation("192.168.1.8", "255.255.240.0", "192.168.15.255");
     }
 
+    #[test]
+    fn calculate_network_address_correctly() {
+        validate_network_address_calculation("192.168.0.8", "255.255.255.0", "192.168.0.0");
+        validate_network_address_calculation("192.168.0.8", "255.255.0.0", "192.168.0.0");
+        validate_network_address_calculation("192.168.1.8", "255.255.0.0", "192.168.0.0");
+        validate_network_address_calculation("192.168.1.8", "255.255.240.0", "192.168.0.0");
+    }
+
     fn validate_mapping(mac: &str, expected_ip: &str) {
         let mac = Mac::create(mac).unwrap();
         let ip = find_ip_from_mac_against_test_data(&mac)
@@ -348,4 +374,19 @@ mod tests {
         let actual_broadcast = super::calculate_broadcast_address(ip, netmask);
         assert_eq!(expected_broadcast, actual_broadcast);
     }
+
+    fn validate_network_address_calculation(ip: &str, mask: &str, expected_network: &str) {
+        let ip = Ipv4Addr::from_str(ip).unwrap();
+        let mask = Ipv4Addr::from_str(mask).unwrap();
+        let broadcast = super::calculate_broadcast_address(ip, mask);
+        let network = NetworkAddress {
+            source_ip_hint: ip,
+            broadcast,
+            mask,
+        };
+
+        let expected_network = Ipv4Addr::from_str(expected_network).unwrap();
+        let actual_network = super::calculate_network_address(&network);
+        assert_eq!(expected_network, actual_network);
+    }
 }

src/wolz.rs

@@ -1,13 +1,15 @@
 use std::{
     io::Write,
-    net::{TcpStream, UdpSocket},
+    net::{Ipv4Addr, TcpStream, UdpSocket},
+    thread,
+    time::Duration,
 };
 
 use crate::{
     config::Config,
     error::{WError, WResult},
     mac::Mac,
-    netif::{self, LocalNetwork},
+    netif::{self, LocalNetwork, get_local_networks},
 };
 
 pub struct Wolz {
@@ -60,10 +62,21 @@ impl Wolz {
     // TODO: Fallback to full network scan if /proc/net/arp file lookup fails
     fn shut_down(&self, mac_or_nick: &str) -> WResult<()> {
         let mac = self.config.get_mac(mac_or_nick)?;
-        let ip = netif::find_ip_with_mac(mac)?;
+        let ip = netif::find_ip_with_mac(&mac);
+        if ip.is_err() {
+            println!("Mac {} IP was not cached, will scan the network", &mac);
+            scan_all_local_networks()?;
+        }
 
+        let ip = netif::find_ip_with_mac(&mac)?;
         println!("Sending shutdown signal to {}:21367", ip);
         let mut stream = TcpStream::connect(format!("{}:21367", ip)).map_err(WError::TcpSocketBindError)?;
+        stream
+            .set_write_timeout(Some(Duration::from_secs(5)))
+            .map_err(WError::TcpSocketWriteError)?;
+        stream
+            .set_read_timeout(Some(Duration::from_secs(5)))
+            .map_err(WError::TcpSocketWriteError)?;
         stream
             .write("shutdown".as_bytes())
             .map_err(WError::TcpSocketWriteError)?;
@@ -72,6 +85,48 @@ impl Wolz {
     }
 }
 
+fn scan_all_local_networks() -> WResult<()> {
+    let local_networks = get_local_networks()?;
+    for n in local_networks {
+        scan_single_network(&n);
+    }
+    let sleep_time = Duration::from_secs(3);
+    println!(
+        "Sleeping for {} seconds for arp cache to get updated",
+        sleep_time.as_secs()
+    );
+    thread::sleep(sleep_time);
+    Ok(())
+}
+
+fn scan_single_network(network: &LocalNetwork) {
+    let start = network.subnet_addr.to_bits() + 1;
+    let end = network.subnet_broadcast.to_bits();
+    let packet = [0u8];
+    let total_addresses = end - start;
+    if total_addresses > 500 {
+        println!(
+            "Network {}/{} is very big, refusing to scan",
+            network.subnet_addr, network.subnet_mask
+        );
+        return;
+    }
+
+    for i in start..end {
+        let s = UdpSocket::bind(format!("{}:0", network.source_ip_hint));
+        if let Ok(s) = s {
+            if let Ok(_) = s.send_to(&packet, format!("{}:80", Ipv4Addr::from_bits(i))) {}
+        }
+    }
+
+    println!(
+        "Finished scanning {} addresses from {} to {}",
+        total_addresses,
+        Ipv4Addr::from_bits(start),
+        Ipv4Addr::from_bits(end)
+    );
+}
+
 fn try_send_magic_bytes(magic_bytes: &[u8; 102], mac: &Mac, local_network: &LocalNetwork) {
     let send_addr = format!("{}:0", local_network.source_ip_hint);
     let broadcast_addr = format!("{}:9", local_network.subnet_broadcast);