Skip to content

Commit

Permalink
Make SystemTime mockable via object
Browse files Browse the repository at this point in the history
- Adds trait SystemTimeClock
- Adds default implementation RealSystemTimeClock
- Adds ServerConfigParameter system_time_clock: Arc<dyn SystemTimeClock>
- Replaces all SystemTime::now calls in proto with SystemTimeClock.now
  calls (there were 2)

This is a backwards-compatible change.
  • Loading branch information
gretchenfrage committed Dec 15, 2024
1 parent 05f6e67 commit e69f6f8
Show file tree
Hide file tree
Showing 4 changed files with 44 additions and 8 deletions.
36 changes: 35 additions & 1 deletion quinn-proto/src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ use crate::{
cid_generator::{ConnectionIdGenerator, HashedConnectionIdGenerator},
crypto::{self, HandshakeTokenKey, HmacKey},
shared::ConnectionId,
Duration, RandomConnectionIdGenerator, VarInt, VarIntBoundsExceeded,
Duration, RandomConnectionIdGenerator, SystemTime, VarInt, VarIntBoundsExceeded,
DEFAULT_SUPPORTED_VERSIONS, MAX_CID_SIZE,
};

Expand Down Expand Up @@ -216,6 +216,8 @@ pub struct ServerConfig {
pub(crate) max_incoming: usize,
pub(crate) incoming_buffer_size: u64,
pub(crate) incoming_buffer_size_total: u64,

pub(crate) system_time_clock: Arc<dyn SystemTimeClock>,
}

impl ServerConfig {
Expand All @@ -239,6 +241,8 @@ impl ServerConfig {
max_incoming: 1 << 16,
incoming_buffer_size: 10 << 20,
incoming_buffer_size_total: 100 << 20,

system_time_clock: Arc::new(RealSystemTimeClock),
}
}

Expand Down Expand Up @@ -334,6 +338,16 @@ impl ServerConfig {
self.incoming_buffer_size_total = incoming_buffer_size_total;
self
}

/// Object to get current [`SystemTime`]
///
/// This exists to allow system time to be mocked in tests, or wherever else desired.
///
/// Defaults to [`RealSystemTimeClock`], which simply calls [`SystemTime::now()`](SystemTime::now).
pub fn system_time_clock(&mut self, system_time_clock: Arc<dyn SystemTimeClock>) -> &mut Self {
self.system_time_clock = system_time_clock;
self
}
}

#[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
Expand Down Expand Up @@ -388,6 +402,7 @@ impl fmt::Debug for ServerConfig {
"incoming_buffer_size_total",
&self.incoming_buffer_size_total,
)
// system_time_clock not debug
.finish_non_exhaustive()
}
}
Expand Down Expand Up @@ -503,3 +518,22 @@ impl From<VarIntBoundsExceeded> for ConfigError {
Self::OutOfBounds
}
}

/// Object to get current [`SystemTime`]
///
/// This exists to allow system time to be mocked in tests, or wherever else desired.
pub trait SystemTimeClock: Send + Sync {
/// Get [`SystemTime::now()`](SystemTime::now) or the mocked equivalent
fn now(&self) -> SystemTime;
}

/// Default implementation of [`SystemTime`]
///
/// Implements `now` by calling [`SystemTime::now()`](SystemTime::now).
pub struct RealSystemTimeClock;

impl SystemTimeClock for RealSystemTimeClock {
fn now(&self) -> SystemTime {
SystemTime::now()
}
}
9 changes: 5 additions & 4 deletions quinn-proto/src/endpoint.rs
Original file line number Diff line number Diff line change
Expand Up @@ -31,8 +31,8 @@ use crate::{
},
token::TokenDecodeError,
transport_parameters::{PreferredAddress, TransportParameters},
Duration, Instant, ResetToken, RetryToken, Side, SystemTime, Transmit, TransportConfig,
TransportError, INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
Duration, Instant, ResetToken, RetryToken, Side, Transmit, TransportConfig, TransportError,
INITIAL_MTU, MAX_CID_SIZE, MIN_INITIAL_SIZE, RESET_TOKEN_SIZE,
};

/// The main entry point to the library
Expand Down Expand Up @@ -506,7 +506,8 @@ impl Endpoint {
&header.token,
) {
Ok(token)
if token.issued + server_config.retry_token_lifetime > SystemTime::now() =>
if token.issued + server_config.retry_token_lifetime
> server_config.system_time_clock.now() =>
{
(Some(header.dst_cid), token.orig_dst_cid)
}
Expand Down Expand Up @@ -774,7 +775,7 @@ impl Endpoint {

let token = RetryToken {
orig_dst_cid: incoming.packet.header.dst_cid,
issued: SystemTime::now(),
issued: server_config.system_time_clock.now(),
}
.encode(
&*server_config.token_key,
Expand Down
2 changes: 1 addition & 1 deletion quinn-proto/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ pub use rustls;
mod config;
pub use config::{
AckFrequencyConfig, ClientConfig, ConfigError, EndpointConfig, IdleTimeout, MtuDiscoveryConfig,
ServerConfig, TransportConfig,
RealSystemTimeClock, ServerConfig, SystemTimeClock, TransportConfig,
};

pub mod crypto;
Expand Down
5 changes: 3 additions & 2 deletions quinn/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,8 +65,9 @@ pub use proto::{
congestion, crypto, AckFrequencyConfig, ApplicationClose, Chunk, ClientConfig, ClosedStream,
ConfigError, ConnectError, ConnectionClose, ConnectionError, ConnectionId,
ConnectionIdGenerator, ConnectionStats, Dir, EcnCodepoint, EndpointConfig, FrameStats,
FrameType, IdleTimeout, MtuDiscoveryConfig, PathStats, ServerConfig, Side, StreamId, Transmit,
TransportConfig, TransportErrorCode, UdpStats, VarInt, VarIntBoundsExceeded, Written,
FrameType, IdleTimeout, MtuDiscoveryConfig, PathStats, RealSystemTimeClock, ServerConfig, Side,
StreamId, SystemTimeClock, Transmit, TransportConfig, TransportErrorCode, UdpStats, VarInt,
VarIntBoundsExceeded, Written,
};
#[cfg(any(feature = "rustls-aws-lc-rs", feature = "rustls-ring"))]
pub use rustls;
Expand Down

0 comments on commit e69f6f8

Please sign in to comment.