reliable shutdown

This commit is contained in:
cool-mist
2026-07-12 09:38:14 +05:30
parent a8cdd6d448
commit 8be77f8cc7
2 changed files with 115 additions and 19 deletions
+57 -16
View File
@@ -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 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<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_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);
}
}
+58 -3
View File
@@ -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);