Lochnair-Labs/sqm-autorate-rust
2
0 Fork 0
Code Issues 5 Pull Requests Actions 14 Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
Nils Andreas Svee committed 2022-08-27 16:27:15 +02:00
commit 2644b17818
9 files changed
+720
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.cargo/config
+10
View File
@@ -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"]
.gitignore
+1
View File
@@ -0,0 +1 @@
/target
Cargo.lock
Generated
+165
View File
@@ -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"
Cargo.toml
+22
View File
@@ -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" }
src/cake.rs
+24
View File
@@ -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
+17
View File
@@ -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
+79
View File
@@ -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
+187
View File
@@ -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
+215
View File
@@ -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(())
}
}