//! Functions for applying the correct weights to relays when choosing

//! a relay at random.

//!

//! The weight to use when picking a relay depends on several factors:

//!

//! - The relay's *apparent bandwidth*.  (This is ideally measured by a set of

//!   bandwidth authorities, but if no bandwidth authorities are running (as on

//!   a test network), we might fall back either to relays' self-declared

//!   values, or we might treat all relays as having equal bandwidth.)

//! - The role that we're selecting a relay to play.  (See [`WeightRole`]).

//! - The flags that a relay has in the consensus, and their scarcity.  If a

//!   relay provides particularly scarce functionality, we might choose not to

//!   use it for other roles, or to use it less commonly for them.

use crate::ConsensusRelays;

use crate::params::NetParameters;

use bitflags::bitflags;

use tor_netdoc::doc::netstatus::{self, MdConsensus, MdRouterStatus, NetParams};

/// Helper: Calculate the function we should use to find initial relay

/// bandwidths.

{

        // If every value is zero, we should just pretend everything has

        // bandwidth == 1.

        // If there are no measured values, then we can look at unmeasured

        // weights.

        // Otherwise, there are measured values; we should look at those only, if

        // any of them is nonzero.

    } else {

        // This is a bit of an ugly case: We have measured values, but they're

        // all zero.  If this happens, the bandwidth authorities exist but they

        // very confused: we should fall back to uniform weighting.

    }

/// Internal: how should we find the base bandwidth of each relay?  This

/// value is global over a whole directory, and depends on the bandwidth

/// weights in the consensus.

#[derive(Copy, Clone, Debug, PartialEq, Eq)]

enum BandwidthFn {

    /// There are no weights at all in the consensus: weight every

    /// relay as 1.

    Uniform,

    /// There are no measured weights in the consensus: count

    /// unmeasured weights as the weights for relays.

    IncludeUnmeasured,

    /// There are measured relays in the consensus; only use those.

    MeasuredOnly,

}

impl BandwidthFn {

    /// Apply this function to the measured or unmeasured bandwidth

    /// of a single relay.

        use BandwidthFn::*;

        use netstatus::RelayWeight::*;

        }

}

/// Possible ways to weight relays when selecting them a random.

///

/// Relays are weighted by a function of their bandwidth that

/// depends on how scarce that "kind" of bandwidth is.  For

/// example, if Exit bandwidth is rare, then Exits should be

/// less likely to get chosen for the middle hop of a path.

#[derive(Clone, Debug, Copy)]

#[non_exhaustive]

pub enum WeightRole {

    /// Selecting a relay to use as a guard

    Guard,

    /// Selecting a relay to use as a middle relay in a circuit.

    Middle,

    /// Selecting a relay to use to deliver traffic to the internet.

    Exit,

    /// Selecting a relay for a one-hop BEGIN_DIR directory request.

    BeginDir,

    /// Selecting a relay with no additional weight beyond its bandwidth.

    Unweighted,

    /// Selecting a relay for use as a hidden service introduction point

    HsIntro,

    /// Selecting a relay for use as a hidden service rendezvous point

    HsRend,

}

/// Description for how to weight a single kind of relay for each WeightRole.

#[derive(Clone, Debug, Copy)]

struct RelayWeight {

    /// How to weight this kind of relay when picking a guard relay.

    as_guard: u32,

    /// How to weight this kind of relay when picking a middle relay.

    as_middle: u32,

    /// How to weight this kind of relay when picking a exit relay.

    as_exit: u32,

    /// How to weight this kind of relay when picking a one-hop BEGIN_DIR.

    as_dir: u32,

}

impl std::ops::Mul<u32> for RelayWeight {

    type Output = Self;

}

impl std::ops::Div<u32> for RelayWeight {

    type Output = Self;

}

impl RelayWeight {

    /// Return the largest weight that we give for this kind of relay.

    // The unwrap() is safe because array is nonempty.

    #[allow(clippy::unwrap_used)]

    /// Return the weight we should give this kind of relay's

    /// bandwidth for a given role.

        }

}

bitflags! {

    /// A kind of relay, for the purposes of selecting a relay by weight.

    ///

    /// Relays can have or lack the Guard flag, the Exit flag, and the

    /// V2Dir flag. All together, this makes 8 kinds of relays.

    #[derive(Clone, Copy, Debug, Eq, PartialEq)]

    struct WeightKind: u8 {

        /// Flag in weightkind for Guard relays.

        const GUARD = 1 << 0;

        /// Flag in weightkind for Exit relays.

        const EXIT = 1 << 1;

        /// Flag in weightkind for V2Dir relays.

        const DIR = 1 << 2;

    }

}

impl WeightKind {

    /// Return the appropriate WeightKind for a relay.

    /// Return the index to use for this kind of a relay within a WeightSet.

}

/// Information derived from a consensus to use when picking relays by

/// weighted bandwidth.

#[derive(Debug, Clone)]

pub(crate) struct WeightSet {

    /// How to find the bandwidth to use when picking a relay by weighted

    /// bandwidth.

    ///

    /// (This tells us us whether to count unmeasured relays, whether

    /// to look at bandwidths at all, etc.)

    bandwidth_fn: BandwidthFn,

    /// Number of bits that we need to right-shift our weighted products

    /// so that their sum won't overflow u64::MAX.

    //

    // TODO: Perhaps we should use f64 to hold our weights instead,

    // so we don't need to keep this ad-hoc fixed-point implementation?

    // If we did so, we won't have to worry about overflows.

    // (When we call choose_multiple_weighted, it already converts into

    // f64 internally.  (Though choose_weighted doesn't.))

    // Before making this change, however,

    // we should think a little about performance and precision.

    shift: u8,

    /// A set of RelayWeight values, indexed by [`WeightKind::idx`], used

    /// to weight different kinds of relays.

    w: [RelayWeight; 8],

}

impl WeightSet {

    /// Find the actual 64-bit weight to use for a given routerstatus when

    /// considering it for a given role.

    ///

    /// NOTE: This function _does not_ consider whether the relay in question

    /// actually matches the given role.  For example, if `role` is Guard

    /// we don't check whether or not `rs` actually has the Guard flag.

    /// Find the 64-bit weight to report for a relay of `kind` whose weight in

    /// the consensus is `relay_weight` when using it for `role`.

        // Note a subtlety here: we multiply the two values _before_

        // we shift, to improve accuracy.  We know that this will be

        // safe, since the inputs are both u32, and so cannot overflow

        // a u64.

    /// Compute the correct WeightSet for a provided MdConsensus.

    /// Compute the correct WeightSet given a bandwidth function, a

    /// weight-scaling parameter, a total amount of bandwidth for all

    /// relays in the consensus, and a set of bandwidth parameters.

        /// Find a single RelayWeight, given the names that its bandwidth

        /// parameters have. The `g` parameter is the weight as a guard, the

        /// `m` parameter is the weight as a middle relay, the `e` parameter is

        /// the weight as an exit, and the `d` parameter is the weight as a

        /// directory.

        #[allow(clippy::many_single_char_names)]

        // Prevent division by zero in case we're called with a bogus

        // input.  (That shouldn't be possible.)

        // For non-V2Dir relays, we have names for most of their weights.

        //

        // (There is no Wge, since we only use Guard relays as guards.  By the

        // same logic, Wme has no reason to exist, but according to the spec it

        // does.)

        // Note that the positions of the elements in this array need to

        // match the values returned by WeightKind.as_idx().

        // This is the largest weight value.

        // The unwrap() is safe because `w` is nonempty.

        #[allow(clippy::unwrap_used)]

        // We want "shift" such that (total * w_max) >> shift <= u64::max

    /// Assert that we have correctly computed our shift values so that

    /// our total weighted bws do not exceed u64::MAX.

        use WeightRole::*;

        }

}

/// The value to return if a weight parameter is absent.

///

/// (If there are no weights at all, then it's correct to set them all to 1,

/// and just use the bandwidths.  If _some_ are present and some are absent,

/// then the spec doesn't say what to do, but this behavior appears

/// reasonable.)

const DFLT_WEIGHT: i32 = 1;

/// Return the weight param named 'kwd' in p.

///

/// Returns DFLT_WEIGHT if there is no such parameter, and 0

/// if `kwd` is "---".

    } else {

    }

/// If `inp` is less than 0, return 0.  Otherwise return `inp` as a u32.

    // (The spec says that we might encounter negative values here, though

    // we never actually generate them, and don't plan to generate them.)

/// Compute a 'shift' value such that `(a * b) >> shift` will be contained

/// inside 64 bits.

/// Return an upper bound for the log2 of n.

///

/// This function overestimates whenever n is a power of two, but that doesn't

/// much matter for the uses we're giving it here.

#[cfg(test)]

mod test {

    // @@ begin test lint list maintained by maint/add_warning @@

    #![allow(clippy::bool_assert_comparison)]

    #![allow(clippy::clone_on_copy)]

    #![allow(clippy::dbg_macro)]

    #![allow(clippy::mixed_attributes_style)]

    #![allow(clippy::print_stderr)]

    #![allow(clippy::print_stdout)]

    #![allow(clippy::single_char_pattern)]

    #![allow(clippy::unwrap_used)]

    #![allow(clippy::unchecked_duration_subtraction)]

    #![allow(clippy::useless_vec)]

    #![allow(clippy::needless_pass_by_value)]

    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->

    use super::*;

    use netstatus::RelayWeight as RW;

    use std::net::SocketAddr;

    use std::time::{Duration, SystemTime};

    use tor_basic_utils::test_rng::testing_rng;

    use tor_netdoc::doc::netstatus::{Lifetime, MdRouterStatusBuilder, RelayFlags};

    #[test]

    fn t_clamp() {

        assert_eq!(clamp_to_pos(32), 32);

        assert_eq!(clamp_to_pos(i32::MAX), i32::MAX as u32);

        assert_eq!(clamp_to_pos(0), 0);

        assert_eq!(clamp_to_pos(-1), 0);

        assert_eq!(clamp_to_pos(i32::MIN), 0);

    }

    #[test]

    fn t_log2() {

        assert_eq!(log2_upper(u64::MAX), 64);

        assert_eq!(log2_upper(0), 0);

        assert_eq!(log2_upper(1), 1);

        assert_eq!(log2_upper(63), 6);

        assert_eq!(log2_upper(64), 7); // a little buggy but harmless.

    }

    #[test]

    fn t_calc_shift() {

        assert_eq!(calculate_shift(1 << 20, 1 << 20), 0);

        assert_eq!(calculate_shift(1 << 50, 1 << 10), 0);

        assert_eq!(calculate_shift(1 << 32, 1 << 33), 3);

        assert!(((1_u64 << 32) >> 3).checked_mul(1_u64 << 33).is_some());

        assert_eq!(calculate_shift(432 << 40, 7777 << 40), 38);

        assert!(

            ((432_u64 << 40) >> 38)

                .checked_mul(7777_u64 << 40)

                .is_some()

        );

    }

    #[test]

    fn t_pick_bwfunc() {

        let empty = [];

        assert_eq!(pick_bandwidth_fn(empty.iter()), BandwidthFn::Uniform);

        let all_zero = [RW::Unmeasured(0), RW::Measured(0), RW::Unmeasured(0)];

        assert_eq!(pick_bandwidth_fn(all_zero.iter()), BandwidthFn::Uniform);

        let all_unmeasured = [RW::Unmeasured(9), RW::Unmeasured(2222)];

        assert_eq!(

            pick_bandwidth_fn(all_unmeasured.iter()),

            BandwidthFn::IncludeUnmeasured

        );

        let some_measured = [

            RW::Unmeasured(10),

            RW::Measured(7),

            RW::Measured(4),

            RW::Unmeasured(0),

        ];

        assert_eq!(

            pick_bandwidth_fn(some_measured.iter()),

            BandwidthFn::MeasuredOnly

        );

        // This corresponds to an open question in

        // `pick_bandwidth_fn`, about what to do when the only nonzero

        // weights are unmeasured.

        let measured_all_zero = [RW::Unmeasured(10), RW::Measured(0)];

        assert_eq!(

            pick_bandwidth_fn(measured_all_zero.iter()),

            BandwidthFn::Uniform

        );

    }

    #[test]

    fn t_apply_bwfn() {

        use BandwidthFn::*;

        use netstatus::RelayWeight::*;

        assert_eq!(Uniform.apply(&Measured(7)), 1);

        assert_eq!(Uniform.apply(&Unmeasured(0)), 1);

        assert_eq!(IncludeUnmeasured.apply(&Measured(7)), 7);

        assert_eq!(IncludeUnmeasured.apply(&Unmeasured(8)), 8);

        assert_eq!(MeasuredOnly.apply(&Measured(9)), 9);

        assert_eq!(MeasuredOnly.apply(&Unmeasured(10)), 0);

    }

    // From a fairly recent Tor consensus.

    const TESTVEC_PARAMS: &str = "Wbd=0 Wbe=0 Wbg=4096 Wbm=10000 Wdb=10000 Web=10000 Wed=10000 Wee=10000 Weg=10000 Wem=10000 Wgb=10000 Wgd=0 Wgg=5904 Wgm=5904 Wmb=10000 Wmd=0 Wme=0 Wmg=4096 Wmm=10000";

    #[test]

    fn t_weightset_basic() {

        let total_bandwidth = 1_000_000_000;

        let params = TESTVEC_PARAMS.parse().unwrap();

        let ws = WeightSet::from_parts(BandwidthFn::MeasuredOnly, total_bandwidth, 10000, &params);

        assert_eq!(ws.bandwidth_fn, BandwidthFn::MeasuredOnly);

        assert_eq!(ws.shift, 0);

        assert_eq!(ws.w[0].as_guard, 5904);

        assert_eq!(ws.w[(WeightKind::GUARD.bits()) as usize].as_guard, 5904);

        assert_eq!(ws.w[(WeightKind::EXIT.bits()) as usize].as_exit, 10000);

        assert_eq!(

            ws.w[(WeightKind::EXIT | WeightKind::GUARD).bits() as usize].as_dir,

            0

        );

        assert_eq!(

            ws.w[(WeightKind::GUARD | WeightKind::DIR).bits() as usize].as_dir,

            4096

        );

        assert_eq!(

            ws.w[(WeightKind::GUARD | WeightKind::DIR).bits() as usize].as_dir,

            4096

        );

        assert_eq!(

            ws.weight_bw_for_role(

                WeightKind::GUARD | WeightKind::DIR,

                &RW::Unmeasured(7777),

                WeightRole::Guard

            ),

            0

        );

        assert_eq!(

            ws.weight_bw_for_role(

                WeightKind::GUARD | WeightKind::DIR,

                &RW::Measured(7777),

                WeightRole::Guard

            ),

            7777 * 5904

        );

        assert_eq!(

            ws.weight_bw_for_role(

                WeightKind::GUARD | WeightKind::DIR,

                &RW::Measured(7777),

                WeightRole::Middle

            ),

            7777 * 4096

        );

        assert_eq!(

            ws.weight_bw_for_role(

                WeightKind::GUARD | WeightKind::DIR,

                &RW::Measured(7777),

                WeightRole::Exit

            ),

            7777 * 10000

        );

        assert_eq!(

            ws.weight_bw_for_role(

                WeightKind::GUARD | WeightKind::DIR,

                &RW::Measured(7777),

                WeightRole::BeginDir

            ),

            7777 * 4096

        );

        assert_eq!(

            ws.weight_bw_for_role(

                WeightKind::GUARD | WeightKind::DIR,

                &RW::Measured(7777),

                WeightRole::Unweighted

            ),

            7777

        );

        // Now try those last few with routerstatuses.

        let rs = rs_builder()

            .set_flags(RelayFlags::GUARD | RelayFlags::V2DIR)

            .weight(RW::Measured(7777))

            .build()

            .unwrap();

        assert_eq!(ws.weight_rs_for_role(&rs, WeightRole::Exit), 7777 * 10000);

        assert_eq!(

            ws.weight_rs_for_role(&rs, WeightRole::BeginDir),

            7777 * 4096

        );

        assert_eq!(ws.weight_rs_for_role(&rs, WeightRole::Unweighted), 7777);

    }

    /// Return a routerstatus builder set up to deliver a routerstatus

    /// with most features disabled.

    fn rs_builder() -> MdRouterStatusBuilder {

        MdConsensus::builder()

            .rs()

            .identity([9; 20].into())

            .add_or_port(SocketAddr::from(([127, 0, 0, 1], 9001)))

            .doc_digest([9; 32])

            .protos("".parse().unwrap())

            .clone()

    }

    #[test]

    fn weight_flags() {

        let rs1 = rs_builder().set_flags(RelayFlags::EXIT).build().unwrap();

        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::EXIT);

        let rs1 = rs_builder().set_flags(RelayFlags::GUARD).build().unwrap();

        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::GUARD);

        let rs1 = rs_builder().set_flags(RelayFlags::V2DIR).build().unwrap();

        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::DIR);

        let rs1 = rs_builder().build().unwrap();

        assert_eq!(WeightKind::for_rs(&rs1), WeightKind::empty());

        let rs1 = rs_builder().set_flags(RelayFlags::all()).build().unwrap();

        assert_eq!(

            WeightKind::for_rs(&rs1),

            WeightKind::EXIT | WeightKind::GUARD | WeightKind::DIR

        );

    }

    #[test]

    fn weightset_from_consensus() {

        use rand::Rng;

        let now = SystemTime::now();

        let one_hour = Duration::new(3600, 0);

        let mut rng = testing_rng();

        let mut bld = MdConsensus::builder();

        bld.consensus_method(34)

            .lifetime(Lifetime::new(now, now + one_hour, now + 2 * one_hour).unwrap())

            .weights(TESTVEC_PARAMS.parse().unwrap());

        // We're going to add a huge amount of unmeasured bandwidth,

        // and a reasonable amount of  measured bandwidth.

        for _ in 0..10 {

            rs_builder()

                .identity(rng.random::<[u8; 20]>().into()) // random id

                .weight(RW::Unmeasured(1_000_000))

                .set_flags(RelayFlags::GUARD | RelayFlags::EXIT)

                .build_into(&mut bld)

                .unwrap();

        }

        for n in 0..30 {

            rs_builder()

                .identity(rng.random::<[u8; 20]>().into()) // random id

                .weight(RW::Measured(1_000 * n))

                .set_flags(RelayFlags::GUARD | RelayFlags::EXIT)

                .build_into(&mut bld)

                .unwrap();

        }

        let consensus = bld.testing_consensus().unwrap();

        let params = NetParameters::default();

        let ws = WeightSet::from_consensus(&consensus, &params);

        assert_eq!(ws.bandwidth_fn, BandwidthFn::MeasuredOnly);

        assert_eq!(ws.shift, 0);

        assert_eq!(ws.w[0].as_guard, 5904);

        assert_eq!(ws.w[5].as_guard, 5904);

        assert_eq!(ws.w[5].as_middle, 4096);

    }

}