use std::{
collections::HashMap,
net::SocketAddr,
time::{Duration, SystemTime},
};
use irtt_proto::{
decode_request, echo_packet_len, encode_echo_reply, encode_open_reply, verify_packet_hmac,
Clock, DecodedParams, DecodedRequestKind, EchoReply, OpenReply, Params, StampAt,
TimestampFields, FLAG_CLOSE, FLAG_OPEN, FLAG_REPLY,
};
use crate::{
clock::{saturating_ns, wall_ns_of, ClockSample, ClockSource, SystemClock},
config::ServerConfig,
error::ServerError,
fill::{echo_payload_len, negotiate_server_fill},
negotiate::negotiate_params,
session::{RateDecision, Session},
token::{OsTokenSource, TokenSource, TOKEN_ATTEMPTS},
};
#[derive(Debug)]
pub struct ServerCore {
config: ServerConfig,
sessions: HashMap<u64, Session>,
tokens: Box<dyn TokenSource>,
clock: Box<dyn ClockSource>,
}
impl ServerCore {
#[must_use]
pub fn new(config: ServerConfig) -> Self {
Self::with_token_source(config, Box::new(OsTokenSource))
}
pub(crate) fn with_token_source(config: ServerConfig, tokens: Box<dyn TokenSource>) -> Self {
Self::with_sources(config, tokens, Box::new(SystemClock::new()))
}
pub(crate) fn with_sources(
config: ServerConfig,
tokens: Box<dyn TokenSource>,
clock: Box<dyn ClockSource>,
) -> Self {
Self {
config,
sessions: HashMap::new(),
tokens,
clock,
}
}
#[must_use]
pub fn config(&self) -> &ServerConfig {
&self.config
}
#[must_use]
pub fn session_count(&self) -> usize {
self.sessions.len()
}
pub(crate) fn session(&self, token: u64) -> Option<&Session> {
self.sessions.get(&token)
}
pub fn handle_datagram(
&mut self,
peer: SocketAddr,
packet: &[u8],
) -> Result<Option<OutboundDatagram>, ServerError> {
self.process(peer, packet, None)
}
pub(crate) fn handle_received_datagram(
&mut self,
peer: SocketAddr,
packet: &[u8],
kernel_rx_timestamp: Option<SystemTime>,
) -> Result<Option<OutboundDatagram>, ServerError> {
self.process(peer, packet, kernel_rx_timestamp)
}
fn process(
&mut self,
peer: SocketAddr,
packet: &[u8],
kernel_rx_timestamp: Option<SystemTime>,
) -> Result<Option<OutboundDatagram>, ServerError> {
let Ok(request) = decode_request(packet) else {
return Ok(None);
};
let kind = self.classify(&request.kind, packet.len(), kernel_rx_timestamp);
if request.hmac_present != self.config.hmac_key().is_some() {
return Ok(None);
}
if let Some(key) = self.config.hmac_key() {
if verify_packet_hmac(key, packet).is_err() {
return Ok(None);
}
}
let now_ns = match kind {
ClassifiedKind::Echo(ref echo) => echo.received_at.mono_ns,
_ => self.clock.sample().mono_ns,
};
self.release_expired_sessions(now_ns);
match kind {
ClassifiedKind::Open { no_test, params } => {
self.handle_open(peer, no_test, params, now_ns)
}
ClassifiedKind::Close { token } => {
self.handle_close(peer, token);
Ok(None)
}
ClassifiedKind::Echo(echo) => self.handle_echo(peer, echo),
}
}
pub(crate) fn maintain(&mut self) {
let now_ns = self.clock.sample().mono_ns;
self.release_expired_sessions(now_ns);
}
fn release_expired_sessions(&mut self, now_ns: i64) {
let idle_timeout_ns = saturating_ns(self.config.idle_timeout());
self.sessions
.retain(|_, session| !session.is_idle_expired(now_ns, idle_timeout_ns));
}
fn classify<'a>(
&mut self,
kind: &DecodedRequestKind<'a>,
datagram_len: usize,
kernel_rx_timestamp: Option<SystemTime>,
) -> ClassifiedKind<'a> {
match *kind {
DecodedRequestKind::Open { no_test, params } => {
ClassifiedKind::Open { no_test, params }
}
DecodedRequestKind::Close { token } => ClassifiedKind::Close { token },
DecodedRequestKind::Echo {
token, sequence, ..
} => {
let received_at = self.clock.sample();
ClassifiedKind::Echo(EchoRequest {
token,
sequence,
datagram_len,
received_at,
receive_wall_ns: preferred_receive_wall_ns(kernel_rx_timestamp, received_at),
})
}
}
}
fn handle_echo(
&mut self,
peer: SocketAddr,
request: EchoRequest,
) -> Result<Option<OutboundDatagram>, ServerError> {
if request.datagram_len > self.config.max_packet_length() {
return Ok(None);
}
let Some(session) = self.sessions.get_mut(&request.token) else {
return Ok(None);
};
if !same_endpoint(session.peer(), peer) {
return Ok(None);
}
let now_ns = request.received_at.mono_ns;
let next_rate = match session.rate_state().charge(now_ns) {
RateDecision::Limited { next } => {
let next_lifetime = session.lifetime_state().refreshed(now_ns);
session.commit_rate_state(next);
session.commit_lifetime_state(next_lifetime);
return Ok(None);
}
RateDecision::Allowed { next } => next,
};
let closing = self.config.max_test_duration().is_some_and(|maximum| {
session
.lifetime_state()
.max_duration_reached(now_ns, saturating_ns(maximum))
});
let next = session.receive_state().accepted(request.sequence);
let next_lifetime = session.lifetime_state().served(now_ns);
let sent_at = self.clock.sample();
let params = session.params();
let payload_len = echo_payload_len(self.config.hmac_key().is_some(), params);
let traffic_class = transport_traffic_class(params);
let reply = EchoReply {
flags: FLAG_REPLY | if closing { FLAG_CLOSE } else { 0 },
token: request.token,
sequence: request.sequence,
recv_count: params.received_stats.has_count().then_some(next.count),
recv_window: params.received_stats.has_window().then_some(next.window),
timestamps: timestamp_fields(
params,
request.received_at,
sent_at,
request.receive_wall_ns,
),
payload: session.fill().payload(payload_len),
};
let packet = encode_echo_reply(&reply, params, self.config.hmac_key())
.map_err(|source| ServerError::ReplyEncoding { source })?;
if closing {
self.sessions.remove(&request.token);
} else {
session.commit_rate_state(next_rate);
session.commit_receive_state(next);
session.commit_lifetime_state(next_lifetime);
}
Ok(Some(OutboundDatagram::new(packet, traffic_class)))
}
fn handle_open(
&mut self,
peer: SocketAddr,
no_test: bool,
encoded_params: &[u8],
now_ns: i64,
) -> Result<Option<OutboundDatagram>, ServerError> {
let Ok(decoded) = Params::decode_with_presence(encoded_params) else {
return Ok(None);
};
if !open_params_are_admissible(&decoded) {
return Ok(None);
}
let mut params = negotiate_params(decoded.params, &self.config);
let fill = negotiate_server_fill(&mut params);
if !negotiated_params_are_admissible(¶ms, &self.config) {
return Ok(None);
}
if no_test {
let reply = OpenReply {
flags: FLAG_OPEN | FLAG_REPLY | FLAG_CLOSE,
token: 0,
params,
};
return self
.encode_reply(&reply)
.map(|packet| Some(OutboundDatagram::unmarked(packet)));
}
if self.sessions.len() >= self.config.max_sessions() {
return Ok(None);
}
let token = self.allocate_token()?;
let reply = OpenReply {
flags: FLAG_OPEN | FLAG_REPLY,
token,
params,
};
let packet = self.encode_reply(&reply)?;
self.sessions.insert(
token,
Session::new(peer, reply.params, fill, &self.config, now_ns),
);
Ok(Some(OutboundDatagram::unmarked(packet)))
}
fn handle_close(&mut self, peer: SocketAddr, token: u64) {
let owns_session = self
.session(token)
.is_some_and(|session| same_endpoint(session.peer(), peer));
if owns_session {
self.sessions.remove(&token);
}
}
fn encode_reply(&self, reply: &OpenReply) -> Result<Vec<u8>, ServerError> {
encode_open_reply(reply, self.config.hmac_key())
.map_err(|source| ServerError::ReplyEncoding { source })
}
fn allocate_token(&mut self) -> Result<u64, ServerError> {
for _ in 0..TOKEN_ATTEMPTS {
let token = self.tokens.next_token()?;
if token != 0 && !self.sessions.contains_key(&token) {
return Ok(token);
}
}
Err(ServerError::TokenExhausted {
attempts: TOKEN_ATTEMPTS,
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutboundDatagram {
bytes: Vec<u8>,
traffic_class: u8,
}
impl OutboundDatagram {
fn new(bytes: Vec<u8>, traffic_class: u8) -> Self {
Self {
bytes,
traffic_class,
}
}
fn unmarked(bytes: Vec<u8>) -> Self {
Self::new(bytes, 0)
}
#[must_use]
pub fn bytes(&self) -> &[u8] {
&self.bytes
}
#[must_use]
pub fn traffic_class(&self) -> u8 {
self.traffic_class
}
#[must_use]
pub fn into_bytes(self) -> Vec<u8> {
self.bytes
}
}
impl AsRef<[u8]> for OutboundDatagram {
fn as_ref(&self) -> &[u8] {
self.bytes()
}
}
fn transport_traffic_class(params: &Params) -> u8 {
u8::try_from(params.dscp).unwrap_or(0)
}
enum ClassifiedKind<'a> {
Open { no_test: bool, params: &'a [u8] },
Close { token: u64 },
Echo(EchoRequest),
}
#[derive(Debug, Clone, Copy)]
struct EchoRequest {
token: u64,
sequence: u32,
datagram_len: usize,
received_at: ClockSample,
receive_wall_ns: i64,
}
pub(crate) const MAX_KERNEL_RX_LAG: Duration = Duration::from_secs(1);
pub(crate) fn preferred_receive_wall_ns(
kernel_rx_timestamp: Option<SystemTime>,
received_at: ClockSample,
) -> i64 {
let userspace_wall_ns = received_at.wall_ns;
let Some(kernel_wall_ns) = kernel_rx_timestamp.and_then(wall_ns_of) else {
return userspace_wall_ns;
};
let lag_ns = i128::from(userspace_wall_ns) - i128::from(kernel_wall_ns);
if (0..=i128::from(saturating_ns(MAX_KERNEL_RX_LAG))).contains(&lag_ns) {
kernel_wall_ns
} else {
userspace_wall_ns
}
}
fn timestamp_fields(
params: &Params,
received: ClockSample,
sent: ClockSample,
receive_wall_ns: i64,
) -> TimestampFields {
let received = received.not_after(sent);
let receive_wall_ns = receive_wall_ns.min(sent.wall_ns);
let per_clock = |sample: ClockSample| {
(
params.clock.has_wall().then_some(sample.wall_ns),
params.clock.has_mono().then_some(sample.mono_ns),
)
};
let arrival = (
params.clock.has_wall().then_some(receive_wall_ns),
params.clock.has_mono().then_some(received.mono_ns),
);
let absent = (None, None);
let (recv, midpoint, send) = match params.stamp_at {
StampAt::None => (absent, absent, absent),
StampAt::Receive => (arrival, absent, absent),
StampAt::Send => (absent, absent, per_clock(sent)),
StampAt::Both => (arrival, absent, per_clock(sent)),
StampAt::Midpoint => (absent, per_clock(received.midpoint(sent)), absent),
};
TimestampFields {
recv_wall: recv.0,
recv_mono: recv.1,
midpoint_wall: midpoint.0,
midpoint_mono: midpoint.1,
send_wall: send.0,
send_mono: send.1,
}
}
fn same_endpoint(session: SocketAddr, peer: SocketAddr) -> bool {
match (session, peer) {
(SocketAddr::V6(session), SocketAddr::V6(peer)) => {
session.ip() == peer.ip()
&& session.port() == peer.port()
&& session.scope_id() == peer.scope_id()
}
(session, peer) => session == peer,
}
}
fn open_params_are_admissible(decoded: &DecodedParams) -> bool {
let admissible = |present: bool, value: i64| !present || value > 0;
admissible(decoded.presence.duration_ns, decoded.params.duration_ns)
&& admissible(decoded.presence.interval_ns, decoded.params.interval_ns)
}
fn negotiated_params_are_admissible(params: &Params, config: &ServerConfig) -> bool {
if params.stamp_at != StampAt::None && params.clock == Clock::Unspecified {
return false;
}
echo_packet_len(config.hmac_key().is_some(), params)
.is_ok_and(|len| len <= config.max_packet_length())
}