use rusty_time_nts::aead::NtsKeys;
use rusty_time_nts::ef::{self, NONCE_LEN, UNIQUE_ID_LEN};
use rusty_time_nts::ke;
use std::collections::VecDeque;
use std::time::Duration;
const TARGET_COOKIES: usize = 8;
pub struct NtsSession {
keys: NtsKeys,
cookies: VecDeque<Vec<u8>>,
pending_unique_id: [u8; UNIQUE_ID_LEN],
pub ntp_server: String,
pub ntp_port: u16,
pub ke_host: String,
}
#[derive(Debug)]
pub enum NtsSessionError {
Ke(ke::KeError),
CookiesExhausted,
Protect(ef::NtsError),
}
impl std::fmt::Display for NtsSessionError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
NtsSessionError::Ke(e) => write!(f, "{e}"),
NtsSessionError::CookiesExhausted => {
write!(f, "NTS cookies exhausted; re-run key establishment")
}
NtsSessionError::Protect(e) => write!(f, "{e}"),
}
}
}
impl NtsSession {
pub fn establish(
ke_host: &str,
ke_port: u16,
timeout: Duration,
extra_roots: &[rusty_time_nts::tls::pki_types::CertificateDer<'static>],
) -> Result<Self, NtsSessionError> {
let result = ke::establish_with_roots(ke_host, ke_port, timeout, extra_roots)
.map_err(NtsSessionError::Ke)?;
Ok(NtsSession {
keys: result.keys,
cookies: result.cookies.into(),
pending_unique_id: [0u8; UNIQUE_ID_LEN],
ntp_server: result.ntp_server,
ntp_port: result.ntp_port,
ke_host: ke_host.to_string(),
})
}
pub fn cookies_held(&self) -> usize {
self.cookies.len()
}
pub fn keys(&self) -> &NtsKeys {
&self.keys
}
pub fn unspent_cookies(&self) -> Vec<Vec<u8>> {
self.cookies.iter().cloned().collect()
}
pub fn resume(
keys: NtsKeys,
cookies: Vec<Vec<u8>>,
ntp_server: String,
ntp_port: u16,
ke_host: String,
) -> Self {
NtsSession {
keys,
cookies: cookies.into(),
pending_unique_id: [0u8; UNIQUE_ID_LEN],
ntp_server,
ntp_port,
ke_host,
}
}
#[cfg(test)]
pub fn for_test(keys: NtsKeys, cookies: Vec<Vec<u8>>) -> Self {
NtsSession {
keys,
cookies: cookies.into(),
pending_unique_id: [0u8; UNIQUE_ID_LEN],
ntp_server: "loopback".into(),
ntp_port: 123,
ke_host: "loopback".into(),
}
}
pub fn protect(&mut self, header: &[u8; 48]) -> Result<Vec<u8>, NtsSessionError> {
let cookie = self
.cookies
.pop_front()
.ok_or(NtsSessionError::CookiesExhausted)?;
let mut unique_id = [0u8; UNIQUE_ID_LEN];
let mut nonce = [0u8; NONCE_LEN];
ke::fill_random(&mut unique_id).map_err(NtsSessionError::Ke)?;
ke::fill_random(&mut nonce).map_err(NtsSessionError::Ke)?;
self.pending_unique_id = unique_id;
let want = TARGET_COOKIES.saturating_sub(self.cookies.len()).min(7);
ef::protect_request(header, &unique_id, &cookie, want, &self.keys.c2s, &nonce)
.map_err(NtsSessionError::Protect)
}
pub fn verify(&mut self, packet: &[u8]) -> Result<(), NtsSessionError> {
let verified = ef::verify_response(packet, &self.keys.s2c, &self.pending_unique_id)
.map_err(NtsSessionError::Protect)?;
for cookie in verified.new_cookies {
if self.cookies.len() < TARGET_COOKIES * 2 {
self.cookies.push_back(cookie);
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn protect_spends_a_cookie_and_refuses_when_empty() {
let mut s = NtsSession {
keys: NtsKeys {
c2s: [1; 32],
s2c: [2; 32],
},
cookies: vec![vec![0xAB; 32]].into(),
pending_unique_id: [0; UNIQUE_ID_LEN],
ntp_server: "example".into(),
ntp_port: 123,
ke_host: "example".into(),
};
let mut header = [0u8; 48];
header[0] = (4 << 3) | 3;
assert_eq!(s.cookies_held(), 1);
let packet = s.protect(&header).expect("protect");
assert_eq!(s.cookies_held(), 0, "cookie must be spent exactly once");
assert!(packet.len() > 48);
assert_ne!(s.pending_unique_id, [0u8; UNIQUE_ID_LEN]);
assert!(matches!(
s.protect(&header),
Err(NtsSessionError::CookiesExhausted)
));
}
#[test]
fn each_request_uses_a_distinct_unique_id() {
let mut s = NtsSession {
keys: NtsKeys {
c2s: [1; 32],
s2c: [2; 32],
},
cookies: vec![vec![0xAB; 32], vec![0xCD; 32]].into(),
pending_unique_id: [0; UNIQUE_ID_LEN],
ntp_server: "example".into(),
ntp_port: 123,
ke_host: "example".into(),
};
let mut header = [0u8; 48];
header[0] = (4 << 3) | 3;
let _ = s.protect(&header).expect("first");
let first = s.pending_unique_id;
let _ = s.protect(&header).expect("second");
assert_ne!(first, s.pending_unique_id, "unique id must not repeat");
}
}