From 8be77f8cc7760c1d42bb43eaa9d2874a6fb38c38 Mon Sep 17 00:00:00 2001 From: cool-mist Date: Sun, 12 Jul 2026 09:38:14 +0530 Subject: [PATCH] reliable shutdown --- src/netif.rs | 73 ++++++++++++++++++++++++++++++++++++++++------------ src/wolz.rs | 61 ++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 115 insertions(+), 19 deletions(-) diff --git a/src/netif.rs b/src/netif.rs index 402c5eb..1085576 100644 --- a/src/netif.rs +++ b/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> { @@ -33,6 +35,8 @@ pub fn get_local_networks() -> WResult> { 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> { Ok(networks) } -pub fn find_ip_with_mac(mac: Mac) -> WResult { +pub fn find_ip_with_mac(mac: &Mac) -> WResult { 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> { 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 broadcast = calculate_broadcast_address(ip, netmask); + let mask = read_sa_ip(sa_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> { // 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_network_address(n_addr: &NetworkAddress) -> Ipv4Addr { + let network_address = n_addr.source_ip_hint.to_bits() & n_addr.mask.to_bits(); + Ipv4Addr::from_bits(network_address) +} + fn calculate_broadcast_address(ip: Ipv4Addr, netmask: Ipv4Addr) -> Ipv4Addr { - let ip_octets = ip.octets(); - let netmask_octets = netmask.octets(); - - 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) + 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); + } } diff --git a/src/wolz.rs b/src/wolz.rs index 8094a88..6feeec4 100644 --- a/src/wolz.rs +++ b/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);