use std::pin::Pin;
use web_time_compat::{Duration, Instant};
use derive_builder::Builder;
use educe::Educe;
use futures::FutureExt;
use futures::future::{self, FusedFuture};
use pin_project::pin_project;
use rand::distr::Distribution;
use tracing::error;
use tor_cell::chancell::msg::{Padding, PaddingNegotiate};
use tor_config::impl_standard_builder;
use tor_error::into_internal;
use tor_rtcompat::SleepProvider;
use tor_units::IntegerMilliseconds;
#[pin_project(project = PaddingTimerProj)]
pub(crate) struct Timer<R: SleepProvider> {
sleep_prov: R,
parameters: Option<PreparedParameters>,
selected_timeout: Option<Duration>,
trigger_at: Option<Instant>,
#[pin]
waker: Option<R::SleepFuture>,
}
#[derive(Debug, Copy, Clone, Eq, PartialEq, Builder)]
#[builder(build_fn(error = "ParametersError", private, name = "build_inner"))]
pub struct Parameters {
#[builder(default = "1500.into()")]
pub(crate) low: IntegerMilliseconds<u32>,
#[builder(default = "9500.into()")]
pub(crate) high: IntegerMilliseconds<u32>,
}
#[derive(Clone, Debug, thiserror::Error)]
#[non_exhaustive]
pub enum ParametersError {
#[error("Cannot construct padding parameters: low bound was above the high bound.")]
InvalidRange,
}
impl ParametersBuilder {
pub fn build(&self) -> Result<Parameters, ParametersError> {
let parameters = self.build_inner()?;
if parameters.low > parameters.high {
return Err(ParametersError::InvalidRange);
}
Ok(parameters)
}
}
impl_standard_builder! { Parameters: !Deserialize + !Builder + !Default }
impl Parameters {
pub fn padding_negotiate_cell(&self) -> Result<PaddingNegotiate, tor_error::Bug> {
let get = |input: IntegerMilliseconds<u32>| {
input
.try_map(TryFrom::try_from)
.map_err(into_internal!("padding negotiate out of range"))
};
Ok(PaddingNegotiate::start(get(self.low)?, get(self.high)?))
}
pub fn default_padding() -> Self {
Parameters::builder().build().expect("build succeeded")
}
pub fn disabled() -> Self {
Parameters {
low: 0.into(),
high: 0.into(),
}
}
}
#[derive(Debug, Clone)]
struct PreparedParameters {
x_distribution_ms: rand::distr::Uniform<u32>,
}
#[derive(Educe)]
#[educe(Debug)]
enum SleepInstructions<'f, R: SleepProvider> {
Immediate {
now: Instant,
},
Forever,
Waker(#[educe(Debug(ignore))] Pin<&'f mut R::SleepFuture>),
}
impl<R: SleepProvider> Timer<R> {
#[allow(dead_code)]
pub(crate) fn new(sleep_prov: R, parameters: Parameters) -> crate::Result<Self> {
let parameters = parameters.prepare()?;
let selected_timeout = parameters.select_timeout();
Ok(Timer {
sleep_prov,
parameters: Some(parameters),
selected_timeout: Some(selected_timeout),
trigger_at: None,
waker: None,
})
}
pub(crate) fn new_disabled(
sleep_prov: R,
parameters: Option<Parameters>,
) -> crate::Result<Self> {
Ok(Timer {
sleep_prov,
parameters: parameters.map(|p| p.prepare()).transpose()?,
selected_timeout: None,
trigger_at: None,
waker: None,
})
}
pub(crate) fn disable(self: &mut Pin<&mut Self>) {
*self.as_mut().project().selected_timeout = None;
}
pub(crate) fn enable(self: &mut Pin<&mut Self>) {
if !self.is_enabled() {
self.as_mut().select_fresh_timeout();
}
}
pub(crate) fn reconfigure(
self: &mut Pin<&mut Self>,
parameters: &Parameters,
) -> crate::Result<()> {
*self.as_mut().project().parameters = Some(parameters.prepare()?);
Ok(())
}
pub(crate) fn is_enabled(&self) -> bool {
self.selected_timeout.is_some()
}
fn select_fresh_timeout(self: Pin<&mut Self>) {
let mut self_ = self.project();
let timeout = self_.parameters.as_ref().map(|p| p.select_timeout());
*self_.selected_timeout = timeout;
*self_.trigger_at = None;
self_.waker.set(None);
}
pub(crate) fn note_cell_sent(self: &mut Pin<&mut Self>) {
let self_ = self.as_mut().project();
*self_.trigger_at = None;
}
fn prepare_to_sleep(mut self: Pin<&mut Self>, now: Option<Instant>) -> SleepInstructions<R> {
let mut self_ = self.as_mut().project();
let timeout = match self_.selected_timeout {
None => return SleepInstructions::Forever,
Some(t) => *t,
};
if self_.waker.is_some() {
let waker = self
.project()
.waker
.as_pin_mut()
.expect("None but we just checked");
return SleepInstructions::Waker(waker);
}
let now = now.unwrap_or_else(|| self_.sleep_prov.now());
let trigger_at = match self_.trigger_at {
Some(t) => t,
None => self_.trigger_at.insert(match now.checked_add(timeout) {
None => {
error!("bug: timeout overflowed computing next channel padding. Disabling.");
self.disable();
return SleepInstructions::Forever;
}
Some(r) => r,
}),
};
let remaining = trigger_at.checked_duration_since(now).unwrap_or_default();
if remaining.is_zero() {
return SleepInstructions::Immediate { now };
}
if self_.waker.is_none() {
self_.waker.set(Some(self_.sleep_prov.sleep(remaining)));
}
let waker = self
.project()
.waker
.as_pin_mut()
.expect("None but we just inserted!");
SleepInstructions::Waker(waker)
}
pub(crate) fn next(self: Pin<&mut Self>) -> impl FusedFuture<Output = Padding> + '_ {
self.next_inner().fuse()
}
async fn next_inner(mut self: Pin<&mut Self>) -> Padding {
let now = loop {
match self.as_mut().prepare_to_sleep(None) {
SleepInstructions::Forever => future::pending().await,
SleepInstructions::Immediate { now } => break now,
SleepInstructions::Waker(waker) => waker.await,
}
self.as_mut().project().waker.set(None);
};
self.as_mut().select_fresh_timeout();
self.as_mut().prepare_to_sleep(Some(now));
Padding::new()
}
}
impl Parameters {
fn prepare(self) -> Result<PreparedParameters, tor_error::Bug> {
Ok(PreparedParameters {
x_distribution_ms: rand::distr::Uniform::new_inclusive(
self.low.as_millis(),
self.high.as_millis(),
)
.map_err(into_internal!("Parameters were not a valid range."))?,
})
}
}
impl PreparedParameters {
fn select_timeout(&self) -> Duration {
let mut rng = rand::rng();
let ms = std::cmp::max(
self.x_distribution_ms.sample(&mut rng),
self.x_distribution_ms.sample(&mut rng),
);
Duration::from_millis(ms.into())
}
}
#[cfg(test)]
mod test {
#![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_time_subtraction)]
#![allow(clippy::useless_vec)]
#![allow(clippy::needless_pass_by_value)]
use super::*;
use futures::future::ready;
use futures::select_biased;
use itertools::{Itertools, izip};
use statrs::distribution::ContinuousCDF;
use tokio::pin;
use tokio_crate::{self as tokio};
use tor_rtcompat::*;
async fn assert_not_ready<R: Runtime>(timer: &mut Pin<&mut Timer<R>>) {
select_biased! {
_ = timer.as_mut().next() => panic!("unexpectedly ready"),
_ = ready(()) => { },
};
}
async fn assert_is_ready<R: Runtime>(timer: &mut Pin<&mut Timer<R>>) {
let _: Padding = select_biased! {
p = timer.as_mut().next() => p,
_ = ready(()) => panic!("pad timer failed to yield"),
};
}
#[test]
fn timer_impl() {
let runtime = tor_rtcompat::tokio::TokioNativeTlsRuntime::create().unwrap();
#[allow(deprecated)] let runtime = tor_rtmock::MockSleepRuntime::new(runtime);
let parameters = Parameters {
low: 1000.into(),
high: 1000.into(),
};
let () = runtime.block_on(async {
let timer = Timer::new(runtime.clone(), parameters).unwrap();
pin!(timer);
assert_eq! { true, timer.is_enabled() }
assert_eq! { timer.as_mut().trigger_at, None };
assert_not_ready(&mut timer).await;
runtime.advance(Duration::from_millis(999)).await;
assert_not_ready(&mut timer).await;
runtime.advance(Duration::from_millis(1)).await;
assert_is_ready(&mut timer).await;
assert_not_ready(&mut timer).await;
runtime.advance(Duration::from_millis(1001)).await;
assert_is_ready(&mut timer).await;
runtime.advance(Duration::from_millis(500)).await;
timer.note_cell_sent();
assert_eq! { timer.as_mut().trigger_at, None };
let () = select_biased! {
_ = ready(()) => { },
_ = timer.as_mut().next() => panic!(),
};
assert_eq! { timer.as_mut().trigger_at, None };
timer.disable();
runtime.advance(Duration::from_millis(2000)).await;
assert_eq! { timer.as_mut().selected_timeout, None };
assert_eq! { false, timer.is_enabled() }
assert_not_ready(&mut timer).await;
timer.enable();
runtime.advance(Duration::from_millis(3000)).await;
assert_eq! { true, timer.is_enabled() }
assert_not_ready(&mut timer).await;
runtime.advance(Duration::from_millis(1000)).await;
assert_is_ready(&mut timer).await;
});
let () = runtime.block_on(async {
let timer = Timer::new(runtime.clone(), parameters).unwrap();
pin!(timer);
assert! { timer.as_mut().selected_timeout.is_some() };
assert! { timer.as_mut().trigger_at.is_none() };
*timer.as_mut().project().selected_timeout = Some(Duration::MAX);
assert_not_ready(&mut timer).await;
dbg!(timer.as_mut().project().trigger_at);
assert_eq! { false, timer.is_enabled() }
});
let () = runtime.block_on(async {
let timer = Timer::new_disabled(runtime.clone(), None).unwrap();
assert! { timer.parameters.is_none() };
pin!(timer);
assert_not_ready(&mut timer).await;
assert! { timer.as_mut().selected_timeout.is_none() };
assert! { timer.as_mut().trigger_at.is_none() };
});
let () = runtime.block_on(async {
let timer = Timer::new_disabled(runtime.clone(), Some(parameters)).unwrap();
assert! { timer.parameters.is_some() };
pin!(timer);
assert_not_ready(&mut timer).await;
runtime.advance(Duration::from_millis(3000)).await;
assert_not_ready(&mut timer).await;
timer.as_mut().enable();
assert_not_ready(&mut timer).await;
runtime.advance(Duration::from_millis(3000)).await;
assert_is_ready(&mut timer).await;
});
}
#[test]
#[allow(clippy::print_stderr)]
fn timeout_distribution() {
#[allow(non_snake_case)]
let mut N = 100_0000;
#[allow(non_upper_case_globals)]
const n: usize = 100;
const P_GOOD: f64 = 0.05; const P_BAD: f64 = 1e-12;
loop {
eprintln!("padding distribution test, n={} N={}", n, N);
let min = 5000;
let max = 17000; assert_eq!(0, (max - min) % (n as u32));
let cdf = (0..=n)
.map(|bi| {
let b = (bi as f64) / (n as f64);
b.powi(2)
})
.collect_vec();
let pdf = cdf
.iter()
.cloned()
.tuple_windows()
.map(|(p, q)| q - p)
.collect_vec();
let exp = pdf.iter().cloned().map(|p| p * f64::from(N)).collect_vec();
assert!(exp.iter().cloned().all(|ei| ei >= 5.));
let mut obs = [0_u32; n];
let params = Parameters {
low: min.into(),
high: (max - 1).into(), }
.prepare()
.unwrap();
for _ in 0..N {
let xx = params.select_timeout();
let ms = xx.as_millis();
let ms = u32::try_from(ms).unwrap();
assert!(ms >= min);
assert!(ms < max);
let bi = ((ms - min) * (n as u32)) / (max - min);
obs[bi as usize] += 1;
}
let chi2 = izip!(&obs, &exp)
.map(|(&oi, &ei)| (f64::from(oi) - ei).powi(2) / ei)
.sum::<f64>();
let chi2_distr = statrs::distribution::ChiSquared::new(n as _).unwrap();
let p = 1. - chi2_distr.cdf(chi2);
eprintln!(
"padding distribution test, n={} N={} chi2={} p={}",
n, N, chi2, p
);
if p >= P_GOOD {
break;
}
for (i, (&oi, &ei)) in izip!(&obs, &exp).enumerate() {
eprintln!("bi={:4} OI={:4} EI={}", i, oi, ei);
}
if p < P_BAD {
panic!("distribution is wrong (p < {:e})", P_BAD);
}
N *= 10;
}
}
#[test]
fn parameters_range() {
let ms100 = IntegerMilliseconds::new(100);
let ms1000 = IntegerMilliseconds::new(1000);
let ms1500 = IntegerMilliseconds::new(1500);
let ms9500 = IntegerMilliseconds::new(9500);
let p = Parameters::builder().build().unwrap();
assert_eq!(
p,
Parameters {
low: ms1500,
high: ms9500
}
);
assert!(p.prepare().is_ok());
let mut pb = Parameters::builder();
pb.low(ms100);
pb.high(ms1000);
let p = pb.build().unwrap();
assert_eq!(
p,
Parameters {
low: ms100,
high: ms1000
}
);
let p = p.prepare().unwrap();
let range = Duration::try_from(ms100).unwrap()..=Duration::try_from(ms1000).unwrap();
for _ in 1..100 {
assert!(range.contains(&p.select_timeout()));
}
let mut pb = Parameters::builder();
pb.low(ms1000);
pb.high(ms1000);
let p = pb.build().unwrap();
assert!(p.prepare().is_ok());
let mut pb = Parameters::builder();
pb.low(ms1000);
pb.high(ms100);
let e = pb.build().unwrap_err();
assert!(matches!(e, ParametersError::InvalidRange));
}
}