tyr/
util.rs

1use std::{net::Ipv4Addr, os::unix::process::ExitStatusExt};
2
3use pnet::datalink;
4
5use crate::error;
6
7/// Generates a random port number within the range of 1024 to 65535.
8///
9/// # Returns
10///
11/// A `u16` representing a random port number.
12pub fn get_random_port() -> u16 {
13    fastrand::u16(1024..65535)
14}
15
16/// Generates a random IPv4 address.
17///
18/// # Returns
19///
20/// An `Ipv4Addr` representing a random IPv4 address.
21pub fn get_random_ip() -> Ipv4Addr {
22    Ipv4Addr::new(
23        fastrand::u8(0..255),
24        fastrand::u8(0..255),
25        fastrand::u8(0..255),
26        fastrand::u8(0..255),
27    )
28}
29
30/// Retrieves a network interface by its name.
31///
32/// # Arguments
33///
34/// * `name` - A string slice that holds the name of the network interface.
35///
36/// # Returns
37///
38/// * `Option<datalink::NetworkInterface>` - An option containing the network interface if found, or None if not found.
39pub fn get_interface(name: &str) -> Option<datalink::NetworkInterface> {
40    datalink::interfaces()
41        .into_iter()
42        .find(|interface| interface.name == name)
43}
44
45/// Retrieves all network interfaces available on the system.
46///
47/// # Returns
48///
49/// * `Vec<datalink::NetworkInterface>` - A vector containing all network interfaces.
50pub fn get_interfaces() -> Vec<datalink::NetworkInterface> {
51    datalink::interfaces()
52}
53
54/// Retrieves the network interface with the default gateway.
55///
56/// # Returns
57///
58/// * `Option<datalink::NetworkInterface>` - An option containing the network interface with the default gateway if found, or None if not found.
59pub fn get_default_interface() -> Option<datalink::NetworkInterface> {
60    datalink::interfaces().into_iter().find(|interface| {
61        interface
62            .ips
63            .iter()
64            .any(|ip| ip.is_ipv4() && ip.ip().is_loopback())
65    })
66}
67
68#[inline(always)]
69/// Checks if the current user is the root user.
70///
71/// Returns `true` if the current user is the root user, `false` otherwise.
72fn is_root() -> bool {
73    nix::unistd::Uid::current().is_root()
74}
75
76/// If the current user is not the root user, reruns the current program using sudo.
77///
78/// Returns `Ok(())` if the program was successfully rerun with sudo, or an `Err` if the rerun failed.
79pub fn rerun_if_not_root() -> crate::Result<()> {
80    if !is_root() {
81        let args: Vec<String> = std::env::args().collect();
82        let program = &args[0];
83        let rest_args = &args[1..];
84
85        let status = std::process::Command::new("sudo")
86            .arg(program)
87            .args(rest_args)
88            .status()?;
89
90        if !status.success() {
91            match status.signal() {
92                Some(nix::libc::SIGINT) => {
93                    std::process::exit(130);
94                }
95                Some(signal) => {
96                    return Err(error::Error::OsError(std::io::Error::new(
97                        std::io::ErrorKind::PermissionDenied,
98                        format!("Failed to grant root privileges: signal {}", signal),
99                    )));
100                }
101                _ => {
102                    return Err(error::Error::OsError(std::io::Error::new(
103                        std::io::ErrorKind::PermissionDenied,
104                        "Failed to grant root privileges",
105                    )));
106                }
107            }
108        }
109    }
110    Ok(())
111}
112
113/// Calculates the number of threads to use based on the rated power.
114///
115/// The rated power is a value from 1 to 4, which mean use a fraction
116/// of the available threads.
117///
118/// # Arguments
119///
120/// * `rated_power` - The rated power value.
121///
122/// # Returns
123///
124/// The number of threads to use.
125///
126/// # Errors
127///
128/// Returns an error if the system fails to determine the available parallelism.
129///
130/// # Panics
131///
132/// Panics if the rated power is not in the range 1 to 4.
133///
134/// # Warning
135///
136/// 0 is also a legal power rating parameter. Using the zero rated power will
137/// put the unit into a full load condition, which will definitely affect the
138/// normal operation of the unit and may even cause unknown consequences.
139///
140/// **USE AT YOUR OWN RISK!**
141pub fn get_num_threads(rated_power: u8) -> crate::Result<usize> {
142    let available_parallelism = std::thread::available_parallelism()?.get();
143    Ok(match rated_power {
144        0 => available_parallelism,
145        1 => available_parallelism / 10,
146        2 => available_parallelism / 5,
147        3 => available_parallelism / 2,
148        4 => available_parallelism * 3 / 4,
149        _ => panic!(
150            "Invalid rated power {}, available ratings are 1 to 4",
151            rated_power
152        ),
153    })
154}
tyr/util.rs

tyr/
util.rs
