RedBear-OS/netstack/src/link/qdisc.rs

//! Traffic Control (qdisc) — mirrors Linux 7.1's `net/sched/`.
//!
//! Reference files:
//! - `net/sched/sch_tbf.c` — Token Bucket Filter (`tbf_enqueue`, `tbf_dequeue`)
//! - `net/sched/sch_prio.c` — Priority queue (`prio_enqueue`, `prio_dequeue`)
//! - `net/sched/sch_api.c` — Qdisc registration and API
//!
//! Two qdiscs are implemented:
//! - **TokenBucket**: rate-limits outgoing packets using a token bucket.
//!   Packets exceeding the rate are dropped. Mirrors `TBF`.
//! - **PriorityQueue**: three FIFO bands prioritized by IP TOS field.
//!   Mirrors `pfifo_fast` (the default Linux qdisc).

use std::collections::VecDeque;

use smoltcp::time::{Duration, Instant};

#[derive(Debug, Clone)]
pub struct TokenBucket {
    rate: u64,
    burst: u64,
    tokens: u64,
    last_update: Instant,
}

impl TokenBucket {
    pub fn new(rate_bps: u64, burst_bytes: u64) -> Self {
        Self { rate: rate_bps, burst: burst_bytes.max(1500), tokens: burst_bytes.max(1500), last_update: Instant::from_millis(0) }
    }
    pub fn rate(&self) -> u64 { self.rate }
    pub fn burst(&self) -> u64 { self.burst }
    pub fn tokens(&self) -> u64 { self.tokens }

    pub fn set_rate(&mut self, rate_bps: u64) {
        self.rate = rate_bps;
    }

    pub fn set_burst(&mut self, burst_bytes: u64) {
        self.burst = burst_bytes.max(1500);
        if self.tokens > self.burst {
            self.tokens = self.burst;
        }
    }

    pub fn consume(&mut self, bytes: u64, now: Instant) -> bool {
        let elapsed = now - self.last_update;
        if elapsed > Duration::ZERO {
            // Saturating mul: rate (bytes/s) * elapsed (ms) / 8000.
            // saturating_mul prevents overflow on long elapsed durations.
            let token_add = self.rate
                .saturating_mul(elapsed.total_millis() as u64)
                / 8000;
            self.tokens = (self.tokens + token_add).min(self.burst);
            self.last_update = now;
        }
        if self.tokens >= bytes {
            self.tokens -= bytes;
            true
        } else {
            false
        }
    }
}

#[derive(Debug, Clone)]
pub struct PriorityQueue {
    bands: [VecDeque<Vec<u8>>; 3],
    max_len: usize,
}

impl PriorityQueue {
    pub fn new(max_len: usize) -> Self { Self { bands: [VecDeque::new(), VecDeque::new(), VecDeque::new()], max_len } }
    pub fn max_len(&self) -> usize { self.max_len }

    fn classify(packet: &[u8]) -> usize {
        if packet.len() < 2 || packet[0] >> 4 != 4 {
            return 1;
        }
        let tos = if packet.len() > 1 { packet[1] } else { 0 };
        let precedence = tos >> 5;
        if precedence >= 6 {
            0
        } else if precedence >= 4 {
            1
        } else {
            2
        }
    }

    pub fn enqueue(&mut self, packet: Vec<u8>) -> bool {
        let band = Self::classify(&packet);
        if self.bands[band].len() >= self.max_len {
            return false;
        }
        self.bands[band].push_back(packet);
        true
    }

    pub fn dequeue(&mut self) -> Option<Vec<u8>> {
        for band in 0..3 {
            if let Some(packet) = self.bands[band].pop_front() {
                return Some(packet);
            }
        }
        None
    }

    pub fn len(&self) -> usize {
        self.bands[0].len() + self.bands[1].len() + self.bands[2].len()
    }

    pub fn flush(&mut self) -> Vec<Vec<u8>> {
        let mut result = Vec::new();
        while let Some(pkt) = self.dequeue() {
            result.push(pkt);
        }
        result
    }
}

#[derive(Debug, Clone)]
pub enum QdiscConfig {
    None,
    TokenBucket(TokenBucket),
    PriorityQueue(PriorityQueue),
}

impl Default for QdiscConfig {
    fn default() -> Self {
        Self::None
    }
}

impl QdiscConfig {
    pub fn enqueue_and_dequeue(
        &mut self,
        packet: Vec<u8>,
        now: Instant,
    ) -> Option<Vec<u8>> {
        match self {
            QdiscConfig::None => Some(packet),
            QdiscConfig::TokenBucket(tb) => {
                if tb.consume(packet.len() as u64, now) {
                    Some(packet)
                } else {
                    None
                }
            }
            QdiscConfig::PriorityQueue(pq) => {
                pq.enqueue(packet);
                pq.dequeue()
            }
        }
    }
}