Post-refactor cleanup, break out some stuff some more in ZSSP.

crypto/benches/benchmark_crypto.rs

@@ -9,8 +9,8 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     let p384_a = P384KeyPair::generate();
     let p384_b = P384KeyPair::generate();
 
-    let kyber_a = pqc_kyber::keypair(&mut random::SecureRandom::default());
+    //let kyber_a = pqc_kyber::keypair(&mut random::SecureRandom::default());
-    let kyber_encap = pqc_kyber::encapsulate(&kyber_a.public, &mut random::SecureRandom::default()).unwrap();
+    //let kyber_encap = pqc_kyber::encapsulate(&kyber_a.public, &mut random::SecureRandom::default()).unwrap();
 
     let x25519_a = X25519KeyPair::generate();
     let x25519_b = X25519KeyPair::generate();
@@ -23,12 +23,12 @@ pub fn criterion_benchmark(c: &mut Criterion) {
         b.iter(|| p384_a.agree(p384_b.public_key()).expect("ecdhp384 failed"))
     });
     group.bench_function("ecdhx25519", |b| b.iter(|| x25519_a.agree(&x25519_b_pub)));
-    group.bench_function("kyber_encapsulate", |b| {
+    //group.bench_function("kyber_encapsulate", |b| {
-        b.iter(|| pqc_kyber::encapsulate(&kyber_a.public, &mut random::SecureRandom::default()).expect("kyber encapsulate failed"))
+    //    b.iter(|| pqc_kyber::encapsulate(&kyber_a.public, &mut random::SecureRandom::default()).expect("kyber encapsulate failed"))
-    });
+    //});
-    group.bench_function("kyber_decapsulate", |b| {
+    //group.bench_function("kyber_decapsulate", |b| {
-        b.iter(|| pqc_kyber::decapsulate(&kyber_encap.0, &kyber_a.secret).expect("kyber decapsulate failed"))
+    //    b.iter(|| pqc_kyber::decapsulate(&kyber_encap.0, &kyber_a.secret).expect("kyber decapsulate failed"))
-    });
+    //});
 
     group.finish();
 }

zssp/src/counter.rs

@@ -0,0 +1,54 @@
+use std::sync::atomic::{AtomicU64, Ordering};
+
+use zerotier_crypto::random;
+
+/// Outgoing packet counter with strictly ordered atomic semantics.
+///
+/// The counter used in packets is actually 32 bits, but using a 64-bit integer internally
+/// lets us more safely implement key lifetime limits without confusing logic to handle 32-bit
+/// wrap-around.
+#[repr(transparent)]
+pub(crate) struct Counter(AtomicU64);
+
+impl Counter {
+    #[inline(always)]
+    pub fn new() -> Self {
+        // Using a random value has no security implication. Zero would be fine. This just
+        // helps randomize packet contents a bit.
+        Self(AtomicU64::new(random::next_u32_secure() as u64))
+    }
+
+    /// Get the value most recently used to send a packet.
+    #[inline(always)]
+    pub fn previous(&self) -> CounterValue {
+        CounterValue(self.0.load(Ordering::SeqCst))
+    }
+
+    /// Get a counter value for the next packet being sent.
+    #[inline(always)]
+    pub fn next(&self) -> CounterValue {
+        CounterValue(self.0.fetch_add(1, Ordering::SeqCst))
+    }
+}
+
+/// A value of the outgoing packet counter.
+#[repr(transparent)]
+#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
+pub(crate) struct CounterValue(u64);
+
+impl CounterValue {
+    /// Get the 32-bit counter value used to build packets.
+    #[inline(always)]
+    pub fn to_u32(&self) -> u32 {
+        self.0 as u32
+    }
+
+    /// Get the counter value after N more uses of the parent counter.
+    ///
+    /// This checks for u64 overflow for the sake of correctness. Be careful if using ZSSP in a
+    /// generational starship where sessions may last for millions of years.
+    #[inline(always)]
+    pub fn counter_value_after_uses(&self, uses: u64) -> Self {
+        Self(self.0.checked_add(uses).unwrap())
+    }
+}

zssp/src/ints.rs

@@ -1,5 +1,3 @@
-use std::sync::atomic::{AtomicU64, Ordering};
-
 use zerotier_crypto::random;
 use zerotier_utils::memory;
 
@@ -61,49 +59,6 @@ impl From<SessionId> for u64 {
     }
 }
 
-/// Outgoing packet counter with strictly ordered atomic semantics.
-#[repr(transparent)]
-pub(crate) struct Counter(AtomicU64);
-
-impl Counter {
-    #[inline(always)]
-    pub fn new() -> Self {
-        // Using a random value has no security implication. Zero would be fine. This just
-        // helps randomize packet contents a bit.
-        Self(AtomicU64::new(random::next_u32_secure() as u64))
-    }
-
-    /// Get the value most recently used to send a packet.
-    #[inline(always)]
-    pub fn previous(&self) -> CounterValue {
-        CounterValue(self.0.load(Ordering::SeqCst))
-    }
-
-    /// Get a counter value for the next packet being sent.
-    #[inline(always)]
-    pub fn next(&self) -> CounterValue {
-        CounterValue(self.0.fetch_add(1, Ordering::SeqCst))
-    }
-}
-
-/// A value of the outgoing packet counter.
-///
-/// The used portion of the packet counter is the least significant 32 bits, but the internal
-/// counter state is kept as a 64-bit integer. This makes it easier to correctly handle
-/// key expiration after usage limits are reached without complicated logic to handle 32-bit
-/// wrapping. Usage limits are below 2^32 so the actual 32-bit counter will not wrap for a
-/// given shared secret key.
-#[repr(transparent)]
-#[derive(Copy, Clone)]
-pub(crate) struct CounterValue(pub u64);
-
-impl CounterValue {
-    #[inline(always)]
-    pub fn to_u32(&self) -> u32 {
-        self.0 as u32
-    }
-}
-
 /// Was this side the one who sent the first offer (Alice) or countered (Bob).
 /// Note that role is not fixed. Either side can take either role. It's just who
 /// initiated first.

zssp/src/lib.rs

@@ -1,4 +1,5 @@
 mod app_layer;
+mod counter;
 mod ints;
 mod tests;
 mod zssp;

zssp/src/zssp.rs

@@ -19,6 +19,7 @@ use zerotier_utils::varint;
 
 use crate::app_layer::ApplicationLayer;
 use crate::constants::*;
+use crate::counter::{Counter, CounterValue};
 use crate::ints::*;
 
 ////////////////////////////////////////////////////////////////
@@ -123,7 +124,7 @@ struct SessionMutableState {
 struct SessionKey {
     secret_fingerprint: [u8; 16],                 // First 128 bits of a SHA384 computed from the secret
     establish_time: i64,                          // Time session key was established
-    establish_counter: u64,                       // Counter value at which session was established
+    establish_counter: CounterValue,              // Counter value at which session was established
     lifetime: KeyLifetime,                        // Key expiration time and counter
     ratchet_key: Secret<64>,                      // Ratchet key for deriving the next session key
     receive_key: Secret<AES_KEY_SIZE>,            // Receive side AES-GCM key
@@ -147,8 +148,8 @@ struct EphemeralOffer {
 
 /// Key lifetime manager state and logic (separate to spotlight and keep clean)
 struct KeyLifetime {
-    rekey_at_or_after_counter: u64,
+    rekey_at_or_after_counter: CounterValue,
-    hard_expire_at_counter: u64,
+    hard_expire_at_counter: CounterValue,
     rekey_at_or_after_timestamp: i64,
 }
 
@@ -1482,24 +1483,22 @@ fn parse_dec_key_offer_after_header(
 impl KeyLifetime {
     fn new(current_counter: CounterValue, current_time: i64) -> Self {
         Self {
-            rekey_at_or_after_counter: current_counter.0
+            rekey_at_or_after_counter: current_counter
-                + REKEY_AFTER_USES
+                .counter_value_after_uses(REKEY_AFTER_USES)
-                + (random::next_u32_secure() % REKEY_AFTER_USES_MAX_JITTER) as u64,
+                .counter_value_after_uses((random::next_u32_secure() % REKEY_AFTER_USES_MAX_JITTER) as u64),
-            hard_expire_at_counter: current_counter.0 + EXPIRE_AFTER_USES,
+            hard_expire_at_counter: current_counter.counter_value_after_uses(EXPIRE_AFTER_USES),
             rekey_at_or_after_timestamp: current_time
                 + REKEY_AFTER_TIME_MS
                 + (random::next_u32_secure() % REKEY_AFTER_TIME_MS_MAX_JITTER) as i64,
         }
     }
 
-    #[inline(always)]
     fn should_rekey(&self, counter: CounterValue, current_time: i64) -> bool {
-        counter.0 >= self.rekey_at_or_after_counter || current_time >= self.rekey_at_or_after_timestamp
+        counter >= self.rekey_at_or_after_counter || current_time >= self.rekey_at_or_after_timestamp
     }
 
-    #[inline(always)]
     fn expired(&self, counter: CounterValue) -> bool {
-        counter.0 >= self.hard_expire_at_counter
+        counter >= self.hard_expire_at_counter
     }
 }
 
@@ -1515,7 +1514,7 @@ impl SessionKey {
         Self {
             secret_fingerprint: secret_fingerprint(key.as_bytes())[..16].try_into().unwrap(),
             establish_time: current_time,
-            establish_counter: current_counter.0,
+            establish_counter: current_counter,
             lifetime: KeyLifetime::new(current_counter, current_time),
             ratchet_key: kbkdf512(key.as_bytes(), KBKDF_KEY_USAGE_LABEL_RATCHETING),
             receive_key,