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::applicationlayer::ApplicationLayer;
pub use crate::error::Error; pub use crate::error::Error;
pub use crate::sessionid::SessionId; 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(), .unwrap(),
)); ));
let mut ts = 0;
for test_loop in 0..256 { for test_loop in 0..256 {
let time_ticks = (test_loop * 10000) as i64;
for host in [&alice_host, &bob_host] { for host in [&alice_host, &bob_host] {
let send_to_other = |data: &mut [u8]| { let send_to_other = |data: &mut [u8]| {
if std::ptr::eq(host, &alice_host) { if std::ptr::eq(host, &alice_host) {
@ -139,8 +139,7 @@ mod tests {
loop { loop {
if let Some(qi) = host.queue.lock().unwrap().pop_back() { if let Some(qi) = host.queue.lock().unwrap().pop_back() {
let qi_len = qi.len(); let qi_len = qi.len();
ts += 1; let r = rc.receive(host, &0, send_to_other, &mut data_buf, qi, mtu_buffer.len(), time_ticks);
let r = rc.receive(host, &0, send_to_other, &mut data_buf, qi, mtu_buffer.len(), ts);
if r.is_ok() { if r.is_ok() {
let r = r.unwrap(); let r = r.unwrap();
match r { match r {
@ -152,13 +151,7 @@ mod tests {
assert!(!data.iter().any(|x| *x != 0x12)); assert!(!data.iter().any(|x| *x != 0x12));
} }
ReceiveResult::OkNewSession(new_session) => { ReceiveResult::OkNewSession(new_session) => {
println!( println!("zssp: new session at {} ({})", host.this_name, u64::from(new_session.id));
"zssp: {} => {} ({}): OkNewSession ({})",
host.other_name,
host.this_name,
qi_len,
u64::from(new_session.id)
);
let mut hs = host.session.lock().unwrap(); let mut hs = host.session.lock().unwrap();
assert!(hs.is_none()); assert!(hs.is_none());
let _ = hs.insert(Arc::new(new_session)); let _ = hs.insert(Arc::new(new_session));
@ -191,11 +184,15 @@ mod tests {
if !security_info.0.eq(key_id.as_ref()) { if !security_info.0.eq(key_id.as_ref()) {
*key_id = security_info.0; *key_id = security_info.0;
println!( println!(
"zssp: new key at {}: fingerprint {} ratchet {} kyber {}", "zssp: new key at {}: fingerprint {} ratchet {} kyber {} latest role {}",
host.this_name, host.this_name,
hex::to_string(key_id.as_ref()), hex::to_string(key_id.as_ref()),
security_info.1, 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()); .is_ok());
} }
if (test_loop % 8) == 0 && test_loop >= 8 && host.this_name.eq("alice") { if (test_loop % 13) == 0 && test_loop > 0 {
session.service(host, send_to_other, &[], mtu_buffer.len(), test_loop as i64, true); 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, 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. /// 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 /// 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 send_key: Secret<AES_KEY_SIZE>, // Send side AES-GCM key
receive_cipher_pool: Mutex<Vec<Box<AesGcm>>>, // Pool of reusable sending ciphers receive_cipher_pool: Mutex<Vec<Box<AesGcm>>>, // Pool of reusable sending ciphers
send_cipher_pool: Mutex<Vec<Box<AesGcm>>>, // Pool of reusable receiving 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 confirmed: bool, // We have confirmed that the other side has this key
jedi: bool, // True if Kyber1024 was used (both sides enabled) 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) 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. /// Note that the role can switch through the course of a session. It's the side that most recently
/// Fragmentation is not considered when authenticating the entire packet. A separate header /// initiated a session or a rekey event. Initiator is Alice, responder is Bob.
/// check code is used to make fragmentation itself more robust, but that's outside the scope
/// of AEAD authentication.
#[derive(Clone, Copy)] #[derive(Clone, Copy)]
#[repr(C, packed)] pub enum Role {
struct CanonicalHeader(pub u64, pub u32); Alice,
Bob,
}
impl<Application: ApplicationLayer> Session<Application> { impl<Application: ApplicationLayer> Session<Application> {
/// Create a new session and send an initial key offer message to the other end. /// 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. /// 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(); let state = self.state.read().unwrap();
state.session_keys[state.cur_session_key_idx] state.session_keys[state.cur_session_key_idx]
.as_ref() .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. /// 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, force_expire: bool,
) { ) {
let state = self.state.read().unwrap(); let state = self.state.read().unwrap();
if (force_expire if state.session_keys[state.cur_session_key_idx].as_ref().map_or(true, |k| {
|| state.session_keys[state.cur_session_key_idx] matches!(k.role, Role::Bob)
.as_ref() && (force_expire || self.send_counter.load(Ordering::Relaxed) >= k.rekey_at_counter || current_time >= k.rekey_at_time)
.map_or(true, |k| k.rekey_needed.load(Ordering::Relaxed) || current_time < k.rekey_at_time)) }) && state
&& state .offer
.offer .as_ref()
.as_ref() .map_or(true, |o| (current_time - o.creation_time) > Application::REKEY_RATE_LIMIT_MS)
.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) { if let Some(remote_s_public) = P384PublicKey::from_bytes(&self.remote_s_public_p384_bytes) {
let mut offer = None; let mut offer = None;
@ -604,6 +594,8 @@ impl<Application: ApplicationLayer> ReceiveContext<Application> {
if other_session_key.ratchet_count < this_ratchet_count { if other_session_key.ratchet_count < this_ratchet_count {
state.cur_session_key_idx = key_index; 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. // Check rate limits.
if let Some(session) = session.as_ref() { 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); return Err(Error::RateLimited);
} }
} else { } else {
@ -969,9 +961,12 @@ impl<Application: ApplicationLayer> ReceiveContext<Application> {
hybrid_kk.is_some(), hybrid_kk.is_some(),
); );
let next_key_index = (next_ratchet_count as usize) & 1;
let mut state = session.state.write().unwrap(); let mut state = session.state.write().unwrap();
let _ = state.remote_session_id.replace(alice_session_id); 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); drop(state);
// Bob now has final key state for this exchange. Yay! Now reply to Alice so she can construct it. // 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, send_key,
receive_cipher_pool: Mutex::new(Vec::with_capacity(2)), receive_cipher_pool: Mutex::new(Vec::with_capacity(2)),
send_cipher_pool: Mutex::new(Vec::with_capacity(2)), send_cipher_pool: Mutex::new(Vec::with_capacity(2)),
rekey_needed: AtomicBool::new(false), role,
confirmed, confirmed,
jedi, jedi,
} }
@ -1460,8 +1455,6 @@ impl SessionKey {
fn get_send_cipher(&self, counter: u64) -> Result<Box<AesGcm>, Error> { fn get_send_cipher(&self, counter: u64) -> Result<Box<AesGcm>, Error> {
if counter < self.expire_at_counter { if counter < self.expire_at_counter {
self.rekey_needed
.store(counter >= self.rekey_at_counter, std::sync::atomic::Ordering::Relaxed);
Ok(self Ok(self
.send_cipher_pool .send_cipher_pool
.lock() .lock()