Skip to content

Commit

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

This is a backwards-compatible change.
  • Loading branch information
gretchenfrage authored and djc committed Dec 17, 2024
1 parent c5f81be commit 51606c6
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) time_source: Arc<dyn TimeSource>,
}

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,

time_source: Arc::new(StdSystemTime),
}
}

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 [`StdSystemTime`], which simply calls [`SystemTime::now()`](SystemTime::now).
pub fn time_source(&mut self, time_source: Arc<dyn TimeSource>) -> &mut Self {
self.time_source = time_source;
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 TimeSource: Send + Sync {
/// Get [`SystemTime::now()`](SystemTime::now) or the mocked equivalent
fn now(&self) -> SystemTime;
}

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

impl TimeSource for StdSystemTime {
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,
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 @@ -504,7 +504,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.time_source.now() =>
{
(Some(header.dst_cid), token.orig_dst_cid)
}
Expand Down Expand Up @@ -769,7 +770,7 @@ impl Endpoint {

let token = RetryToken {
orig_dst_cid: incoming.packet.header.dst_cid,
issued: SystemTime::now(),
issued: server_config.time_source.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,
ServerConfig, StdSystemTime, TimeSource, 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, ServerConfig, Side, StdSystemTime,
StreamId, TimeSource, 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 51606c6

Please sign in to comment.