Rework defragmentation, and it now tolerates very poor link quality pretty well.

utils/src/gatherarray.rs

@@ -38,7 +38,7 @@ impl<T, const C: usize> GatherArray<T, C> {
 
     /// Add an item to the array if we don't have this index anymore, returning complete array if all parts are here.
     #[inline(always)]
-    pub fn add(&mut self, index: u8, value: T) -> Option<ArrayVec<T, C>> {
+    pub fn add_return_when_satisfied(&mut self, index: u8, value: T) -> Option<ArrayVec<T, C>> {
         if index < self.goal {
             let mut have = self.have_bits;
             let got = 1u64.wrapping_shl(index as u32);
@@ -91,9 +91,9 @@ mod tests {
         for goal in 2u8..64u8 {
             let mut m = GatherArray::<u8, 64>::new(goal);
             for x in 0..(goal - 1) {
-                assert!(m.add(x, x).is_none());
+                assert!(m.add_return_when_satisfied(x, x).is_none());
             }
-            let r = m.add(goal - 1, goal - 1).unwrap();
+            let r = m.add_return_when_satisfied(goal - 1, goal - 1).unwrap();
             for x in 0..goal {
                 assert_eq!(r.as_ref()[x as usize], x);
             }

zssp/src/applicationlayer.rs

@@ -6,6 +6,8 @@
  * https://www.zerotier.com/
  */
 
+use std::hash::Hash;
+
 use zerotier_crypto::p384::P384KeyPair;
 
 /// Trait to implement to integrate the session into an application.
@@ -65,6 +67,12 @@ pub trait ApplicationLayer: Sized {
     /// for a short period of time when assembling fragmented packets on the receive path.
     type IncomingPacketBuffer: AsRef<[u8]> + AsMut<[u8]>;
 
+    /// Opaque type for whatever constitutes a physical path to the application.
+    ///
+    /// A physical path could be an IP address or IP plus device in the case of UDP, a socket in the
+    /// case of TCP, etc.
+    type PhysicalPath: PartialEq + Eq + Hash + Clone;
+
     /// Get a reference to this host's static public key blob.
     ///
     /// This must contain a NIST P-384 public key but can contain other information. In ZeroTier this

zssp/src/fragged.rs

@@ -0,0 +1,103 @@
+use std::mem::{needs_drop, size_of, zeroed, MaybeUninit};
+use std::ptr::slice_from_raw_parts;
+
+/// Fast packet defragmenter
+pub struct Fragged<Fragment, const MAX_FRAGMENTS: usize> {
+    have: u64,
+    counter: u64,
+    frags: [MaybeUninit<Fragment>; MAX_FRAGMENTS],
+}
+
+pub struct Assembled<Fragment, const MAX_FRAGMENTS: usize>([MaybeUninit<Fragment>; MAX_FRAGMENTS], usize);
+
+impl<Fragment, const MAX_FRAGMENTS: usize> AsRef<[Fragment]> for Assembled<Fragment, MAX_FRAGMENTS> {
+    #[inline(always)]
+    fn as_ref(&self) -> &[Fragment] {
+        unsafe { &*slice_from_raw_parts(self.0.as_ptr().cast::<Fragment>(), self.1) }
+    }
+}
+
+impl<Fragment, const MAX_FRAGMENTS: usize> Drop for Assembled<Fragment, MAX_FRAGMENTS> {
+    #[inline(always)]
+    fn drop(&mut self) {
+        for i in 0..self.1 {
+            unsafe {
+                self.0.get_unchecked_mut(i).assume_init_drop();
+            }
+        }
+    }
+}
+
+impl<Fragment, const MAX_FRAGMENTS: usize> Fragged<Fragment, MAX_FRAGMENTS> {
+    pub fn new() -> Self {
+        debug_assert!(MAX_FRAGMENTS <= 64);
+        debug_assert_eq!(size_of::<MaybeUninit<Fragment>>(), size_of::<Fragment>());
+        debug_assert_eq!(
+            size_of::<[MaybeUninit<Fragment>; MAX_FRAGMENTS]>(),
+            size_of::<[Fragment; MAX_FRAGMENTS]>()
+        );
+        unsafe { zeroed() }
+    }
+
+    pub fn assemble(
+        &mut self,
+        counter: u64,
+        fragment: Fragment,
+        fragment_no: u8,
+        fragment_count: u8,
+    ) -> Option<Assembled<Fragment, MAX_FRAGMENTS>> {
+        if fragment_no < fragment_count && (fragment_count as usize) <= MAX_FRAGMENTS {
+            debug_assert!((fragment_count as usize) <= MAX_FRAGMENTS);
+            debug_assert!((fragment_no as usize) < MAX_FRAGMENTS);
+
+            let mut have = self.have;
+            if counter != self.counter {
+                self.counter = counter;
+                if needs_drop::<Fragment>() {
+                    let mut i = 0;
+                    while have != 0 {
+                        if (have & 1) != 0 {
+                            debug_assert!(i < MAX_FRAGMENTS);
+                            unsafe { self.frags.get_unchecked_mut(i).assume_init_drop() };
+                        }
+                        have = have.wrapping_shr(1);
+                        i += 1;
+                    }
+                } else {
+                    have = 0;
+                }
+            }
+
+            unsafe {
+                self.frags.get_unchecked_mut(fragment_no as usize).write(fragment);
+            }
+
+            let want = 0xffffffffffffffffu64.wrapping_shr((64 - fragment_count) as u32);
+            have |= 1u64.wrapping_shl(fragment_no as u32);
+            if (have & want) == want {
+                self.have = 0;
+                return Some(Assembled(unsafe { std::mem::transmute_copy(&self.frags) }, fragment_count as usize));
+            } else {
+                self.have = have;
+            }
+        }
+        return None;
+    }
+}
+
+impl<Fragment, const MAX_FRAGMENTS: usize> Drop for Fragged<Fragment, MAX_FRAGMENTS> {
+    fn drop(&mut self) {
+        if needs_drop::<Fragment>() {
+            let mut have = self.have;
+            let mut i = 0;
+            while have != 0 {
+                if (have & 1) != 0 {
+                    debug_assert!(i < MAX_FRAGMENTS);
+                    unsafe { self.frags.get_unchecked_mut(i).assume_init_drop() };
+                }
+                have = have.wrapping_shr(1);
+                i += 1;
+            }
+        }
+    }
+}

zssp/src/lib.rs

@@ -8,6 +8,7 @@
 
 mod applicationlayer;
 mod error;
+mod fragged;
 mod proto;
 mod sessionid;
 mod zssp;

zssp/src/main.rs

@@ -1,9 +1,12 @@
+use std::iter::ExactSizeIterator;
+use std::str::FromStr;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::mpsc;
 use std::thread;
 use std::time::Duration;
 
 use zerotier_crypto::p384::{P384KeyPair, P384PublicKey};
+use zerotier_crypto::random;
 use zerotier_crypto::secret::Secret;
 use zerotier_utils::hex;
 use zerotier_utils::ms_monotonic;
@@ -23,8 +26,8 @@ impl zssp::ApplicationLayer for TestApplication {
     const RETRY_INTERVAL: i64 = 500;
 
     type Data = ();
-
     type IncomingPacketBuffer = Vec<u8>;
+    type PhysicalPath = usize;
 
     fn get_local_s_public_blob(&self) -> &[u8] {
         self.identity_key.public_key_bytes()
@@ -37,6 +40,7 @@ impl zssp::ApplicationLayer for TestApplication {
 
 fn alice_main(
     run: &AtomicBool,
+    packet_success_rate: u32,
     alice_app: &TestApplication,
     bob_app: &TestApplication,
     alice_out: mpsc::SyncSender<Vec<u8>>,
@@ -46,7 +50,7 @@ fn alice_main(
     let mut data_buf = [0u8; 65536];
     let mut next_service = ms_monotonic() + 500;
     let mut last_ratchet_count = 0;
-    let test_data = [1u8; 10000];
+    let test_data = [1u8; TEST_MTU * 10];
     let mut up = false;
 
     let alice_session = context
@@ -71,7 +75,7 @@ fn alice_main(
         loop {
             let pkt = alice_in.try_recv();
             if let Ok(pkt) = pkt {
-                //println!("bob >> alice {}", pkt.len());
+                if (random::xorshift64_random() as u32) <= packet_success_rate {
                     match context.receive(
                         alice_app,
                         || true,
@@ -79,6 +83,7 @@ fn alice_main(
                         |_, b| {
                             let _ = alice_out.send(b.to_vec());
                         },
+                        &0,
                         &mut data_buf,
                         pkt,
                         TEST_MTU,
@@ -96,6 +101,9 @@ fn alice_main(
                         Ok(zssp::ReceiveResult::Rejected) => {}
                         Err(e) => {
                             println!("[alice] ERROR {}", e.to_string());
+                            //run.store(false, Ordering::SeqCst);
+                            //break;
+                        }
                     }
                 }
             } else {
@@ -116,12 +124,14 @@ fn alice_main(
                         let _ = alice_out.send(b.to_vec());
                     },
                     &mut data_buf[..TEST_MTU],
-                    &test_data[..1400 + ((zerotier_crypto::random::xorshift64_random() as usize) % (test_data.len() - 1400))],
+                    &test_data[..1400 + ((random::xorshift64_random() as usize) % (test_data.len() - 1400))],
                 )
                 .is_ok());
         } else {
             if alice_session.established() {
                 up = true;
+            } else {
+                thread::sleep(Duration::from_millis(10));
             }
         }
 
@@ -140,6 +150,7 @@ fn alice_main(
 
 fn bob_main(
     run: &AtomicBool,
+    packet_success_rate: u32,
     _alice_app: &TestApplication,
     bob_app: &TestApplication,
     bob_out: mpsc::SyncSender<Vec<u8>>,
@@ -160,7 +171,7 @@ fn bob_main(
         let current_time = ms_monotonic();
 
         if let Ok(pkt) = pkt {
-            //println!("alice >> bob {}", pkt.len());
+            if (random::xorshift64_random() as u32) <= packet_success_rate {
                 match context.receive(
                     bob_app,
                     || true,
@@ -168,6 +179,7 @@ fn bob_main(
                     |_, b| {
                         let _ = bob_out.send(b.to_vec());
                     },
+                    &0,
                     &mut data_buf,
                     pkt,
                     TEST_MTU,
@@ -196,6 +208,9 @@ fn bob_main(
                     Ok(zssp::ReceiveResult::Rejected) => {}
                     Err(e) => {
                         println!("[bob] ERROR {}", e.to_string());
+                        //run.store(false, Ordering::SeqCst);
+                        //break;
+                    }
                 }
             }
         }
@@ -211,10 +226,12 @@ fn bob_main(
         let speed_metric_elapsed = current_time - last_speed_metric;
         if speed_metric_elapsed >= 1000 {
             last_speed_metric = current_time;
+            if transferred > 0 {
                 println!(
                     "[bob] throughput: {} MiB/sec (combined input and output)",
                     ((transferred as f64) / 1048576.0) / ((speed_metric_elapsed as f64) / 1000.0)
                 );
+            }
             transferred = 0;
         }
 
@@ -240,9 +257,16 @@ fn main() {
     let (alice_out, bob_in) = mpsc::sync_channel::<Vec<u8>>(1024);
     let (bob_out, alice_in) = mpsc::sync_channel::<Vec<u8>>(1024);
 
+    let args = std::env::args();
+    let packet_success_rate = if args.len() <= 1 {
+        u32::MAX
+    } else {
+        ((u32::MAX as f64) * f64::from_str(args.last().unwrap().as_str()).unwrap()) as u32
+    };
+
     thread::scope(|ts| {
-        let alice_thread = ts.spawn(|| alice_main(&run, &alice_app, &bob_app, alice_out, alice_in));
+        let alice_thread = ts.spawn(|| alice_main(&run, packet_success_rate, &alice_app, &bob_app, alice_out, alice_in));
-        let bob_thread = ts.spawn(|| bob_main(&run, &alice_app, &bob_app, bob_out, bob_in));
+        let bob_thread = ts.spawn(|| bob_main(&run, packet_success_rate, &alice_app, &bob_app, bob_out, bob_in));
 
         thread::sleep(Duration::from_secs(60 * 10));
 

zssp/src/proto.rs

@@ -24,24 +24,29 @@ pub const MIN_TRANSPORT_MTU: usize = 128;
 /// Maximum combined size of static public blob and metadata.
 pub const MAX_INIT_PAYLOAD_SIZE: usize = MAX_NOISE_HANDSHAKE_SIZE - ALICE_NOISE_XK_ACK_MIN_SIZE;
 
+/// Version 0: Noise_XK with NIST P-384 plus Kyber1024 hybrid exchange on session init.
 pub(crate) const SESSION_PROTOCOL_VERSION: u8 = 0x00;
 
-pub(crate) const COUNTER_WINDOW_MAX_OOO: usize = 16;
+/// Maximum window over which packets may be reordered.
+pub(crate) const COUNTER_WINDOW_MAX_OOO: usize = 32;
+
+/// Maximum number of counter steps that the counter is allowed to skip ahead.
 pub(crate) const COUNTER_WINDOW_MAX_SKIP_AHEAD: u64 = 16777216;
 
-pub(crate) const PACKET_TYPE_DATA: u8 = 0;
+pub(crate) const PACKET_TYPE_NOP: u8 = 0;
-pub(crate) const PACKET_TYPE_ALICE_NOISE_XK_INIT: u8 = 1;
+pub(crate) const PACKET_TYPE_DATA: u8 = 1;
-pub(crate) const PACKET_TYPE_BOB_NOISE_XK_ACK: u8 = 2;
+pub(crate) const PACKET_TYPE_ALICE_NOISE_XK_INIT: u8 = 2;
-pub(crate) const PACKET_TYPE_ALICE_NOISE_XK_ACK: u8 = 3;
+pub(crate) const PACKET_TYPE_BOB_NOISE_XK_ACK: u8 = 3;
-pub(crate) const PACKET_TYPE_REKEY_INIT: u8 = 4;
+pub(crate) const PACKET_TYPE_ALICE_NOISE_XK_ACK: u8 = 4;
-pub(crate) const PACKET_TYPE_REKEY_ACK: u8 = 5;
+pub(crate) const PACKET_TYPE_REKEY_INIT: u8 = 5;
+pub(crate) const PACKET_TYPE_REKEY_ACK: u8 = 6;
 
 pub(crate) const HEADER_SIZE: usize = 16;
 pub(crate) const HEADER_PROTECT_ENCRYPT_START: usize = 6;
 pub(crate) const HEADER_PROTECT_ENCRYPT_END: usize = 22;
 
-pub(crate) const KBKDF_KEY_USAGE_LABEL_KEX_ENCRYPTION: u8 = b'X'; // intermediate keys used in key exchanges
+pub(crate) const KBKDF_KEY_USAGE_LABEL_INIT_ENCRYPTION: u8 = b'x'; // AES-CTR encryption during initial setup
-pub(crate) const KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION: u8 = b'x'; // intermediate keys used in key exchanges
+pub(crate) const KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION: u8 = b'X'; // HMAC-SHA384 during initial setup
 pub(crate) const KBKDF_KEY_USAGE_LABEL_AES_GCM_ALICE_TO_BOB: u8 = b'A'; // AES-GCM in A->B direction
 pub(crate) const KBKDF_KEY_USAGE_LABEL_AES_GCM_BOB_TO_ALICE: u8 = b'B'; // AES-GCM in B->A direction
 pub(crate) const KBKDF_KEY_USAGE_LABEL_RATCHET: u8 = b'R'; // Key used in derivatin of next session key
@@ -50,6 +55,7 @@ pub(crate) const MAX_FRAGMENTS: usize = 48; // hard protocol max: 63
 pub(crate) const MAX_NOISE_HANDSHAKE_FRAGMENTS: usize = 16; // enough room for p384 + ZT identity + kyber1024 + tag/hmac/etc.
 pub(crate) const MAX_NOISE_HANDSHAKE_SIZE: usize = MAX_NOISE_HANDSHAKE_FRAGMENTS * MIN_TRANSPORT_MTU;
 
+/// Size of keys used during derivation, mixing, etc. process.
 pub(crate) const BASE_KEY_SIZE: usize = 64;
 
 pub(crate) const AES_256_KEY_SIZE: usize = 32;
@@ -162,7 +168,6 @@ impl RekeyAck {
 pub(crate) trait ProtocolFlatBuffer {}
 impl ProtocolFlatBuffer for AliceNoiseXKInit {}
 impl ProtocolFlatBuffer for BobNoiseXKAck {}
-//impl ProtocolFlatBuffer for NoiseXKAliceStaticAck {}
 impl ProtocolFlatBuffer for RekeyInit {}
 impl ProtocolFlatBuffer for RekeyAck {}
 

zssp/src/zssp.rs

@@ -15,19 +15,18 @@ use std::sync::atomic::{AtomicI64, AtomicU64, Ordering};
 use std::sync::{Arc, Mutex, RwLock, Weak};
 
 use zerotier_crypto::aes::{Aes, AesCtr, AesGcm};
-use zerotier_crypto::hash::{hmac_sha512, HMACSHA384, HMAC_SHA384_SIZE, SHA384, SHA384_HASH_SIZE};
+use zerotier_crypto::hash::{hmac_sha512, HMACSHA384, HMAC_SHA384_SIZE, SHA384};
 use zerotier_crypto::p384::{P384KeyPair, P384PublicKey, P384_ECDH_SHARED_SECRET_SIZE};
 use zerotier_crypto::secret::Secret;
 use zerotier_crypto::{random, secure_eq};
 
 use zerotier_utils::arrayvec::ArrayVec;
-use zerotier_utils::gatherarray::GatherArray;
-use zerotier_utils::ringbuffermap::RingBufferMap;
 
-use pqc_kyber::{KYBER_CIPHERTEXTBYTES, KYBER_SECRETKEYBYTES, KYBER_SSBYTES};
+use pqc_kyber::{KYBER_SECRETKEYBYTES, KYBER_SSBYTES};
 
 use crate::applicationlayer::ApplicationLayer;
 use crate::error::Error;
+use crate::fragged::Fragged;
 use crate::proto::*;
 use crate::sessionid::SessionId;
 
@@ -37,7 +36,12 @@ use crate::sessionid::SessionId;
 /// defragment incoming packets that are not yet associated with a session.
 pub struct Context<Application: ApplicationLayer> {
     max_incomplete_session_queue_size: usize,
-    defrag: Mutex<RingBufferMap<u64, GatherArray<Application::IncomingPacketBuffer, MAX_NOISE_HANDSHAKE_FRAGMENTS>, 256, 256>>,
+    defrag: Mutex<
+        HashMap<
+            (Application::PhysicalPath, u64),
+            Arc<Mutex<(Fragged<Application::IncomingPacketBuffer, MAX_NOISE_HANDSHAKE_FRAGMENTS>, i64)>>,
+        >,
+    >,
     sessions: RwLock<SessionsById<Application>>,
 }
 
@@ -80,7 +84,7 @@ pub struct Session<Application: ApplicationLayer> {
     receive_window: [AtomicU64; COUNTER_WINDOW_MAX_OOO],
     header_protection_cipher: Aes,
     state: RwLock<State>,
-    defrag: Mutex<RingBufferMap<u64, GatherArray<Application::IncomingPacketBuffer, MAX_FRAGMENTS>, 16, 16>>,
+    defrag: [Mutex<Fragged<Application::IncomingPacketBuffer, MAX_FRAGMENTS>>; COUNTER_WINDOW_MAX_OOO],
 }
 
 /// Most of the mutable parts of a session state.
@@ -91,20 +95,16 @@ struct State {
     current_offer: Offer,
 }
 
-/// State related to an incoming session not yet fully established.
 struct IncomingIncompleteSession {
     timestamp: i64,
-    request_hash: [u8; SHA384_HASH_SIZE],
     alice_session_id: SessionId,
     bob_session_id: SessionId,
     noise_es_ee: Secret<BASE_KEY_SIZE>,
-    bob_hk_ciphertext: [u8; KYBER_CIPHERTEXTBYTES],
     hk: Secret<KYBER_SSBYTES>,
     header_protection_key: Secret<AES_HEADER_PROTECTION_KEY_SIZE>,
     bob_noise_e_secret: P384KeyPair,
 }
 
-/// State related to an outgoing session attempt.
 struct OutgoingSessionInit {
     last_retry_time: AtomicI64,
     alice_noise_e_secret: P384KeyPair,
@@ -114,14 +114,19 @@ struct OutgoingSessionInit {
     init_packet: [u8; AliceNoiseXKInit::SIZE],
 }
 
-/// Latest outgoing offer, either an outgoing attempt or a rekey attempt.
+struct OutgoingSessionAck {
+    last_retry_time: AtomicI64,
+    ack: [u8; MAX_NOISE_HANDSHAKE_SIZE],
+    ack_size: usize,
+}
+
 enum Offer {
     None,
     NoiseXKInit(Box<OutgoingSessionInit>),
-    RekeyInit(P384KeyPair, [u8; RekeyInit::SIZE], AtomicI64),
+    NoiseXKAck(Box<OutgoingSessionAck>),
+    RekeyInit(P384KeyPair, i64),
 }
 
-/// An ephemeral session key with expiration info.
 struct SessionKey {
     ratchet_key: Secret<BASE_KEY_SIZE>,           // Key used in derivation of the next session key
     receive_key: Secret<AES_256_KEY_SIZE>,        // Receive side AES-GCM key
@@ -134,6 +139,7 @@ struct SessionKey {
     expire_at_counter: u64,                       // Hard error when this counter value is reached or exceeded
     ratchet_count: u64,                           // Number of rekey events
     bob: bool,                                    // Was this side "Bob" in this exchange?
+    confirmed: bool,                              // Is this key confirmed by the other side?
 }
 
 impl<Application: ApplicationLayer> Context<Application> {
@@ -141,7 +147,7 @@ impl<Application: ApplicationLayer> Context<Application> {
     pub fn new(max_incomplete_session_queue_size: usize) -> Self {
         Self {
             max_incomplete_session_queue_size,
-            defrag: Mutex::new(RingBufferMap::new(random::next_u32_secure())),
+            defrag: Mutex::new(HashMap::new()),
             sessions: RwLock::new(SessionsById {
                 active: HashMap::with_capacity(64),
                 incoming: HashMap::with_capacity(64),
@@ -173,19 +179,14 @@ impl<Application: ApplicationLayer> Context<Application> {
             for (id, s) in sessions.active.iter() {
                 if let Some(session) = s.upgrade() {
                     let state = session.state.read().unwrap();
-                    match &state.current_offer {
+                    if match &state.current_offer {
-                        Offer::None => {
+                        Offer::None => true,
-                            if let Some(key) = state.keys[state.current_key].as_ref() {
-                                if key.bob
-                                    && (current_time >= key.rekey_at_time
-                                        || session.send_counter.load(Ordering::Relaxed) >= key.rekey_at_counter)
-                                {
-                                    drop(state);
-                                    session.initiate_rekey(|b| send(&session, b), current_time);
-                                }
-                            }
-                        }
                         Offer::NoiseXKInit(offer) => {
+                            // If there's an outstanding attempt to open a session, retransmit this periodically
+                            // in case the initial packet doesn't make it. Note that we currently don't have
+                            // retransmission for the intermediate steps, so a new session may still fail if the
+                            // packet loss rate is huge. The application layer has its own logic to keep trying
+                            // under those conditions.
                             if offer.last_retry_time.load(Ordering::Relaxed) < retry_cutoff {
                                 offer.last_retry_time.store(current_time, Ordering::Relaxed);
                                 let _ = send_with_fragmentation(
@@ -199,11 +200,37 @@ impl<Application: ApplicationLayer> Context<Application> {
                                     None,
                                 );
                             }
+                            false
                         }
-                        Offer::RekeyInit(_, rekey_packet, last_retry_time) => {
+                        Offer::NoiseXKAck(ack) => {
-                            if last_retry_time.load(Ordering::Relaxed) < retry_cutoff {
+                            // We also keep retransmitting the final ACK until we get a valid DATA or NOP packet
-                                last_retry_time.store(current_time, Ordering::Relaxed);
+                            // from Bob, otherwise we could get a half open session.
-                                send(&session, &mut (rekey_packet.clone()));
+                            if ack.last_retry_time.load(Ordering::Relaxed) < retry_cutoff {
+                                ack.last_retry_time.store(current_time, Ordering::Relaxed);
+                                let _ = send_with_fragmentation(
+                                    |b| send(&session, b),
+                                    &mut (ack.ack.clone())[..ack.ack_size],
+                                    mtu,
+                                    PACKET_TYPE_ALICE_NOISE_XK_ACK,
+                                    state.remote_session_id,
+                                    0,
+                                    2,
+                                    Some(&session.header_protection_cipher),
+                                );
+                            }
+                            false
+                        }
+                        Offer::RekeyInit(_, last_rekey_attempt_time) => *last_rekey_attempt_time < retry_cutoff,
+                    } {
+                        // Check whether we need to rekey if there is no pending offer or if the last rekey
+                        // offer was before retry_cutoff (checked in the 'match' above).
+                        if let Some(key) = state.keys[state.current_key].as_ref() {
+                            if key.bob
+                                && (current_time >= key.rekey_at_time
+                                    || session.send_counter.load(Ordering::Relaxed) >= key.rekey_at_counter)
+                            {
+                                drop(state);
+                                session.initiate_rekey(|b| send(&session, b), current_time);
                             }
                         }
                     }
@@ -297,7 +324,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                         init_packet: [0u8; AliceNoiseXKInit::SIZE],
                     })),
                 }),
-                defrag: Mutex::new(RingBufferMap::new(random::xorshift64_random() as u32)),
+                defrag: std::array::from_fn(|_| Mutex::new(Fragged::new())),
             });
 
             sessions.active.insert(local_session_id, Arc::downgrade(&session));
@@ -310,7 +337,7 @@ impl<Application: ApplicationLayer> Context<Application> {
             let init_packet = if let Offer::NoiseXKInit(offer) = &mut state.current_offer {
                 &mut offer.init_packet
             } else {
-                panic!();
+                panic!(); // should be impossible
             };
 
             let init: &mut AliceNoiseXKInit = byte_array_as_proto_buffer_mut(init_packet).unwrap();
@@ -321,12 +348,12 @@ impl<Application: ApplicationLayer> Context<Application> {
             init.header_protection_key = header_protection_key.0;
 
             aes_ctr_crypt_one_time_use_key(
-                kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_ENCRYPTION>(noise_es.as_bytes()).as_bytes(),
+                kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_ENCRYPTION>(noise_es.as_bytes()).as_bytes(),
                 &mut init_packet[AliceNoiseXKInit::ENC_START..AliceNoiseXKInit::AUTH_START],
             );
 
             let hmac = hmac_sha384_2(
-                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(noise_es.as_bytes()).as_bytes(),
+                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(noise_es.as_bytes()).as_bytes(),
                 &create_message_nonce(PACKET_TYPE_ALICE_NOISE_XK_INIT, 1),
                 &init_packet[HEADER_SIZE..AliceNoiseXKInit::AUTH_START],
             );
@@ -385,6 +412,7 @@ impl<Application: ApplicationLayer> Context<Application> {
         mut check_allow_incoming_session: CheckAllowIncomingSession,
         mut check_accept_session: CheckAcceptSession,
         mut send: SendFunction,
+        source: &Application::PhysicalPath,
         data_buf: &'b mut [u8],
         mut incoming_packet_buf: Application::IncomingPacketBuffer,
         mtu: usize,
@@ -414,11 +442,10 @@ impl<Application: ApplicationLayer> Context<Application> {
 
                 if session.check_receive_window(incoming_counter) {
                     if fragment_count > 1 {
-                        if fragment_count <= (MAX_FRAGMENTS as u8) && fragment_no < fragment_count {
+                        let mut fragged = session.defrag[(incoming_counter as usize) % COUNTER_WINDOW_MAX_OOO].lock().unwrap();
-                            let mut defrag = session.defrag.lock().unwrap();
+                        if let Some(assembled_packet) = fragged.assemble(incoming_counter, incoming_packet_buf, fragment_no, fragment_count)
-                            let fragment_gather_array = defrag.get_or_create_mut(&incoming_counter, || GatherArray::new(fragment_count));
+                        {
-                            if let Some(assembled_packet) = fragment_gather_array.add(fragment_no, incoming_packet_buf) {
+                            drop(fragged);
-                                drop(defrag); // release lock
                             return self.process_complete_incoming_packet(
                                 app,
                                 &mut send,
@@ -437,9 +464,6 @@ impl<Application: ApplicationLayer> Context<Application> {
                         } else {
                             return Ok(ReceiveResult::Ok);
                         }
-                        } else {
-                            return Err(Error::InvalidPacket);
-                        }
                     } else {
                         return self.process_complete_incoming_packet(
                             app,
@@ -476,10 +500,19 @@ impl<Application: ApplicationLayer> Context<Application> {
 
         let (key_index, packet_type, fragment_count, fragment_no, incoming_counter) = parse_packet_header(&incoming_packet);
         if fragment_count > 1 {
+            let fragged_m = {
                 let mut defrag = self.defrag.lock().unwrap();
-            let fragment_gather_array = defrag.get_or_create_mut(&incoming_counter, || GatherArray::new(fragment_count));
+                defrag
-            if let Some(assembled_packet) = fragment_gather_array.add(fragment_no, incoming_packet_buf) {
+                    .entry((source.clone(), incoming_counter))
-                drop(defrag); // release lock
+                    .or_insert_with(|| Arc::new(Mutex::new((Fragged::new(), current_time))))
+                    .clone()
+            };
+            let mut fragged = fragged_m.lock().unwrap();
+            if let Some(assembled_packet) = fragged
+                .0
+                .assemble(incoming_counter, incoming_packet_buf, fragment_no, fragment_count)
+            {
+                self.defrag.lock().unwrap().remove(&(source.clone(), incoming_counter));
                 return self.process_complete_incoming_packet(
                     app,
                     &mut send,
@@ -543,7 +576,7 @@ impl<Application: ApplicationLayer> Context<Application> {
         // Generate incoming message nonce for decryption and authentication.
         let incoming_message_nonce = create_message_nonce(packet_type, incoming_counter);
 
-        if packet_type == PACKET_TYPE_DATA {
+        if packet_type <= PACKET_TYPE_DATA {
             if let Some(session) = session {
                 let state = session.state.read().unwrap();
                 if let Some(key) = state.keys[key_index].as_ref() {
@@ -590,22 +623,41 @@ impl<Application: ApplicationLayer> Context<Application> {
                             // Update the current key to point to this key if it's newer, since having received
                             // a packet encrypted with it proves that the other side has successfully derived it
                             // as well.
-                            if state.current_key == key_index {
+                            if state.current_key == key_index && key.confirmed {
                                 drop(state);
                             } else {
-                                let key_created_at_counter = key.created_at_counter;
+                                let current_key_created_at_counter = key.created_at_counter;
+
                                 drop(state);
                                 let mut state = session.state.write().unwrap();
+
+                                if state.current_key != key_index {
                                     if let Some(other_session_key) = state.keys[state.current_key].as_ref() {
-                                    if other_session_key.created_at_counter < key_created_at_counter {
+                                        if other_session_key.created_at_counter < current_key_created_at_counter {
                                             state.current_key = key_index;
                                         }
                                     } else {
                                         state.current_key = key_index;
                                     }
+                                } else {
+                                    state.keys[key_index].as_mut().unwrap().confirmed = true;
+                                }
+
+                                // If we got a valid data packet from Bob, this means we can cancel any offers
+                                // that are still oustanding for initialization.
+                                match &state.current_offer {
+                                    Offer::NoiseXKInit(_) | Offer::NoiseXKAck(_) => {
+                                        state.current_offer = Offer::None;
+                                    }
+                                    _ => {}
+                                }
                             }
 
+                            if packet_type == PACKET_TYPE_DATA {
                                 return Ok(ReceiveResult::OkData(session, &mut data_buf[..data_len]));
+                            } else {
+                                println!("nop");
+                            }
                         } else {
                             return Err(Error::OutOfSequence);
                         }
@@ -650,26 +702,10 @@ impl<Application: ApplicationLayer> Context<Application> {
                     if incoming_counter != 1 || session.is_some() {
                         return Err(Error::OutOfSequence);
                     }
-
+                    if pkt_assembled.len() != AliceNoiseXKInit::SIZE {
-                    // Hash the init packet so we can check to see if it's just being retransmitted. Alice may
+                        return Err(Error::InvalidPacket);
-                    // attempt to retransmit this packet until she receives a response.
-                    let request_hash = SHA384::hash(&pkt_assembled);
-
-                    let (alice_session_id, mut bob_session_id, noise_es_ee, bob_hk_ciphertext, header_protection_key, bob_noise_e);
-                    if let Some(incoming) = incoming {
-                        // If we've already seen this exact packet before, just recall the same state so we send the
-                        // same response.
-                        if secure_eq(&request_hash, &incoming.request_hash) {
-                            alice_session_id = incoming.alice_session_id;
-                            bob_session_id = incoming.bob_session_id;
-                            noise_es_ee = incoming.noise_es_ee.clone();
-                            bob_hk_ciphertext = incoming.bob_hk_ciphertext;
-                            header_protection_key = incoming.header_protection_key.clone();
-                            bob_noise_e = *incoming.bob_noise_e_secret.public_key_bytes();
-                        } else {
-                            return Err(Error::FailedAuthentication);
                     }
-                    } else {
+
                     // Otherwise parse the packet, authenticate, generate keys, etc. and record state in an
                     // incoming state object until this phase of the negotiation is done.
                     let pkt: &AliceNoiseXKInit = byte_array_as_proto_buffer(pkt_assembled)?;
@@ -680,7 +716,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                     if !secure_eq(
                         &pkt.hmac_es,
                         &hmac_sha384_2(
-                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(noise_es.as_bytes()).as_bytes(),
+                            kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(noise_es.as_bytes()).as_bytes(),
                             &incoming_message_nonce,
                             &pkt_assembled[HEADER_SIZE..AliceNoiseXKInit::AUTH_START],
                         ),
@@ -695,32 +731,32 @@ impl<Application: ApplicationLayer> Context<Application> {
 
                     // Decrypt encrypted part of payload.
                     aes_ctr_crypt_one_time_use_key(
-                            kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_ENCRYPTION>(noise_es.as_bytes()).as_bytes(),
+                        kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_ENCRYPTION>(noise_es.as_bytes()).as_bytes(),
                         &mut pkt_assembled[AliceNoiseXKInit::ENC_START..AliceNoiseXKInit::AUTH_START],
                     );
 
                     let pkt: &AliceNoiseXKInit = byte_array_as_proto_buffer(pkt_assembled)?;
-                        alice_session_id = SessionId::new_from_bytes(&pkt.alice_session_id).ok_or(Error::InvalidPacket)?;
+                    let alice_session_id = SessionId::new_from_bytes(&pkt.alice_session_id).ok_or(Error::InvalidPacket)?;
-                        header_protection_key = Secret(pkt.header_protection_key);
+                    let header_protection_key = Secret(pkt.header_protection_key);
 
                     // Create Bob's ephemeral keys and derive noise_es_ee by agreeing with Alice's. Also create
                     // a Kyber ciphertext to send back to Alice.
                     let bob_noise_e_secret = P384KeyPair::generate();
-                        bob_noise_e = bob_noise_e_secret.public_key_bytes().clone();
+                    let bob_noise_e = bob_noise_e_secret.public_key_bytes().clone();
-                        noise_es_ee = Secret(hmac_sha512(
+                    let noise_es_ee = Secret(hmac_sha512(
                         noise_es.as_bytes(),
                         bob_noise_e_secret
                             .agree(&alice_noise_e)
                             .ok_or(Error::FailedAuthentication)?
                             .as_bytes(),
                     ));
-                        let (hk_ct, hk) = pqc_kyber::encapsulate(&pkt.alice_hk_public, &mut random::SecureRandom::default())
+                    let (bob_hk_ciphertext, hk) = pqc_kyber::encapsulate(&pkt.alice_hk_public, &mut random::SecureRandom::default())
                         .map_err(|_| Error::FailedAuthentication)
                         .map(|(ct, hk)| (ct, Secret(hk)))?;
-                        bob_hk_ciphertext = hk_ct;
 
                     let mut sessions = self.sessions.write().unwrap();
 
+                    let mut bob_session_id;
                     loop {
                         bob_session_id = SessionId::random();
                         if !sessions.active.contains_key(&bob_session_id) && !sessions.incoming.contains_key(&bob_session_id) {
@@ -752,17 +788,17 @@ impl<Application: ApplicationLayer> Context<Application> {
                         bob_session_id,
                         Arc::new(IncomingIncompleteSession {
                             timestamp: current_time,
-                                request_hash,
                             alice_session_id,
                             bob_session_id,
                             noise_es_ee: noise_es_ee.clone(),
-                                bob_hk_ciphertext,
                             hk,
                             bob_noise_e_secret,
                             header_protection_key: Secret(pkt.header_protection_key),
                         }),
                     );
-                    }
+                    debug_assert!(!sessions.active.contains_key(&bob_session_id));
+
+                    drop(sessions);
 
                     // Create Bob's ephemeral counter-offer reply.
                     let mut ack_packet = [0u8; BobNoiseXKAck::SIZE];
@@ -774,13 +810,13 @@ impl<Application: ApplicationLayer> Context<Application> {
 
                     // Encrypt main section of reply.
                     aes_ctr_crypt_one_time_use_key(
-                        kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_ENCRYPTION>(noise_es_ee.as_bytes()).as_bytes(),
+                        kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_ENCRYPTION>(noise_es_ee.as_bytes()).as_bytes(),
                         &mut ack_packet[BobNoiseXKAck::ENC_START..BobNoiseXKAck::AUTH_START],
                     );
 
                     // Add HMAC-SHA384 to reply packet.
                     let reply_hmac = hmac_sha384_2(
-                        kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(noise_es_ee.as_bytes()).as_bytes(),
+                        kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(noise_es_ee.as_bytes()).as_bytes(),
                         &create_message_nonce(PACKET_TYPE_BOB_NOISE_XK_ACK, 1),
                         &ack_packet[HEADER_SIZE..BobNoiseXKAck::AUTH_START],
                     );
@@ -812,9 +848,18 @@ impl<Application: ApplicationLayer> Context<Application> {
                     if incoming_counter != 1 || incoming.is_some() {
                         return Err(Error::OutOfSequence);
                     }
+                    if pkt_assembled.len() != BobNoiseXKAck::SIZE {
+                        return Err(Error::InvalidPacket);
+                    }
 
                     if let Some(session) = session {
                         let state = session.state.read().unwrap();
+
+                        // This doesn't make sense if the session is up.
+                        if state.keys[state.current_key].is_some() {
+                            return Err(Error::OutOfSequence);
+                        }
+
                         if let Offer::NoiseXKInit(outgoing_offer) = &state.current_offer {
                             let pkt: &BobNoiseXKAck = byte_array_as_proto_buffer(pkt_assembled)?;
 
@@ -830,7 +875,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                             ));
 
                             let noise_es_ee_kex_hmac_key =
-                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(noise_es_ee.as_bytes());
+                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(noise_es_ee.as_bytes());
 
                             // Authenticate Bob's reply and the validity of bob_noise_e.
                             if !secure_eq(
@@ -846,7 +891,7 @@ impl<Application: ApplicationLayer> Context<Application> {
 
                             // Decrypt encrypted portion of message.
                             aes_ctr_crypt_one_time_use_key(
-                                kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_ENCRYPTION>(noise_es_ee.as_bytes()).as_bytes(),
+                                kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_ENCRYPTION>(noise_es_ee.as_bytes()).as_bytes(),
                                 &mut pkt_assembled[BobNoiseXKAck::ENC_START..BobNoiseXKAck::AUTH_START],
                             );
                             let pkt: &BobNoiseXKAck = byte_array_as_proto_buffer(pkt_assembled)?;
@@ -915,7 +960,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                                 // key exchange. Bob won't be able to do this until he decrypts and parses Alice's
                                 // identity, so the first HMAC is to let him authenticate that first.
                                 let hmac_es_ee_se_hk_psk = hmac_sha384_2(
-                                    kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(noise_es_ee_se_hk_psk.as_bytes())
+                                    kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(noise_es_ee_se_hk_psk.as_bytes())
                                         .as_bytes(),
                                     &reply_message_nonce,
                                     &reply_buffer[HEADER_SIZE..reply_len],
@@ -933,9 +978,15 @@ impl<Application: ApplicationLayer> Context<Application> {
                                         current_time,
                                         2,
                                         false,
+                                        false,
                                     ));
+                                    debug_assert!(state.keys[1].is_none());
                                     state.current_key = 0;
-                                    state.current_offer = Offer::None;
+                                    state.current_offer = Offer::NoiseXKAck(Box::new(OutgoingSessionAck {
+                                        last_retry_time: AtomicI64::new(current_time),
+                                        ack: reply_buffer,
+                                        ack_size: reply_len,
+                                    }));
                                 }
 
                                 send_with_fragmentation(
@@ -980,18 +1031,13 @@ impl<Application: ApplicationLayer> Context<Application> {
                     }
 
                     if let Some(incoming) = incoming {
-                        // Check timeout, negotiations aren't allowed to take longer than this.
-                        if (current_time - incoming.timestamp) > Application::INCOMING_SESSION_NEGOTIATION_TIMEOUT_MS {
-                            return Err(Error::UnknownLocalSessionId);
-                        }
-
                         // Check the first HMAC to verify against the currently known noise_es_ee key, which verifies
                         // that this reply is part of this session.
                         let auth_start = pkt_assembled.len() - ALICE_NOISE_XK_ACK_AUTH_SIZE;
                         if !secure_eq(
                             &pkt_assembled[auth_start..pkt_assembled.len() - HMAC_SHA384_SIZE],
                             &hmac_sha384_2(
-                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(incoming.noise_es_ee.as_bytes())
+                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(incoming.noise_es_ee.as_bytes())
                                     .as_bytes(),
                                 &incoming_message_nonce,
                                 &pkt_assembled[HEADER_SIZE..auth_start],
@@ -1065,7 +1111,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                         if !secure_eq(
                             &pkt_assembly_buffer_copy[auth_start + HMAC_SHA384_SIZE..pkt_assembled.len()],
                             &hmac_sha384_2(
-                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_KEX_AUTHENTICATION>(noise_es_ee_se_hk_psk.as_bytes())
+                                kbkdf::<HMAC_SHA384_SIZE, KBKDF_KEY_USAGE_LABEL_INIT_AUTHENTICATION>(noise_es_ee_se_hk_psk.as_bytes())
                                     .as_bytes(),
                                 &incoming_message_nonce,
                                 &pkt_assembly_buffer_copy[HEADER_SIZE..auth_start + HMAC_SHA384_SIZE],
@@ -1084,22 +1130,30 @@ impl<Application: ApplicationLayer> Context<Application> {
                             state: RwLock::new(State {
                                 remote_session_id: Some(incoming.alice_session_id),
                                 keys: [
-                                    Some(SessionKey::new::<Application>(noise_es_ee_se_hk_psk, 1, current_time, 2, true)),
+                                    Some(SessionKey::new::<Application>(
+                                        noise_es_ee_se_hk_psk,
+                                        1,
+                                        current_time,
+                                        2,
+                                        true,
+                                        true,
+                                    )),
                                     None,
                                 ],
                                 current_key: 0,
                                 current_offer: Offer::None,
                             }),
-                            defrag: Mutex::new(RingBufferMap::new(random::xorshift64_random() as u32)),
+                            defrag: std::array::from_fn(|_| Mutex::new(Fragged::new())),
                         });
 
-                        // Promote this from an incomplete session to an established session.
                         {
                             let mut sessions = self.sessions.write().unwrap();
                             sessions.incoming.remove(&incoming.bob_session_id);
                             sessions.active.insert(incoming.bob_session_id, Arc::downgrade(&session));
                         }
 
+                        let _ = session.send_nop(|b| send(Some(&session), b));
+
                         return Ok(ReceiveResult::OkNewSession(session));
                     } else {
                         return Err(Error::UnknownLocalSessionId);
@@ -1177,6 +1231,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                                                 current_time,
                                                 counter,
                                                 false,
+                                                false,
                                             ));
 
                                             return Ok(ReceiveResult::Ok);
@@ -1202,7 +1257,7 @@ impl<Application: ApplicationLayer> Context<Application> {
 
                     if let Some(session) = session {
                         let state = session.state.read().unwrap();
-                        if let Offer::RekeyInit(alice_e_secret, _, _) = &state.current_offer {
+                        if let Offer::RekeyInit(alice_e_secret, _) = &state.current_offer {
                             if let Some(key) = state.keys[key_index].as_ref() {
                                 // Only the current "Bob" initiates rekeys and expects this ACK.
                                 if key.bob {
@@ -1234,6 +1289,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                                                     current_time,
                                                     session.send_counter.load(Ordering::Acquire),
                                                     true,
+                                                    true,
                                                 ));
                                                 state.current_key = next_key_index; // this is an ACK so it's confirmed
                                                 state.current_offer = Offer::None;
@@ -1325,6 +1381,32 @@ impl<Application: ApplicationLayer> Session<Application> {
         return Err(Error::SessionNotEstablished);
     }
 
+    /// Send a NOP to the other side (e.g. for keep alive).
+    pub fn send_nop<SendFunction: FnMut(&mut [u8])>(&self, mut send: SendFunction) -> Result<(), Error> {
+        let state = self.state.read().unwrap();
+        if let Some(remote_session_id) = state.remote_session_id {
+            if let Some(session_key) = state.keys[state.current_key].as_ref() {
+                let counter = self.get_next_outgoing_counter().ok_or(Error::MaxKeyLifetimeExceeded)?.get();
+                let mut nop = [0u8; HEADER_SIZE + AES_GCM_TAG_SIZE];
+                let mut c = session_key.get_send_cipher(counter)?;
+                c.reset_init_gcm(&create_message_nonce(PACKET_TYPE_NOP, counter));
+                nop[HEADER_SIZE..].copy_from_slice(&c.finish_encrypt());
+                session_key.return_send_cipher(c);
+                set_packet_header(
+                    &mut nop,
+                    1,
+                    0,
+                    PACKET_TYPE_NOP,
+                    u64::from(remote_session_id),
+                    state.current_key,
+                    counter,
+                );
+                send(&mut nop);
+            }
+        }
+        return Err(Error::SessionNotEstablished);
+    }
+
     /// Check whether this session is established.
     pub fn established(&self) -> bool {
         let state = self.state.read().unwrap();
@@ -1381,7 +1463,7 @@ impl<Application: ApplicationLayer> Session<Application> {
                             .encrypt_block_in_place(&mut rekey_buf[HEADER_PROTECT_ENCRYPT_START..HEADER_PROTECT_ENCRYPT_END]);
                         send(&mut rekey_buf);
 
-                        self.state.write().unwrap().current_offer = Offer::RekeyInit(rekey_e, rekey_buf, AtomicI64::new(current_time));
+                        self.state.write().unwrap().current_offer = Offer::RekeyInit(rekey_e, current_time);
                     }
                 }
             }
@@ -1448,15 +1530,11 @@ fn set_packet_header(
 #[inline(always)]
 fn parse_packet_header(incoming_packet: &[u8]) -> (usize, u8, u8, u8, u64) {
     let raw_header_a = u16::from_le_bytes(incoming_packet[6..8].try_into().unwrap());
-    let key_index = (raw_header_a & 1) as usize;
-    let packet_type = (raw_header_a.wrapping_shr(1) & 7) as u8;
-    let fragment_count = ((raw_header_a.wrapping_shr(4) & 63) + 1) as u8;
-    let fragment_no = raw_header_a.wrapping_shr(10) as u8;
     (
-        key_index,
+        (raw_header_a & 1) as usize,
-        packet_type,
+        (raw_header_a.wrapping_shr(1) & 7) as u8,
-        fragment_count,
+        ((raw_header_a.wrapping_shr(4) & 63) + 1) as u8,
-        fragment_no,
+        raw_header_a.wrapping_shr(10) as u8,
         u64::from_le_bytes(incoming_packet[8..16].try_into().unwrap()),
     )
 }
@@ -1532,11 +1610,12 @@ impl SessionKey {
         ratchet_count: u64,
         current_time: i64,
         current_counter: u64,
-        role_is_bob: bool,
+        bob: bool,
+        confirmed: bool,
     ) -> Self {
         let a2b = kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_AES_GCM_ALICE_TO_BOB>(key.as_bytes());
         let b2a = kbkdf::<AES_256_KEY_SIZE, KBKDF_KEY_USAGE_LABEL_AES_GCM_BOB_TO_ALICE>(key.as_bytes());
-        let (receive_key, send_key) = if role_is_bob {
+        let (receive_key, send_key) = if bob {
             (a2b, b2a)
         } else {
             (b2a, a2b)
@@ -1557,7 +1636,8 @@ impl SessionKey {
             rekey_at_counter: current_counter.checked_add(Application::REKEY_AFTER_USES).unwrap(),
             expire_at_counter: current_counter.checked_add(Application::EXPIRE_AFTER_USES).unwrap(),
             ratchet_count,
-            bob: role_is_bob,
+            bob,
+            confirmed,
         }
     }