Initial commit

src/cake.rs

@@ -0,0 +1,24 @@
+#[allow(dead_code)]
+pub enum TcaCake {
+    Unspec = 0,
+    Pad = 1,
+    BaseRate64 = 2,
+    DiffservMode = 3,
+    ATM = 4,
+    FlowMode = 5,
+    Overhead = 6,
+    RTT = 7,
+    Target = 8,
+    Autorate = 9,
+    Memory = 10,
+    NAT = 11,
+    Raw = 12,
+    Wash = 13,
+    MPU = 14,
+    Ingress = 15,
+    AckFilter = 16,
+    SplitGso = 17,
+    FwMark = 18,
+    FwMarkStore = 19,
+    SCE = 20,
+}

src/config.rs

@@ -0,0 +1,17 @@
+use std::env;
+
+#[allow(dead_code)]
+pub(crate) struct Config {
+    pub(crate) download_interface: String,
+    pub(crate) upload_interface: String,
+}
+
+impl Config {
+    pub fn new() -> Self {
+        Self {
+            download_interface: env::var("DOWNLOAD_INTERFACE")
+                .expect("A download interface is required"),
+            upload_interface: env::var("UPLOAD_INTERFACE").expect("A upload interface is required"),
+        }
+    }
+}

src/main.rs

@@ -0,0 +1,79 @@
+extern crate core;
+
+mod cake;
+mod config;
+mod netlink;
+mod pinger;
+
+use std::net::IpAddr;
+use std::str::FromStr;
+use std::{thread, time};
+
+use crate::config::Config;
+use crate::netlink::{find_interface, find_qdisc, set_qdisc_rate};
+use crate::pinger::{Pinger, PingerICMPEcho};
+
+#[macro_export]
+macro_rules! error {
+    ($a:expr,$b:expr) => {
+        match $a {
+            //Some(e) => return e,
+            Ok(e) => e,
+            Err(e) => {
+                println!("{}: {}", $b, $a);
+            }
+        }
+    };
+}
+
+fn main() {
+    println!("starting up");
+
+    let config = Config::new();
+    test(config);
+
+    let mut reflectors: Vec<IpAddr> = Vec::new();
+
+    reflectors.push(IpAddr::from_str("1.0.0.1").unwrap());
+    reflectors.push(IpAddr::from_str("1.1.1.1").unwrap());
+    reflectors.push(IpAddr::from_str("9.9.9.9").unwrap());
+    reflectors.push(IpAddr::from_str("9.9.9.10").unwrap());
+
+    let pinger = PingerICMPEcho::new(reflectors);
+    let pinger_receiver = pinger.clone();
+    let pinger_sender = pinger.clone();
+
+    let receiver_handle = thread::spawn(move || pinger_receiver.receive_loop());
+    let sender_handle = thread::spawn(move || pinger_sender.sender_loop());
+
+    receiver_handle.join().unwrap();
+    sender_handle.join().unwrap();
+}
+
+fn test(config: Config) {
+    if true {
+        return;
+    }
+
+    let start = time::Instant::now();
+    let ifindex =
+        find_interface(config.upload_interface.as_str()).expect("Couldn't find interface");
+    let find_if_time = time::Instant::now().duration_since(start);
+    println!("time get intf: {}", find_if_time.as_micros());
+
+    let start = time::Instant::now();
+    let qdisc = find_qdisc(ifindex).expect("Couldn't find qdisc");
+    let find_qdisc_time = time::Instant::now().duration_since(start);
+
+    println!("time get qdisc: {}", find_qdisc_time.as_micros());
+
+    let start = time::Instant::now();
+    set_qdisc_rate(qdisc, 1000 * 1000 / 8).unwrap();
+    let set_qdisc_time = time::Instant::now().duration_since(start);
+    println!("time set qdisc: {}", set_qdisc_time.as_micros());
+
+    println!(
+        "time total: {}",
+        find_if_time.as_micros() + find_qdisc_time.as_micros() + set_qdisc_time.as_micros()
+    );
+}

src/netlink.rs

@@ -0,0 +1,187 @@
+use crate::cake::TcaCake;
+
+use neli::consts::nl::{NlmF, NlmFFlags};
+use neli::consts::rtnl::{Arphrd, IffFlags, Ifla, RtAddrFamily, Rtm, Tca};
+use neli::consts::socket::NlFamily;
+use neli::err::NlError;
+use neli::nl::{NlPayload, Nlmsghdr};
+use neli::rtnl::{Ifinfomsg, Rtattr, Tcmsg};
+use neli::socket::NlSocketHandle;
+use neli::types::{Buffer, RtBuffer};
+use std::error::Error;
+use std::fmt::Display;
+
+use std::fmt;
+
+#[derive(Default, Debug)]
+struct NoQdiscFoundError;
+
+impl Error for NoQdiscFoundError {
+    fn source(&self) -> Option<&(dyn Error + 'static)> {
+        // The compiler transparently casts `&sqlx::Error` into a `&dyn Error`
+        None
+    }
+}
+
+// Implement std::fmt::Display for AppError
+impl Display for NoQdiscFoundError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "Couldn't find a CAKE qdisc for interface") // user-facing output
+    }
+}
+
+#[allow(dead_code)]
+pub struct Qdisc {
+    ifindex: i32,
+    handle: u32,
+    family: u8,
+    refcnt: u32,
+    parent: u32,
+}
+
+pub fn find_interface(ifname: &str) -> Result<i32, Box<dyn Error>> {
+    let mut socket = NlSocketHandle::connect(NlFamily::Route, None, &[]).unwrap();
+    let mut attrs = RtBuffer::new();
+
+    const RTEXT_FILTER_VF: i32 = 1 << 0;
+
+    let attr_ifname = Rtattr::new(None, Ifla::Ifname, ifname).unwrap();
+    let attr_ext_mask = Rtattr::new(None, Ifla::ExtMask, RTEXT_FILTER_VF).unwrap();
+    attrs.push(attr_ifname);
+    attrs.push(attr_ext_mask);
+
+    let if_msg = Ifinfomsg::new(
+        RtAddrFamily::Unspecified,
+        Arphrd::None,
+        -1,
+        IffFlags::empty(),
+        IffFlags::empty(),
+        attrs,
+    );
+
+    let nlhdr = Nlmsghdr::new(
+        None,
+        Rtm::Getlink,
+        NlmFFlags::new(&[NlmF::Request, NlmF::Ack]),
+        None,
+        None,
+        NlPayload::Payload(if_msg),
+    );
+
+    socket.send(nlhdr)?;
+
+    for response in socket.iter(false) {
+        let header: Nlmsghdr<Rtm, Ifinfomsg> = response?;
+
+        if header.nl_type != Rtm::Newlink {
+            return Err(Box::new(NlError::msg("Netlink error retrieving link")));
+        }
+
+        if let NlPayload::Payload(p) = header.nl_payload {
+            return Ok(p.ifi_index);
+        }
+    }
+
+    // we shouldn't reach here
+    Ok(-1)
+}
+
+pub fn find_qdisc(ifindex: i32) -> Result<Qdisc, Box<dyn Error>> {
+    let mut socket = NlSocketHandle::connect(NlFamily::Route, None, &[]).unwrap();
+    let tc_msg = Tcmsg::new(
+        u8::from(RtAddrFamily::Unspecified),
+        0,
+        0,
+        0,
+        0,
+        RtBuffer::new(),
+    );
+
+    let nlhdr = Nlmsghdr::new(
+        None,
+        Rtm::Getqdisc,
+        NlmFFlags::new(&[NlmF::Request, NlmF::Dump]),
+        None,
+        None,
+        NlPayload::Payload(tc_msg),
+    );
+
+    socket.send(nlhdr)?;
+
+    for response in socket.iter(false) {
+        let header: Nlmsghdr<Rtm, Tcmsg> = response?;
+
+        if let NlPayload::Payload(p) = header.nl_payload {
+            if header.nl_type != Rtm::Newqdisc {
+                return Err(Box::new(NlError::msg("Netlink error retrieving qdisc")));
+            }
+
+            if p.tcm_ifindex == ifindex {
+                let mut _type = "";
+
+                for attr in p.rtattrs.iter() {
+                    if attr.rta_type == Tca::Kind {
+                        let buff = attr.rta_payload.as_ref();
+                        _type = std::str::from_utf8(buff)
+                            .expect("Found invalid UTF-8 string")
+                            .trim_end_matches('\0'); // Null terminator is valid UTF-8, but breaks comparison, so we remove it
+                    }
+                }
+
+                let type_to_look_for = "cake";
+
+                if _type.eq(type_to_look_for) {
+                    let qdisc = Qdisc {
+                        ifindex: p.tcm_ifindex as i32,
+                        handle: p.tcm_handle,
+                        family: p.tcm_family,
+                        refcnt: p.tcm_info,
+                        parent: p.tcm_parent,
+                    };
+
+                    return Ok(qdisc);
+                }
+            }
+        }
+    }
+
+    return Err(Box::new(NoQdiscFoundError));
+}
+
+pub fn set_qdisc_rate(qdisc: Qdisc, bandwidth: u64) -> Result<(), Box<dyn Error>> {
+    let mut socket = NlSocketHandle::connect(NlFamily::Route, None, &[]).unwrap();
+
+    let mut attrs = RtBuffer::new();
+
+    let attr_type = Rtattr::new(None, Tca::Kind, "cake").unwrap();
+    let mut attr_options = Rtattr::new(None, Tca::Options, Buffer::from(Vec::new())).unwrap();
+    attr_options.add_nested_attribute(&Rtattr::new(
+        None,
+        TcaCake::BaseRate64 as u16,
+        bandwidth,
+    )?)?;
+
+    attrs.push(attr_type);
+    attrs.push(attr_options);
+
+    let tc_msg = Tcmsg::new(
+        u8::from(RtAddrFamily::Unspecified),
+        qdisc.ifindex,
+        0,
+        qdisc.parent,
+        0,
+        attrs,
+    );
+
+    let nlhdr = Nlmsghdr::new(
+        None,
+        Rtm::Newqdisc,
+        NlmFFlags::new(&[NlmF::Request, NlmF::Ack]),
+        None,
+        None,
+        NlPayload::Payload(tc_msg),
+    );
+
+    socket.send(nlhdr)?;
+    Ok(())
+}

src/pinger.rs

@@ -0,0 +1,215 @@
+use byteorder::{BigEndian, ReadBytesExt};
+use nix::sys::socket::{
+    recvfrom, sendto, AddressFamily, MsgFlags, SockFlag, SockProtocol, SockType, SockaddrIn,
+    SockaddrIn6, SockaddrLike,
+};
+use nix::sys::time::TimeValLike;
+use nix::time::{clock_gettime, ClockId};
+use std::error::Error;
+use std::mem::size_of;
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6};
+use std::os::unix::io::RawFd;
+use std::str::FromStr;
+use std::{thread, time};
+
+struct ICMP {
+    icmp_type: u8,
+    code: u8,
+    checksum: u16,
+    identifier: u16,
+    sequence: u16,
+    payload: u64,
+}
+
+impl ICMP {
+    fn to_bytes(&self) -> Box<[u8; 16]> {
+        let checksum = self.checksum.to_be_bytes();
+        let identifier = self.identifier.to_be_bytes();
+        let sequence = self.sequence.to_be_bytes();
+        let payload = self.payload.to_be_bytes();
+
+        let buf = [
+            self.icmp_type,
+            self.code,
+            checksum[0],
+            checksum[1],
+            identifier[0],
+            identifier[1],
+            sequence[0],
+            sequence[1],
+            payload[0],
+            payload[1],
+            payload[2],
+            payload[3],
+            payload[4],
+            payload[5],
+            payload[6],
+            payload[7],
+        ];
+
+        return Box::new(buf);
+    }
+}
+
+impl From<&[u8]> for ICMP {
+    fn from(buf: &[u8]) -> Self {
+        ICMP {
+            icmp_type: buf[0],
+            code: buf[1],
+            checksum: (&buf[2..3]).read_u16::<BigEndian>().unwrap(),
+            identifier: (&buf[4..5]).read_u16::<BigEndian>().unwrap(),
+            sequence: (&buf[6..7]).read_u16::<BigEndian>().unwrap(),
+            payload: (&buf[8..16]).read_u64::<BigEndian>().unwrap(),
+        }
+    }
+}
+
+impl From<&mut [u8]> for ICMP {
+    fn from(buf: &mut [u8]) -> Self {
+        ICMP {
+            icmp_type: buf[0],
+            code: buf[1],
+            checksum: (&buf[2..3]).read_u16::<BigEndian>().unwrap(),
+            identifier: (&buf[4..5]).read_u16::<BigEndian>().unwrap(),
+            sequence: (&buf[6..7]).read_u16::<BigEndian>().unwrap(),
+            payload: (&buf[8..16]).read_u64::<BigEndian>().unwrap(),
+        }
+    }
+}
+
+pub trait Pinger {
+    fn new(reflectors: Vec<IpAddr>) -> Self;
+    fn receive_loop(&self);
+    fn receive_ping(&self) -> Result<(), Box<dyn Error>>;
+    fn sender_loop(&self);
+    fn send_ping(&self, reflector: &IpAddr, seq: u16) -> Result<(), Box<dyn Error>>;
+}
+
+#[derive(Clone)]
+pub struct PingerICMPEcho {
+    reflectors: Vec<IpAddr>,
+    socket: RawFd,
+}
+
+impl PingerICMPEcho {
+    fn calculate_checksum(buffer: &mut [u8]) -> u16 {
+        let mut sum = 0u32;
+        for word in buffer.chunks(2) {
+            let mut part = u16::from(word[0]) << 8;
+            if word.len() > 1 {
+                part += u16::from(word[1]);
+            }
+            sum = sum.wrapping_add(u32::from(part));
+        }
+
+        while (sum >> 16) > 0 {
+            sum = (sum & 0xffff) + (sum >> 16);
+        }
+
+        !sum as u16
+    }
+}
+
+impl Pinger for PingerICMPEcho {
+    fn new(reflectors: Vec<IpAddr>) -> Self {
+        let socket = nix::sys::socket::socket(
+            AddressFamily::Inet,
+            SockType::Raw,
+            SockFlag::empty(),
+            SockProtocol::ICMP,
+        )
+        .expect("Couldn't open a socket");
+
+        PingerICMPEcho { reflectors, socket }
+    }
+
+    fn receive_loop(&self) {
+        loop {
+            let mut v: Vec<u8> = vec![0; size_of::<ICMP>()];
+            let buf = v.as_mut_slice();
+
+            println!("buf len: {}", buf.len());
+
+            let res = recvfrom::<SockaddrIn>(self.socket, buf).unwrap();
+            let remote = res.1.unwrap();
+
+            let ip = IpAddr::from(remote.ip().to_ne_bytes());
+
+            if self.reflectors.contains(&ip) {
+                println!("ip: {}", ip.to_string());
+                println!("usize: {}", res.0);
+                println!("{:?}", buf);
+
+                let hdr: ICMP = buf.into();
+                println!("Type: {}", hdr.icmp_type);
+            }
+        }
+    }
+
+    fn receive_ping(&self) -> Result<(), Box<dyn Error>> {
+        // Process ping RTT
+        /*let ping_rtt = SystemTime::now()
+        .duration_since(ping_start_time)
+        .unwrap_or(Duration::from_secs(0));*/
+        todo!()
+    }
+
+    fn sender_loop(&self) {
+        let mut seq: u16 = 0;
+        let tick_duration_ms: u16 = 500;
+        let sleep_duration = time::Duration::from_millis(tick_duration_ms as u64);
+
+        loop {
+            for reflector in self.reflectors.iter() {
+                self.send_ping(reflector, seq).expect("TODO: panic message");
+            }
+
+            thread::sleep(sleep_duration);
+
+            seq += 1;
+
+            if seq > std::u16::MAX {
+                seq = 0;
+            }
+        }
+    }
+
+    fn send_ping(&self, reflector: &IpAddr, seq: u16) -> Result<(), Box<dyn Error>> {
+        let time = clock_gettime(ClockId::CLOCK_MONOTONIC).unwrap();
+
+        let mut hdr = ICMP {
+            icmp_type: 8,
+            code: 0,
+            checksum: 0,
+            identifier: 0xBABE,
+            sequence: seq,
+            payload: time.num_milliseconds() as u64,
+        };
+
+        let mut bytes = hdr.to_bytes();
+        let buf: &mut [u8] = bytes.as_mut_slice();
+
+        hdr.checksum = PingerICMPEcho::calculate_checksum(buf);
+        let bytes = hdr.to_bytes();
+        let buf = bytes.as_slice();
+
+        let addr: Box<dyn SockaddrLike>;
+
+        match reflector.is_ipv4() {
+            true => {
+                let ip4 = Ipv4Addr::from_str(&*reflector.to_string()).unwrap();
+                let sock4 = SocketAddrV4::new(ip4, 0);
+                addr = Box::new(SockaddrIn::from(sock4));
+            }
+            false => {
+                let ip6 = Ipv6Addr::from_str(&*reflector.to_string()).unwrap();
+                let sock6 = SocketAddrV6::new(ip6, 0, 0, 0);
+                addr = Box::new(SockaddrIn6::from(sock6));
+            }
+        }
+
+        sendto(self.socket, buf, addr.as_ref(), MsgFlags::empty()).expect("TODO: F message");
+
+        Ok(())
+    }
+}