saving experimental changes

performance/src/frag_cache.rs

@@ -2,6 +2,7 @@ use std::collections::hash_map::RandomState;
 use std::hash::{BuildHasher, Hash, Hasher};
 use std::mem::MaybeUninit;
 
+use crate::application::CryptoLayer;
 use crate::crypto::AES_GCM_NONCE_SIZE;
 use crate::fragged::Assembled;
 use crate::proto::{MAX_FRAGMENTS, MAX_UNASSOCIATED_FRAGMENTS, MAX_UNASSOCIATED_PACKETS, MAX_UNASSOCIATED_PACKET_SIZE};
@@ -15,18 +16,18 @@ struct PacketMetadata {
     creation_time: i64,
 }
 
-pub(crate) struct UnassociatedFragCache<Fragment> {
+pub(crate) struct UnassociatedFragCache<Crypto: CryptoLayer> {
     dos_salt: RandomState,
     frags_first_unused: usize,
     frags_unused_size: usize,
     map: [PacketMetadata; MAX_UNASSOCIATED_PACKETS],
-    frags: [MaybeUninit<Fragment>; MAX_UNASSOCIATED_FRAGMENTS],
+    frags: [MaybeUninit<Crypto::IncomingPacketBuffer>; MAX_UNASSOCIATED_FRAGMENTS],
     map_idx: [u32; MAX_UNASSOCIATED_FRAGMENTS],
 }
 /// A combination of a hash table cache and a ring buffer for unassociated fragments.
 /// Designed specifically to be extremely DDOS resistant.
 /// This datastructure takes raw unauthenticated fragments straight from the network.
-impl<Fragment> UnassociatedFragCache<Fragment> {
+impl<Crypto: CryptoLayer> UnassociatedFragCache<Crypto> {
     pub(crate) fn new() -> Self {
         Self {
             dos_salt: RandomState::new(),
@@ -46,24 +47,24 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
     }
     /// Add a fragment and return an assembled packet container if all fragments have been received.
     /// Will check that aad is the same for all fragments.
+    /// Returns true if the fragment is a part of a new fragment.
     pub(crate) fn assemble(
         &mut self,
         nonce: &[u8; AES_GCM_NONCE_SIZE],
         remote_address: impl Hash,
         fragment_size: usize,
-        fragment: Fragment,
+        fragment: Crypto::IncomingPacketBuffer,
         fragment_no: usize,
         fragment_count: usize,
-        timeout_interval: i64,
         current_time: i64,
-        ret_assembled: &mut Assembled<Fragment>,
-    ) {
+        ret_assembled: &mut Assembled<Crypto::IncomingPacketBuffer>,
+    ) -> Option<i64> {
         debug_assert!(MAX_FRAGMENTS < MAX_UNASSOCIATED_FRAGMENTS);
         if fragment_no >= fragment_count
             || fragment_count > MAX_FRAGMENTS
             || fragment_size > MAX_UNASSOCIATED_PACKET_SIZE
         {
-            return;
+            return None;
         }
 
         let mut hasher = self.dos_salt.build_hasher();
@@ -94,7 +95,7 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
         } else if self.map[idx0].key == 0 || self.map[idx1].key == 0 {
             if (fragment_count as usize) > self.frags_unused_size {
                 // There are not enough free fragment slots so attempt to expire a bunch of entries.
-                self.check_for_expiry(timeout_interval, current_time);
+                let _ = self.check_for_expiry_inner(Crypto::SETTINGS.resend_time as i64, current_time);
             }
             if self.map[idx0].key == 0 {
                 idx0
@@ -103,20 +104,21 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
             }
         } else {
             // No room for a new entry so attempt to expire a bunch of entries.
-            self.check_for_expiry(timeout_interval, current_time);
+            let _ = self.check_for_expiry_inner(Crypto::SETTINGS.resend_time as i64, current_time);
             if self.map[idx0].key == 0 {
                 idx0
             } else if self.map[idx1].key == 0 {
                 idx1
             } else {
                 // Give up and drop the fragment.
-                return;
+                return None;
             }
         };
-
+        let mut new_expiry = None;
         if self.map[idx].key == 0 {
             // This is a new entry so initialize it.
             if (fragment_count as usize) <= self.frags_unused_size {
+                new_expiry = Some(current_time + Crypto::SETTINGS.fragment_assembly_timeout as i64);
                 let entry = &mut self.map[idx];
                 entry.key = key;
                 entry.frags_idx = self.frags_first_unused as u32;
@@ -132,7 +134,7 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
                 }
             } else {
                 // If there are not enough free fragment slots by this point we just drop the fragment.
-                return;
+                return None;
             }
         }
         let entry = &mut self.map[idx];
@@ -160,8 +162,12 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
                 self.invalidate::<false>(idx);
             }
         }
+        new_expiry
     }
-    pub(crate) fn check_for_expiry(&mut self, timeout: i64, current_time: i64) {
+    pub(crate) fn check_for_expiry(&mut self, current_time: i64) -> i64 {
+        self.check_for_expiry_inner(Crypto::SETTINGS.fragment_assembly_timeout as i64, current_time)
+    }
+    fn check_for_expiry_inner(&mut self, timeout: i64, current_time: i64) -> i64 {
         while self.frags_unused_size < self.frags.len() {
             // Check if we can drop the entry at the start of the ring buffer.
             let frag_idx = (self.frags_first_unused + self.frags_unused_size) % self.frags.len();
@@ -169,12 +175,14 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
             debug_assert!(map_idx < self.map.len());
 
             let entry = &mut self.map[map_idx];
-            if entry.creation_time + timeout < current_time {
+            let expiry = entry.creation_time + timeout;
+            if expiry <= current_time {
                 self.invalidate::<true>(map_idx);
             } else {
-                break;
+                return expiry;
             }
         }
+        return i64::MAX;
     }
 
     fn invalidate<const DROP: bool>(&mut self, idx: usize) {
@@ -210,7 +218,7 @@ impl<Fragment> UnassociatedFragCache<Fragment> {
         }
     }
 }
-impl<Fragment> Drop for UnassociatedFragCache<Fragment> {
+impl<Crypto: CryptoLayer> Drop for UnassociatedFragCache<Crypto> {
     fn drop(&mut self) {
         for i in 0..self.map.len() {
             if self.map[i].key != 0 {
@@ -233,11 +241,28 @@ fn test_cache() {
         drop(x);
         r.wrapping_mul(0x2545F4914F6CDD1Du64)
     }
+    use crate::crypto_impl::*;
+    struct Crypto {}
+    impl CryptoLayer for Crypto {
+        type Rng = rand_core::OsRng;
+        type PrpEnc = OpenSSLAes256Enc;
+        type PrpDec = OpenSSLAes256Dec;
+        type Aead = OpenSSLAesGcm;
+        type AeadPool = OpenSSLAesGcmPool;
+        type Hash = CrateSha512;
+        type Hmac = CrateHmacSha512;
+        type PublicKey = CrateP384PublicKey;
+        type KeyPair = CrateP384KeyPair;
+        type Kem = CrateKyber1024PrivateKey;
 
-    let mut cache = UnassociatedFragCache::new();
+        type SessionData = ();
+        type IncomingPacketBuffer = Vec<u8>;
+    }
+
+    let mut cache = UnassociatedFragCache::<Crypto>::new();
     let mut assembled = Assembled::new();
 
-    let mut time = 1;
+    let mut time = 0;
     let mut in_progress = Vec::new();
     let mut in_progress_fragments = 0;
     // A basic fuzzer for testing the cache.
@@ -267,11 +292,10 @@ fn test_cache() {
                         fragment,
                         j,
                         fragment_count,
-                        1,
                         time,
                         &mut assembled,
                     );
-                    time += 1;
+                    time += 200;
                 }
             }
             if drop >= fragment_count {
@@ -310,11 +334,10 @@ fn test_cache() {
                             fragment,
                             no as usize,
                             fragment_count as usize,
-                            1000,
                             time,
                             &mut assembled,
                         );
-                        time += 1;
+                        time += 200;
                         in_progress_fragments -= 1;
 
                         if packet.len() > 0 {

performance/src/handshake_cache.rs

@@ -12,7 +12,7 @@ pub(crate) struct UnassociatedHandshakeCache<Application: CryptoLayer> {
 /// SoA format
 struct CacheInner<Crypto: CryptoLayer> {
     local_ids: [Option<NonZeroU32>; MAX_UNASSOCIATED_HANDSHAKE_STATES],
-    timeouts: [i64; MAX_UNASSOCIATED_HANDSHAKE_STATES],
+    expiries: [i64; MAX_UNASSOCIATED_HANDSHAKE_STATES],
     handshakes: [Option<Arc<StateB2<Crypto>>>; MAX_UNASSOCIATED_HANDSHAKE_STATES],
 }
 
@@ -25,7 +25,7 @@ impl<Application: CryptoLayer> UnassociatedHandshakeCache<Application> {
             has_pending: AtomicBool::new(false),
             cache: RwLock::new(CacheInner {
                 local_ids: std::array::from_fn(|_| None),
-                timeouts: std::array::from_fn(|_| 0),
+                expiries: std::array::from_fn(|_| 0),
                 handshakes: std::array::from_fn(|_| None),
             }),
         }
@@ -39,55 +39,64 @@ impl<Application: CryptoLayer> UnassociatedHandshakeCache<Application> {
         }
         None
     }
-    pub(crate) fn insert(&self, local_id: NonZeroU32, state: Arc<StateB2<Application>>, current_time: i64) {
+    pub(crate) fn insert(
+        &self,
+        local_id: NonZeroU32,
+        state: Arc<StateB2<Application>>,
+        current_time: i64,
+    ) -> Option<i64> {
         let mut cache = self.cache.write().unwrap();
         let mut idx = 0;
         for i in 0..cache.local_ids.len() {
-            if cache.local_ids[i].is_none() || cache.timeouts[i] < current_time {
+            if cache.local_ids[i].is_none() || cache.expiries[i] <= current_time {
                 idx = i;
                 break;
             } else if cache.local_ids[i] == Some(local_id) {
-                return;
+                return None;
             }
         }
+        let expiry = current_time + Application::SETTINGS.fragment_assembly_timeout as i64;
         cache.local_ids[idx] = Some(local_id);
-        cache.timeouts[idx] = current_time + Application::SETTINGS.fragment_assembly_timeout as i64;
+        cache.expiries[idx] = expiry;
         cache.handshakes[idx] = Some(state);
         self.has_pending.store(true, Ordering::Release);
+        return Some(expiry);
     }
     pub(crate) fn remove(&self, local_id: NonZeroU32) -> bool {
         let mut cache = self.cache.write().unwrap();
         for (i, id) in cache.local_ids.iter().enumerate() {
             if *id == Some(local_id) {
                 cache.local_ids[i] = None;
-                cache.timeouts[i] = 0;
+                cache.expiries[i] = 0;
                 cache.handshakes[i] = None;
                 return true;
             }
         }
         false
     }
-    pub(crate) fn service(&self, current_time: i64) {
+    pub(crate) fn service(&self, current_time: i64) -> i64 {
         // Only check for expiration if we have a pending packet.
         // This check is allowed to have false positives for simplicity's sake.
+        let mut next_service_time = i64::MAX;
         if self.has_pending.swap(false, Ordering::Acquire) {
             // Check for packet expiration
             let mut cache = self.cache.write().unwrap();
-            let mut has_pending = false;
             for i in 0..cache.local_ids.len() {
                 if cache.local_ids[i].is_some() {
-                    if cache.timeouts[i] < current_time {
+                    let expiry = cache.expiries[i];
+                    if expiry <= current_time {
                         cache.local_ids[i] = None;
-                        cache.timeouts[i] = 0;
+                        cache.expiries[i] = 0;
                         cache.handshakes[i] = None;
                     } else {
-                        has_pending = true;
+                        next_service_time = next_service_time.min(expiry);
                     }
                 }
             }
-            if has_pending {
+            if next_service_time < i64::MAX {
                 self.has_pending.store(true, Ordering::Release);
             }
         }
+        next_service_time
     }
 }

performance/src/zeta.rs

@@ -420,6 +420,9 @@ pub(crate) fn trans_to_a1<Crypto: CryptoLayer, App: ApplicationLayer<Crypto = Cr
 
     session_map.insert(kid_recv, Arc::downgrade(&session));
     session_queue.push_reserved(queue_idx, Arc::downgrade(&session), Reverse(resend_timer));
+    let _ = ctx.reduce_next_service_time(resend_timer);
+    drop(session_map);
+    drop(session_queue);
 
     send(&mut x1, None);
 
@@ -447,7 +450,7 @@ pub(crate) fn received_x1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     n: &[u8; AES_GCM_NONCE_SIZE],
     x1: &mut [u8],
     send: impl FnOnce(&mut [u8], Option<&Crypto::PrpEnc>),
-) -> Result<(), ReceiveError> {
+) -> Result<bool, ReceiveError> {
     use FaultType::*;
     //    <- s
     //    ...
@@ -566,7 +569,7 @@ pub(crate) fn received_x1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
 
     set_header(&mut x2, kid_send.get(), &to_nonce(PACKET_TYPE_HANDSHAKE_RESPONSE, c));
 
-    ctx.unassociated_handshake_states.insert(
+    let next_service_time = ctx.unassociated_handshake_states.insert(
         kid_recv,
         Arc::new(StateB2 {
             ratchet_state,
@@ -580,12 +583,16 @@ pub(crate) fn received_x1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
         }),
         app.time(),
     );
+    let mut reduced_service_time = false;
+    if let Some(next_service_time) = next_service_time {
+        reduced_service_time = ctx.reduce_next_service_time(next_service_time);
+    }
 
     send(
         &mut x2,
         Some(&Crypto::PrpEnc::new(&hk_send[..AES_256_KEY_SIZE].try_into().unwrap())),
     );
-    Ok(())
+    Ok(reduced_service_time)
 }
 /// Corresponds to Transition Algorithm 3 found in Section 4.3.
 pub(crate) fn received_x2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypto = Crypto>>(
@@ -596,7 +603,7 @@ pub(crate) fn received_x2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     n: &[u8; AES_GCM_NONCE_SIZE],
     x2: &mut [u8],
     send: impl FnOnce(&mut [u8], Option<&Crypto::PrpEnc>),
-) -> Result<bool, ReceiveError> {
+) -> Result<(bool, bool), ReceiveError> {
     use FaultType::*;
     //    <- e, ee, ekem1, psk
     //    -> s, se
@@ -770,8 +777,9 @@ pub(crate) fn received_x2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
             .lock()
             .unwrap()
             .change_priority(session.queue_idx, Reverse(resend_timer));
+        let reduced = ctx.reduce_next_service_time(resend_timer);
 
-        Ok(x3)
+        Ok((x3, reduced))
     })();
 
     match result {
@@ -781,10 +789,10 @@ pub(crate) fn received_x2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
             let current_time = app.time();
             timeout_trans(app, ctx, session, kex_lock, state, current_time, send);
         }
-        Ok(ref mut packet) => send(packet, Some(&session.state.read().unwrap().hk_send)),
+        Ok((ref mut packet, _)) => send(packet, Some(&session.state.read().unwrap().hk_send)),
         _ => {}
     }
-    result.map(|_| should_warn_missing_ratchet)
+    result.map(|(_, reduced_service_time)| (should_warn_missing_ratchet, reduced_service_time))
 }
 fn send_control<Crypto: CryptoLayer, const CAP: usize>(
     session: &Arc<Session<Crypto>>,
@@ -816,7 +824,7 @@ pub(crate) fn received_x3_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     kid: NonZeroU32,
     x3: &mut [u8],
     send: impl FnOnce(&mut [u8], Option<&Crypto::PrpEnc>),
-) -> Result<(Arc<Session<Crypto>>, bool), ReceiveError> {
+) -> Result<(Arc<Session<Crypto>>, bool, bool), ReceiveError> {
     use FaultType::*;
     //    -> s, se
     if x3.len() < HANDSHAKE_COMPLETION_MIN_SIZE {
@@ -916,7 +924,7 @@ pub(crate) fn received_x3_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
                 )
                 .map_err(|_| ReceiveError::RatchetStorageError)?;
 
-                let session = {
+                let (session, reduced_service_time) = {
                     let mut session_map = ctx.session_map.write().unwrap();
                     use std::collections::hash_map::Entry::*;
                     let entry = match session_map.entry(zeta.kid_recv) {
@@ -966,7 +974,7 @@ pub(crate) fn received_x3_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
                     session_queue.push_reserved(queue_idx, Arc::downgrade(&session), Reverse(resend_timer));
                     entry.insert(Arc::downgrade(&session));
 
-                    session
+                    (session, ctx.reduce_next_service_time(resend_timer))
                 };
                 let state = session.state.read().unwrap();
                 let mut c1 = ArrayVec::<u8, HEADERED_KEY_CONFIRMATION_SIZE>::new();
@@ -974,7 +982,7 @@ pub(crate) fn received_x3_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
                 send_control(&session, &state, PACKET_TYPE_KEY_CONFIRM, c1, send);
                 drop(state);
 
-                Ok((session, should_warn_missing_ratchet))
+                Ok((session, should_warn_missing_ratchet, reduced_service_time))
             }
             Err(()) => Err(ReceiveError::RatchetStorageError),
         }
@@ -994,7 +1002,7 @@ pub(crate) fn received_c1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     n: &[u8; AES_GCM_NONCE_SIZE],
     c1: &[u8],
     send: impl FnOnce(&mut [u8], Option<&Crypto::PrpEnc>),
-) -> Result<bool, ReceiveError> {
+) -> Result<(bool, bool), ReceiveError> {
     use FaultType::*;
 
     if c1.len() != KEY_CONFIRMATION_SIZE {
@@ -1024,6 +1032,7 @@ pub(crate) fn received_c1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     if !session.window.update(c) {
         return Err(fault!(ExpiredCounter, true));
     }
+    let mut reduced_service_time = false;
 
     let just_establised = is_other && matches!(&state.beta, ZetaAutomata::A3 { .. });
     if is_other {
@@ -1058,6 +1067,7 @@ pub(crate) fn received_c1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
                 .lock()
                 .unwrap()
                 .change_priority(session.queue_idx, Reverse(timeout_timer));
+            reduced_service_time = ctx.reduce_next_service_time(timeout_timer);
             state = session.state.read().unwrap();
         }
     }
@@ -1068,7 +1078,7 @@ pub(crate) fn received_c1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
         return Err(fault!(OutOfSequence, true));
     }
 
-    Ok(just_establised)
+    Ok((just_establised, reduced_service_time))
 }
 /// Corresponds to the trivial Transition Algorithm described for processing C_2 packets found in
 /// Section 4.3.
@@ -1079,7 +1089,7 @@ pub(crate) fn received_c2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     kid: NonZeroU32,
     n: &[u8; AES_GCM_NONCE_SIZE],
     c2: &[u8],
-) -> Result<(), ReceiveError> {
+) -> Result<bool, ReceiveError> {
     use FaultType::*;
 
     if c2.len() != ACKNOWLEDGEMENT_SIZE {
@@ -1120,7 +1130,8 @@ pub(crate) fn received_c2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
         .lock()
         .unwrap()
         .change_priority(session.queue_idx, Reverse(timeout_timer));
-    Ok(())
+
+    Ok(ctx.reduce_next_service_time(timeout_timer))
 }
 /// Corresponds to the trivial Transition Algorithm described for processing D packets found in
 /// Section 4.3.
@@ -1346,7 +1357,7 @@ pub(crate) fn received_k1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     n: &[u8; AES_GCM_NONCE_SIZE],
     k1: &mut [u8],
     send: impl FnOnce(&mut [u8], Option<&Crypto::PrpEnc>),
-) -> Result<(), ReceiveError> {
+) -> Result<bool, ReceiveError> {
     use FaultType::*;
     //    -> s
     //    <- s
@@ -1478,13 +1489,14 @@ pub(crate) fn received_k1_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
             .lock()
             .unwrap()
             .change_priority(session.queue_idx, Reverse(resend_timer));
+        let reduced_service_time = ctx.reduce_next_service_time(resend_timer);
         let state = session.state.read().unwrap();
 
         if !send_control(session, &state, PACKET_TYPE_REKEY_COMPLETE, k2, send) {
             return Err(fault!(OutOfSequence, true));
         }
 
-        Ok(())
+        Ok(reduced_service_time)
     })();
 
     if matches!(result, Err(ReceiveError::ByzantineFault { .. })) {
@@ -1501,7 +1513,7 @@ pub(crate) fn received_k2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     n: &[u8; AES_GCM_NONCE_SIZE],
     k2: &mut [u8],
     send: impl FnOnce(&mut [u8], Option<&Crypto::PrpEnc>),
-) -> Result<(), ReceiveError> {
+) -> Result<bool, ReceiveError> {
     use FaultType::*;
     //    <- e, ee, se
     if k2.len() != REKEY_SIZE {
@@ -1601,6 +1613,7 @@ pub(crate) fn received_k2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
                 .lock()
                 .unwrap()
                 .change_priority(session.queue_idx, Reverse(resend_timer));
+            let reduced_service_time = ctx.reduce_next_service_time(resend_timer);
             let state = session.state.read().unwrap();
 
             let mut c1 = ArrayVec::<u8, HEADERED_KEY_CONFIRMATION_SIZE>::new();
@@ -1609,7 +1622,7 @@ pub(crate) fn received_k2_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
                 return Err(fault!(OutOfSequence, true));
             }
 
-            Ok(())
+            Ok(reduced_service_time)
         } else {
             unreachable!()
         }
@@ -1627,7 +1640,7 @@ pub(crate) fn send_payload<Crypto: CryptoLayer>(
     payload: &[u8],
     mut send: impl FnMut(&mut [u8]) -> bool,
     mtu_sized_buffer: &mut [u8],
-) -> Result<(), SendError> {
+) -> Result<bool, SendError> {
     use SendError::*;
     let mtu = mtu_sized_buffer.len();
     if mtu < MIN_TRANSPORT_MTU {
@@ -1673,7 +1686,7 @@ pub(crate) fn send_payload<Crypto: CryptoLayer>(
         state.hk_send.encrypt_in_place(header_auth.try_into().unwrap());
 
         if !send(&mut mtu_sized_buffer[..HEADER_SIZE + fragment_len]) {
-            return Ok(());
+            return Ok(false);
         }
         i = j;
     }
@@ -1692,7 +1705,7 @@ pub(crate) fn send_payload<Crypto: CryptoLayer>(
     state.hk_send.encrypt_in_place(header_auth.try_into().unwrap());
 
     if !send(&mut mtu_sized_buffer[..HEADER_SIZE + fragment_len]) {
-        return Ok(());
+        return Ok(false);
     }
 
     drop(state);
@@ -1705,8 +1718,10 @@ pub(crate) fn send_payload<Crypto: CryptoLayer>(
             .lock()
             .unwrap()
             .change_priority(session.queue_idx, Reverse(i64::MIN));
+        Ok(ctx.reduce_next_service_time(i64::MIN))
+    } else {
+        Ok(false)
     }
-    Ok(())
 }
 /// Corresponds to Algorithm 10 found in Section 4.3.
 pub(crate) fn receive_payload_in_place<Crypto: CryptoLayer>(

performance/src/zssp.rs

@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use std::hash::Hash;
 use std::io::Write;
 use std::num::NonZeroU32;
+use std::sync::atomic::{AtomicI64, Ordering};
 use std::sync::{Arc, Mutex, RwLock, Weak};
 
 use arrayvec::ArrayVec;
@@ -44,19 +45,27 @@ impl<Crypto: CryptoLayer> Clone for Context<Crypto> {
 }
 
 pub(crate) type SessionMap<Crypto> = RwLock<HashMap<NonZeroU32, Weak<Session<Crypto>>>>;
+
 pub(crate) type SessionQueue<Crypto> = IndexedBinaryHeap<Weak<Session<Crypto>>, Reverse<i64>>;
+
 pub struct ContextInner<Crypto: CryptoLayer> {
     pub rng: Mutex<Crypto::Rng>,
+    pub next_service_time: AtomicI64,
     pub(crate) s_secret: Crypto::KeyPair,
     /// `session_queue -> state_machine_lock -> state -> session_map`
     pub(crate) session_queue: Mutex<SessionQueue<Crypto>>,
     /// `session_queue -> state_machine_lock -> state -> session_map`
     pub(crate) session_map: SessionMap<Crypto>,
-    pub(crate) unassociated_defrag_cache: Mutex<UnassociatedFragCache<Crypto::IncomingPacketBuffer>>,
+    pub(crate) unassociated_defrag_cache: Mutex<UnassociatedFragCache<Crypto>>,
     pub(crate) unassociated_handshake_states: UnassociatedHandshakeCache<Crypto>,
 
     pub(crate) challenge: ChallengeContext,
 }
+impl<Crypto: CryptoLayer> ContextInner<Crypto> {
+    pub(crate) fn reduce_next_service_time(&self, time: i64) -> bool {
+        self.next_service_time.fetch_min(time, Ordering::Relaxed) > time
+    }
+}
 
 fn parse_fragment_header(incoming_fragment: &[u8]) -> Result<(usize, usize, [u8; AES_GCM_NONCE_SIZE]), ReceiveError> {
     let fragment_no = incoming_fragment[FRAGMENT_NO_IDX] as usize;
@@ -114,6 +123,7 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
         Self(Arc::new(ContextInner {
             rng: Mutex::new(rng),
             s_secret: static_secret_key,
+            next_service_time: AtomicI64::new(i64::MAX),
             session_map: RwLock::new(HashMap::new()),
             challenge,
             session_queue: Mutex::new(IndexedBinaryHeap::new()),
@@ -144,7 +154,7 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
         static_remote_key: Crypto::PublicKey,
         session_data: Crypto::SessionData,
         identity: &[u8],
-    ) -> Result<Arc<Session<Crypto>>, OpenError> {
+    ) -> Result<(Arc<Session<Crypto>>, Option<i64>), OpenError> {
         mtu = mtu.max(MIN_TRANSPORT_MTU);
         if identity.len() > IDENTITY_MAX_SIZE {
             return Err(OpenError::IdentityTooLarge);
@@ -189,7 +199,7 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
         remote_address: &impl Hash,
         mut incoming_fragment_buf: Crypto::IncomingPacketBuffer,
         output_buffer: impl Write,
-    ) -> Result<ReceiveOk<Crypto>, ReceiveError> {
+    ) -> Result<(ReceiveOk<Crypto>, Option<i64>), ReceiveError> {
         use crate::result::FaultType::*;
         let ctx = &self.0;
         send_unassociated_mtu = send_unassociated_mtu.max(MIN_TRANSPORT_MTU);
@@ -315,7 +325,7 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
                     match packet_type {
                         PACKET_TYPE_HANDSHAKE_RESPONSE => {
                             log!(app, ReceivedRawX2);
-                            let should_warn_missing_ratchet = received_x2_trans(
+                            let (should_warn_missing_ratchet, reduced) = received_x2_trans(
                                 &mut app,
                                 ctx,
                                 &session,
@@ -480,17 +490,19 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
 
             let mut buffer = ArrayVec::<u8, HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE>::new();
             let assembled_packet = if fragment_count > 1 {
-                self.0.unassociated_defrag_cache.lock().unwrap().assemble(
+                let next_service_time = self.0.unassociated_defrag_cache.lock().unwrap().assemble(
                     &nonce,
                     remote_address,
                     incoming_fragment.len() - HEADER_SIZE,
                     incoming_fragment_buf,
                     fragment_no,
                     fragment_count,
-                    Crypto::SETTINGS.resend_time as i64,
                     app.time(),
                     &mut fragment_buffer,
                 );
+                if let Some(next_service_time) = next_service_time {
+                    ctx.reduce_next_service_time(next_service_time);
+                }
                 if fragment_buffer.is_empty() {
                     return Ok(ReceiveOk::Unassociated);
                 } else {
@@ -616,13 +628,13 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
         let ctx = &self.0;
         let mut session_queue = ctx.session_queue.lock().unwrap();
         let current_time = app.time();
-        let mut next_service_time = current_time + Crypto::SETTINGS.fragment_assembly_timeout as i64;
+        let mut queue_service_time = i64::MAX;
         // This update system takes heavy advantage of the fact that sessions only need to be updated
         // either roughly every second or roughly every hour. That big gap allows for minor optimizations.
         // If the gap changes (unlikely) this code may need to be rewritten.
         while let Some((session, Reverse(timer), queue_idx)) = session_queue.peek() {
-            if *timer >= current_time {
-                next_service_time = next_service_time.min(*timer);
+            if *timer > current_time {
+                queue_service_time = queue_service_time.min(*timer);
                 break;
             }
             let session = match session.upgrade() {
@@ -639,21 +651,31 @@ impl<Crypto: CryptoLayer> Context<Crypto> {
                 }
             });
             if let Some(next_timer) = result {
-                next_service_time = next_service_time.min(next_timer);
+                queue_service_time = queue_service_time.min(next_timer);
                 session_queue.change_priority(queue_idx, Reverse(next_timer));
             } else {
                 session.expire_inner(Some(ctx), Some(&mut session_queue));
             }
         }
+        // This is the only place where `ctx.next_service_time` can be increased. This only works
+        // correctly because we are holding the `session_queue` lock and we are guaranteed to run
+        // the service code for the other two systems which are not currently locked.
+        ctx.next_service_time.store(queue_service_time, Ordering::Relaxed);
         drop(session_queue);
 
-        self.0
+        let defrag_service_time = self
+            .0
             .unassociated_defrag_cache
             .lock()
             .unwrap()
-            .check_for_expiry(Crypto::SETTINGS.fragment_assembly_timeout as i64, current_time);
-        self.0.unassociated_handshake_states.service(current_time);
+            .check_for_expiry(current_time);
+        let handshake_service_time = self.0.unassociated_handshake_states.service(current_time);
 
-        next_service_time - current_time
+        ctx.next_service_time
+            .fetch_min(defrag_service_time.min(handshake_service_time), Ordering::Relaxed);
+
+        queue_service_time = queue_service_time.min(current_time + Crypto::SETTINGS.fragment_assembly_timeout as i64);
+
+        queue_service_time - current_time
     }
 }