From 2644b17818824b64f74e95110fa5b75823810ac1 Mon Sep 17 00:00:00 2001 From: Nils Andreas Svee Date: Sat, 27 Aug 2022 16:27:15 +0200 Subject: [PATCH] Initial commit --- .cargo/config | 10 +++ .gitignore | 1 + Cargo.lock | 165 +++++++++++++++++++++++++++++++++++++ Cargo.toml | 22 +++++ src/cake.rs | 24 ++++++ src/config.rs | 17 ++++ src/main.rs | 79 ++++++++++++++++++ src/netlink.rs | 187 ++++++++++++++++++++++++++++++++++++++++++ src/pinger.rs | 215 +++++++++++++++++++++++++++++++++++++++++++++++++ 9 files changed, 720 insertions(+) create mode 100644 .cargo/config create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 src/cake.rs create mode 100644 src/config.rs create mode 100644 src/main.rs create mode 100644 src/netlink.rs create mode 100644 src/pinger.rs diff --git a/.cargo/config b/.cargo/config new file mode 100644 index 0000000..67ecc3f --- /dev/null +++ b/.cargo/config @@ -0,0 +1,10 @@ +[target.mips64-unknown-linux-muslabi64] +linker = "mips64-openwrt-linux-gcc" +rustflags = ["-C", "target-feature=-crt-static"] + +[target.mipsel-unknown-linux-musl] +linker = "mipsel-openwrt-linux-gcc" +rustflags = ["-C", "target-feature=-crt-static"] + +[target.x86_64-unknown-linux-gnu] +rustflags = ["-C", "target-feature=-crt-static"] \ No newline at end of file diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..799914f --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,165 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "autocfg" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa" + +[[package]] +name = "bitflags" +version = "1.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bef38d45163c2f1dde094a7dfd33ccf595c92905c8f8f4fdc18d06fb1037718a" + +[[package]] +name = "byteorder" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" + +[[package]] +name = "cake-autorate-rust" +version = "0.1.0" +dependencies = [ + "byteorder", + "libc", + "neli", + "nix", +] + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "either" +version = "1.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f107b87b6afc2a64fd13cac55fe06d6c8859f12d4b14cbcdd2c67d0976781be" + +[[package]] +name = "libc" +version = "0.2.132" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8371e4e5341c3a96db127eb2465ac681ced4c433e01dd0e938adbef26ba93ba5" + +[[package]] +name = "log" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "abb12e687cfb44aa40f41fc3978ef76448f9b6038cad6aef4259d3c095a2382e" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "memoffset" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aa361d4faea93603064a027415f07bd8e1d5c88c9fbf68bf56a285428fd79ce" +dependencies = [ + "autocfg", +] + +[[package]] +name = "neli" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f04de495602387a2c53a1eddd39f1b8aed297af3ed3b003910654f1f35ec4660" +dependencies = [ + "byteorder", + "libc", + "log", + "neli-proc-macros", +] + +[[package]] +name = "neli-proc-macros" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c168194d373b1e134786274020dae7fc5513d565ea2ebb9bc9ff17ffb69106d4" +dependencies = [ + "either", + "proc-macro2", + "quote", + "serde", + "syn", +] + +[[package]] +name = "nix" +version = "0.25.0" +dependencies = [ + "autocfg", + "bitflags", + "cfg-if", + "libc", + "memoffset", + "pin-utils", +] + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "proc-macro2" +version = "1.0.43" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a2ca2c61bc9f3d74d2886294ab7b9853abd9c1ad903a3ac7815c58989bb7bab" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbe448f377a7d6961e30f5955f9b8d106c3f5e449d493ee1b125c1d43c2b5179" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "serde" +version = "1.0.143" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e8e5d5b70924f74ff5c6d64d9a5acd91422117c60f48c4e07855238a254553" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.143" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3d8e8de557aee63c26b85b947f5e59b690d0454c753f3adeb5cd7835ab88391" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "syn" +version = "1.0.99" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58dbef6ec655055e20b86b15a8cc6d439cca19b667537ac6a1369572d151ab13" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "unicode-ident" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4f5b37a154999a8f3f98cc23a628d850e154479cd94decf3414696e12e31aaf" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..b27ec7a --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "cake-autorate-rust" +version = "0.1.0" +edition = "2021" + +# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html + +[profile.release] +codegen-units = 1 +lto = true +opt-level = 'z' +strip = "symbols" + + +[dependencies] +byteorder = "1.4.3" +libc = "0.2.132" +neli = "0.6.2" +nix = { version = "0.25.0", features = ["net", "socket", "time"] } + +[patch.crates-io] +nix = { path = "../nix" } \ No newline at end of file diff --git a/src/cake.rs b/src/cake.rs new file mode 100644 index 0000000..b20e4ff --- /dev/null +++ b/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, +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..e4e850f --- /dev/null +++ b/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"), + } + } +} diff --git a/src/main.rs b/src/main.rs new file mode 100644 index 0000000..6b3d1a3 --- /dev/null +++ b/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 = 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() + ); +} diff --git a/src/netlink.rs b/src/netlink.rs new file mode 100644 index 0000000..b69d552 --- /dev/null +++ b/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> { + 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 = 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> { + 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 = 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> { + 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(()) +} diff --git a/src/pinger.rs b/src/pinger.rs new file mode 100644 index 0000000..77ce175 --- /dev/null +++ b/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::().unwrap(), + identifier: (&buf[4..5]).read_u16::().unwrap(), + sequence: (&buf[6..7]).read_u16::().unwrap(), + payload: (&buf[8..16]).read_u64::().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::().unwrap(), + identifier: (&buf[4..5]).read_u16::().unwrap(), + sequence: (&buf[6..7]).read_u16::().unwrap(), + payload: (&buf[8..16]).read_u64::().unwrap(), + } + } +} + +pub trait Pinger { + fn new(reflectors: Vec) -> Self; + fn receive_loop(&self); + fn receive_ping(&self) -> Result<(), Box>; + fn sender_loop(&self); + fn send_ping(&self, reflector: &IpAddr, seq: u16) -> Result<(), Box>; +} + +#[derive(Clone)] +pub struct PingerICMPEcho { + reflectors: Vec, + 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) -> 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 = vec![0; size_of::()]; + let buf = v.as_mut_slice(); + + println!("buf len: {}", buf.len()); + + let res = recvfrom::(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> { + // 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> { + 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; + + 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(()) + } +}