Fix tests, remove some dead code, fix a bug that prevented rekey rate limiting from working.

This commit is contained in:
Adam Ierymenko 2023-01-06 20:39:20 -05:00
parent 73e6be7959
commit d2b49130b2
3 changed files with 37 additions and 47 deletions

View file

@ -9,4 +9,4 @@ pub mod constants;
pub use crate::applicationlayer::ApplicationLayer;
pub use crate::error::Error;
pub use crate::sessionid::SessionId;
pub use crate::zssp::{ReceiveContext, ReceiveResult, Session};
pub use crate::zssp::{ReceiveContext, ReceiveResult, Role, Session};

View file

@ -119,8 +119,8 @@ mod tests {
.unwrap(),
));
let mut ts = 0;
for test_loop in 0..256 {
let time_ticks = (test_loop * 10000) as i64;
for host in [&alice_host, &bob_host] {
let send_to_other = |data: &mut [u8]| {
if std::ptr::eq(host, &alice_host) {
@ -139,8 +139,7 @@ mod tests {
loop {
if let Some(qi) = host.queue.lock().unwrap().pop_back() {
let qi_len = qi.len();
ts += 1;
let r = rc.receive(host, &0, send_to_other, &mut data_buf, qi, mtu_buffer.len(), ts);
let r = rc.receive(host, &0, send_to_other, &mut data_buf, qi, mtu_buffer.len(), time_ticks);
if r.is_ok() {
let r = r.unwrap();
match r {
@ -152,13 +151,7 @@ mod tests {
assert!(!data.iter().any(|x| *x != 0x12));
}
ReceiveResult::OkNewSession(new_session) => {
println!(
"zssp: {} => {} ({}): OkNewSession ({})",
host.other_name,
host.this_name,
qi_len,
u64::from(new_session.id)
);
println!("zssp: new session at {} ({})", host.this_name, u64::from(new_session.id));
let mut hs = host.session.lock().unwrap();
assert!(hs.is_none());
let _ = hs.insert(Arc::new(new_session));
@ -191,11 +184,15 @@ mod tests {
if !security_info.0.eq(key_id.as_ref()) {
*key_id = security_info.0;
println!(
"zssp: new key at {}: fingerprint {} ratchet {} kyber {}",
"zssp: new key at {}: fingerprint {} ratchet {} kyber {} latest role {}",
host.this_name,
hex::to_string(key_id.as_ref()),
security_info.1,
security_info.2
security_info.3,
match security_info.2 {
Role::Alice => "A",
Role::Bob => "B",
}
);
}
}
@ -208,8 +205,8 @@ mod tests {
)
.is_ok());
}
if (test_loop % 8) == 0 && test_loop >= 8 && host.this_name.eq("alice") {
session.service(host, send_to_other, &[], mtu_buffer.len(), test_loop as i64, true);
if (test_loop % 13) == 0 && test_loop > 0 {
session.service(host, send_to_other, &[], mtu_buffer.len(), time_ticks, true);
}
}
}

View file

@ -43,15 +43,6 @@ pub enum ReceiveResult<'a, H: ApplicationLayer> {
Ignored,
}
/// Was this side the one who sent the first offer (Alice) or countered (Bob).
///
/// Note that the role can switch through the course of a session. It's the side that most recently
/// initiated a session or a rekey event. Initiator is Alice, responder is Bob.
pub enum Role {
Alice,
Bob,
}
/// State information to associate with receiving contexts such as sockets or remote paths/endpoints.
///
/// This holds the data structures used to defragment incoming packets that are not associated with an
@ -102,7 +93,7 @@ struct SessionKey {
send_key: Secret<AES_KEY_SIZE>, // Send side AES-GCM key
receive_cipher_pool: Mutex<Vec<Box<AesGcm>>>, // Pool of reusable sending ciphers
send_cipher_pool: Mutex<Vec<Box<AesGcm>>>, // Pool of reusable receiving ciphers
rekey_needed: AtomicBool, // We have reached or exceeded the counter
role: Role, // Was this side Alice or Bob?
confirmed: bool, // We have confirmed that the other side has this key
jedi: bool, // True if Kyber1024 was used (both sides enabled)
}
@ -118,15 +109,15 @@ struct EphemeralOffer {
alice_hk_keypair: Option<pqc_kyber::Keypair>, // Kyber1024 key pair (PQ hybrid ephemeral key for Alice)
}
/// "Canonical header" for generating 96-bit AES-GCM nonce and for inclusion in key exchange HMACs.
/// Was this side the one who sent the first offer (Alice) or countered (Bob).
///
/// This is basically the actual header but with fragment count and fragment total set to zero.
/// Fragmentation is not considered when authenticating the entire packet. A separate header
/// check code is used to make fragmentation itself more robust, but that's outside the scope
/// of AEAD authentication.
/// Note that the role can switch through the course of a session. It's the side that most recently
/// initiated a session or a rekey event. Initiator is Alice, responder is Bob.
#[derive(Clone, Copy)]
#[repr(C, packed)]
struct CanonicalHeader(pub u64, pub u32);
pub enum Role {
Alice,
Bob,
}
impl<Application: ApplicationLayer> Session<Application> {
/// Create a new session and send an initial key offer message to the other end.
@ -297,11 +288,11 @@ impl<Application: ApplicationLayer> Session<Application> {
}
/// Get the shared key fingerprint, ratchet count, and whether Kyber was used, or None if not yet established.
pub fn status(&self) -> Option<([u8; 16], u64, bool)> {
pub fn status(&self) -> Option<([u8; 16], u64, Role, bool)> {
let state = self.state.read().unwrap();
state.session_keys[state.cur_session_key_idx]
.as_ref()
.map(|k| (k.secret_fingerprint, k.ratchet_count, k.jedi))
.map(|k| (k.secret_fingerprint, k.ratchet_count, k.role, k.jedi))
}
/// This function needs to be called on each session at least every SERVICE_INTERVAL milliseconds.
@ -322,14 +313,13 @@ impl<Application: ApplicationLayer> Session<Application> {
force_expire: bool,
) {
let state = self.state.read().unwrap();
if (force_expire
|| state.session_keys[state.cur_session_key_idx]
.as_ref()
.map_or(true, |k| k.rekey_needed.load(Ordering::Relaxed) || current_time < k.rekey_at_time))
&& state
.offer
.as_ref()
.map_or(true, |o| (current_time - o.creation_time) > Application::REKEY_RATE_LIMIT_MS)
if state.session_keys[state.cur_session_key_idx].as_ref().map_or(true, |k| {
matches!(k.role, Role::Bob)
&& (force_expire || self.send_counter.load(Ordering::Relaxed) >= k.rekey_at_counter || current_time >= k.rekey_at_time)
}) && state
.offer
.as_ref()
.map_or(true, |o| (current_time - o.creation_time) > Application::REKEY_RATE_LIMIT_MS)
{
if let Some(remote_s_public) = P384PublicKey::from_bytes(&self.remote_s_public_p384_bytes) {
let mut offer = None;
@ -604,6 +594,8 @@ impl<Application: ApplicationLayer> ReceiveContext<Application> {
if other_session_key.ratchet_count < this_ratchet_count {
state.cur_session_key_idx = key_index;
}
} else {
state.cur_session_key_idx = key_index;
}
}
}
@ -687,7 +679,7 @@ impl<Application: ApplicationLayer> ReceiveContext<Application> {
// Check rate limits.
if let Some(session) = session.as_ref() {
if (current_time - session.state.read().unwrap().last_remote_offer) < Application::REKEY_RATE_LIMIT_MS {
if (session.state.read().unwrap().last_remote_offer + Application::REKEY_RATE_LIMIT_MS) > current_time {
return Err(Error::RateLimited);
}
} else {
@ -969,9 +961,12 @@ impl<Application: ApplicationLayer> ReceiveContext<Application> {
hybrid_kk.is_some(),
);
let next_key_index = (next_ratchet_count as usize) & 1;
let mut state = session.state.write().unwrap();
let _ = state.remote_session_id.replace(alice_session_id);
let _ = state.session_keys[(next_ratchet_count as usize) & 1].replace(session_key);
let _ = state.session_keys[next_key_index].replace(session_key);
state.last_remote_offer = current_time;
drop(state);
// Bob now has final key state for this exchange. Yay! Now reply to Alice so she can construct it.
@ -1452,7 +1447,7 @@ impl SessionKey {
send_key,
receive_cipher_pool: Mutex::new(Vec::with_capacity(2)),
send_cipher_pool: Mutex::new(Vec::with_capacity(2)),
rekey_needed: AtomicBool::new(false),
role,
confirmed,
jedi,
}
@ -1460,8 +1455,6 @@ impl SessionKey {
fn get_send_cipher(&self, counter: u64) -> Result<Box<AesGcm>, Error> {
if counter < self.expire_at_counter {
self.rekey_needed
.store(counter >= self.rekey_at_counter, std::sync::atomic::Ordering::Relaxed);
Ok(self
.send_cipher_pool
.lock()