Merge pull request #34 from zerotier/dev

Fixed Issue #33 along with other minor changes

.drone.yml

@@ -16,7 +16,7 @@ steps:
         from_secret: codecov_token
     commands:
       - cargo build
-      - CARGO_INCREMENTAL=0 RUSTFLAGS='-Cinstrument-coverage' LLVM_PROFILE_FILE='coverage/cargo-test-%p-%m.profraw' cargo test
+      - CARGO_INCREMENTAL=0 RUSTFLAGS='-Cinstrument-coverage' LLVM_PROFILE_FILE='coverage/cargo-test-%p-%m.profraw' cargo test --all-targets
       - mkdir -p target/coverage
       - grcov . --binary-path ./target/debug/deps/ -s . -t lcov --branch --ignore-not-existing --ignore '../*' --ignore "/*" -o target/coverage/tests.lcov
       - codecov -B ${DRONE_BRANCH} -b ${DRONE_BUILD_NUMBER} -f target/coverage/tests.lcov -F amd64
@@ -42,7 +42,7 @@ steps:
         from_secret: codecov_token
     commands:
       - cargo build
-      - CARGO_INCREMENTAL=0 RUSTFLAGS='-Cinstrument-coverage' LLVM_PROFILE_FILE='coverage/cargo-test-%p-%m.profraw' cargo test
+      - CARGO_INCREMENTAL=0 RUSTFLAGS='-Cinstrument-coverage' LLVM_PROFILE_FILE='coverage/cargo-test-%p-%m.profraw' cargo test --all-targets
       - mkdir -p target/coverage
       - grcov . --binary-path ./target/debug/deps/ -s . -t lcov --branch --ignore-not-existing --ignore '../*' --ignore "/*" -o target/coverage/tests.lcov
       - codecov -B ${DRONE_BRANCH} -b ${DRONE_BUILD_NUMBER} -f target/coverage/tests.lcov -F arm64

performance/Cargo.toml

@@ -3,7 +3,7 @@ authors = ["ZeroTier, Inc. <contact@zerotier.com>", "Adam Ierymenko <adam.ieryme
 edition = "2021"
 license = "MPL-2.0"
 name = "zssp"
-version = "0.3.1"
+version = "0.4.0"
 
 [lib]
 name = "zssp"

performance/examples/basic_test.rs

@@ -1,3 +1,4 @@
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::iter::ExactSizeIterator;
 use std::str::FromStr;
@@ -10,7 +11,7 @@ use rand_core::OsRng;
 use rand_core::RngCore;
 
 use zssp::application::{
-    AcceptAction, ApplicationLayer, CryptoLayer, IncomingSessionAction, RatchetState, RatchetStates, RatchetUpdate,
+    AcceptAction, ApplicationLayer, CompareAndSwap, CryptoLayer, IncomingSessionAction, RatchetState, RatchetStates,
     Settings, RATCHET_SIZE,
 };
 use zssp::crypto::P384KeyPair;
@@ -114,14 +115,29 @@ impl ApplicationLayer<TestApplication> for &TestApplication {
         &mut self,
         remote_static_key: &CrateP384PublicKey,
         session_data: &u128,
-        update_data: RatchetUpdate<'_>,
-    ) -> Result<(), std::io::Error> {
+        update_data: CompareAndSwap<'_>,
+    ) -> Result<bool, std::io::Error> {
         let mut ratchets = self.ratchets.lock().unwrap();
-        ratchets.peer_map.insert(*session_data, update_data.to_states());
+        match ratchets.peer_map.entry(*session_data) {
+            Entry::Occupied(mut entry) => {
+                if update_data.compare(&entry.get()) {
+                    entry.insert(update_data.to_new_states());
+                } else {
+                    return Ok(false);
+                }
+            }
+            Entry::Vacant(entry) => {
+                if update_data.cur_is_initial_states() {
+                    entry.insert(update_data.to_new_states());
+                } else {
+                    return Ok(false);
+                }
+            }
+        }
 
         if let Some(rf) = update_data.added_fingerprint() {
-            ratchets.rf_map.insert(*rf, update_data.state1.clone());
-            println!("[{}] new ratchet #{}", self.name, update_data.state1.chain_len());
+            ratchets.rf_map.insert(*rf, update_data.new_state1.clone());
+            println!("[{}] new ratchet #{}", self.name, update_data.new_state1.chain_len());
         }
         if let Some(rf) = update_data.deleted_fingerprint1() {
             ratchets.rf_map.remove(rf);
@@ -129,7 +145,7 @@ impl ApplicationLayer<TestApplication> for &TestApplication {
         if let Some(rf) = update_data.deleted_fingerprint2() {
             ratchets.rf_map.remove(rf);
         }
-        Ok(())
+        Ok(true)
     }
 
     fn time(&mut self) -> i64 {
@@ -211,8 +227,6 @@ fn alice_main(
                             }
                         }
                     }
-                    //} else if OsRng.next_u32() | 1 > 0 {
-                    //    let _ = recursive_out.send(pkt);
                 }
             } else {
                 break;
@@ -310,8 +324,6 @@ fn bob_main(
                         }
                     }
                 }
-                //} else if OsRng.next_u32() | 1 > 0 {
-                //    let _ = recursive_out.try_send(pkt);
             }
         }
 
@@ -403,6 +415,16 @@ fn main() {
 }
 
 #[test]
-fn test_main() {
-    core(2, u32::MAX / 2)
+fn test_50() {
+    core(10, u32::MAX / 2)
+}
+
+#[test]
+fn test_75() {
+    core(10, u32::MAX / 4 * 3)
+}
+
+#[test]
+fn test_99() {
+    core(10, u32::MAX / 100 * 99)
 }

performance/examples/benchmark.rs

@@ -8,7 +8,7 @@ use rand_core::OsRng;
 use rand_core::RngCore;
 
 use zssp::application::{
-    AcceptAction, ApplicationLayer, IncomingSessionAction, RatchetState, RatchetStates, RatchetUpdate, RATCHET_SIZE,
+    AcceptAction, ApplicationLayer, CompareAndSwap, IncomingSessionAction, RatchetState, RatchetStates, RATCHET_SIZE,
 };
 use zssp::crypto::P384KeyPair;
 use zssp::crypto_impl::*;
@@ -53,7 +53,6 @@ impl AsRef<[u8]> for PooledVec {
 #[allow(unused)]
 impl DefaultCrypto for TestApplication {
     type SessionData = ();
-    type LookupData = ();
     type IncomingPacketBuffer = PooledVec;
 }
 
@@ -106,9 +105,9 @@ impl ApplicationLayer<TestApplication> for &TestApplication {
         &mut self,
         remote_static_key: &CrateP384PublicKey,
         session_data: &(),
-        update_data: RatchetUpdate<'_>,
-    ) -> Result<(), std::io::Error> {
-        Ok(())
+        update_data: CompareAndSwap<'_>,
+    ) -> Result<bool, std::io::Error> {
+        Ok(true)
     }
 
     fn time(&mut self) -> i64 {

performance/src/application.rs

@@ -277,8 +277,9 @@ pub trait ApplicationLayer<C: CryptoLayer>: Sized {
     /// This function will be called whenever Alice attempts to open a session, or Bob attempts
     /// to verify Alice's identity.
     ///
-    /// If the peer's ratchet states could not be could, this function should return
-    /// `RatchetState::new_initial_states()`.
+    /// If the peer's ratchet states could not be found, this function should return `None`.
+    /// A return value of `None` is equivalent to a return value of
+    /// `Some(RatchetState::new_initial_states())`.
     ///
     /// If a one-time-password has been pre-shared with this peer, `RatchetState::new_otp_states(...)`
     /// should be pre-saved to the storage backend as if it is a normal ratchet state.
@@ -295,11 +296,16 @@ pub trait ApplicationLayer<C: CryptoLayer>: Sized {
         session_data: &C::SessionData,
         fingerprint_data: Option<&C::FingerprintData>,
     ) -> Result<Option<RatchetStates>, std::io::Error>;
-    /// Atomically commit the update specified by `update_data` to storage, or return an error if
-    /// the update could not be made.
-    /// The implementor is free to choose how to apply these updates to storage.
+    /// Atomically compare-and-swap (a.k.a. compare-exchange) `update` to storage.
     ///
-    /// If this returns `Err(())`, the packet which triggered this function to be called will be
+    /// If `update.cur_state1` and `update.cur_state2` are currently in storage, they must
+    /// be swapped with `update.new_state1` and `update.new_state2`.
+    /// Otherwise, storage must remain unchanged.
+    ///
+    /// `Ok(true)` must be returned if the compare-and-swap was successful.
+    /// `Ok(false)` must be returned if comparison failed and the swap was cancelled.
+    ///
+    /// If this returns `Err`, the packet which triggered this function to be called will be
     /// dropped, and no session state will be mutated, preserving synchronization. The remote peer
     /// will eventually resend that packet and so this function will be called again.
     ///
@@ -311,12 +317,20 @@ pub trait ApplicationLayer<C: CryptoLayer>: Sized {
     /// This function may also save state to volatile storage, in which case all peers which connect
     /// to us will have to allow downgrade across the board.
     /// Otherwise, when we restart, we will not be allowed to reconnect.
+    ///
+    /// # Security
+    /// Implementations must not perform comparison operations (equals, less than, etc.) directly
+    /// on two ratchet keys. Instead, comparison operations must be performed indirectly
+    /// upon their ratchet fingerprints. If two ratchet states have the same ratchet fingerprint,
+    /// it should be assumed that they also have the same ratchet key.
+    ///
+    /// The implementations of `PartialEq` for `RatchetState` and `RatchetStates` do this by default.
     fn save_ratchet_state(
         &mut self,
         remote_static_key: &C::PublicKey,
         session_data: &C::SessionData,
-        update_data: RatchetUpdate<'_>,
-    ) -> Result<(), std::io::Error>;
+        update: CompareAndSwap<'_>,
+    ) -> Result<bool, std::io::Error>;
 
     /// Receives a stream of events that occur during an execution of ZSSP.
     /// These are provided for debugging, logging or metrics purposes, and must be used for

performance/src/crypto/aes.rs

@@ -47,29 +47,28 @@ pub trait Aes256Dec: Sized + Send + Sync {
     fn decrypt_in_place(&self, block: &mut [u8; AES_256_BLOCK_SIZE]);
 }
 
-pub trait AesGcmEncContext {
-    fn encrypt(&mut self, input: &[u8], output: &mut [u8]);
-
-    fn finish(self) -> [u8; AES_GCM_TAG_SIZE];
-}
-
-pub trait AesGcmDecContext {
-    fn decrypt_in_place(&mut self, data: &mut [u8]);
-
-    #[must_use]
-    fn finish(self, tag: &[u8; AES_GCM_TAG_SIZE]) -> bool;
-}
-
 /// A trait for implementing AES-GCM-256 in a way that allows for extremely high throughput.
 /// One instance of this trait is created whenever a new pair of noise keys are created,
-/// and it handles all data encryption that passes through that session.
+/// and it handles all data encryption and decryption that passes through that session.
+///
+/// It is highly recommended to implement this trait such that encryption and decryption
+/// are hardware accelerated and parallelized.
+/// ZSSP's throughput is near 90% determined by this trait.
 ///
 /// Instances must securely delete their keys when dropped.
 pub trait HighThroughputAesGcmPool: Send + Sync {
-    type EncContext<'a>: AesGcmEncContext
+    /// This type represents the state needed to stream a single plaintext for
+    /// AES-GCM Authenticated Encryption.
+    /// It should be set to whatever object, struct, handle or pointer your chosen library
+    /// uses for its stream encryption API.
+    type EncContext<'a>
     where
         Self: 'a;
-    type DecContext<'a>: AesGcmDecContext
+    /// This type represents the state needed to stream a single ciphertext for
+    /// AES-GCM Authenticated Decryption.
+    /// It should be set to whatever object, struct, handle or pointer your chosen library
+    /// uses for its stream decryption API.
+    type DecContext<'a>
     where
         Self: 'a;
 
@@ -86,17 +85,65 @@ pub trait HighThroughputAesGcmPool: Send + Sync {
     /// `nonce` must be set as the AEAD nonce.
     /// There is no additional associated data to be used.
     fn start_dec<'a>(&'a self, nonce: &[u8; AES_GCM_NONCE_SIZE]) -> Self::DecContext<'a>;
+
+    /// Stream-encrypt `input` using the specified encryption context `enc`, and write the
+    /// resulting ciphertext to `output`.
+    ///
+    /// `input` and `output` are guaranteed to have the same length.
+    ///
+    /// Be sure to update the internal state of `enc` so the authentication tag can be correctly
+    /// computed.
+    fn encrypt<'a>(&'a self, enc: &mut Self::EncContext<'a>, input: &[u8], output: &mut [u8]);
+
+    /// Stream-decrypt `data` using the specified encryption context `dec`.
+    /// The resulting plaintext should be written back into `data`.
+    ///
+    /// Be sure to update the internal state of `dec` so the authentication tag can be correctly
+    /// computed.
+    fn decrypt_in_place<'a>(&'a self, dec: &mut Self::DecContext<'a>, data: &mut [u8]);
+
+    /// Finish streaming to `enc` and output the resulting authentication tag.
+    ///
+    /// Perform any pooling or cleanup required on `enc`.
+    /// For many libraries it is much faster to return `enc` to some pool to be reused
+    /// by `start_enc`, rather than dropping it.
+    /// Be sure to perform your own benchmark.
+    fn finish_enc<'a>(&'a self, enc: Self::EncContext<'a>) -> [u8; AES_GCM_TAG_SIZE];
+    /// Finish streaming to `dec` and check that the expected authentication tag matches `tag`.
+    /// Make sure that comparing `tag` to the expected authentication tag is performed
+    /// in constant-time. Many libraries provide functions which will do this for you.
+    /// Output `true` only if `tag` is correct.
+    ///
+    /// Afterwards, perform any pooling or cleanup required on `dec`.
+    /// For many libraries it is much faster to return `dec` to some pool to be reused
+    /// by `start_dec`, rather than dropping it.
+    /// Be sure to perform your own benchmark.
+    #[must_use]
+    fn finish_dec<'a>(&'a self, dec: Self::DecContext<'a>, tag: &[u8; AES_GCM_TAG_SIZE]) -> bool;
 }
 
 /// A trait for implementing AES-GCM-256 to handle the more varied, but much lower throughput
 /// requirements of a Noise handshake.
 pub trait LowThroughputAesGcm {
+    /// A pure function (no side effects) that implements AESGCM AEAD enryption.
+    ///
+    /// Encryption must be performed on `data` in-place.
+    /// The initial plaintext of `data` must be overwriten with its ciphertext.
+    ///
+    /// The resulting GCM authentication tag must be returned.
     fn encrypt_in_place(
         key: &[u8; AES_256_KEY_SIZE],
         nonce: &[u8; AES_GCM_NONCE_SIZE],
         aad: &[u8],
         data: &mut [u8],
     ) -> [u8; AES_GCM_TAG_SIZE];
+    /// A pure function (no side effects) that implements AESGCM AEAD decryption.
+    ///
+    /// Encryption must be performed on `data` in-place.
+    /// The initial ciphertext of `data` must be overwriten with its plaintext.
+    ///
+    /// This function must check that the expected authentication tag matches `tag`,
+    /// and only return `true` if they match. This must be done in constant-time.
     #[must_use]
     fn decrypt_in_place(
         key: &[u8; AES_256_KEY_SIZE],

performance/src/crypto_impl/mod.rs

@@ -35,10 +35,22 @@ pub use openssl::*;
 #[cfg(feature = "openssl-sys")]
 pub use openssl_sys;
 
+/// Any type which implements this trait will also auto-implement the `CryptoLayer` trait,
+/// using the default set of cryptography implementations.
+///
+/// If you want to use this library, and don't care which specific implementations it uses,
+/// then implement this trait instead of the `CryptoLayer` trait.
 #[cfg(feature = "default-crypto")]
 pub trait DefaultCrypto {
+    /// Type for arbitrary opaque object for use by the application that is attached to
+    /// each session.
     type SessionData;
-    type LookupData;
+    /// Data type for incoming packet buffers.
+    ///
+    /// This can be something like `Vec<u8>` or `Box<[u8]>` or it can be something like a pooled
+    /// reusable buffer that automatically returns to its pool when ZSSP is done with it. ZSSP may
+    /// hold these for a short period of time when assembling fragmented packets on the receive
+    /// path.
     type IncomingPacketBuffer: AsMut<[u8]> + AsRef<[u8]>;
 }
 #[cfg(feature = "default-crypto")]
@@ -53,8 +65,8 @@ impl<C: DefaultCrypto> crate::application::CryptoLayer for C {
     type PublicKey = CrateP384PublicKey;
     type KeyPair = CrateP384KeyPair;
     type Kem = CrateKyber1024PrivateKey;
+    type FingerprintData = ();
 
     type SessionData = C::SessionData;
-    type FingerprintData = C::LookupData;
     type IncomingPacketBuffer = C::IncomingPacketBuffer;
 }

performance/src/crypto_impl/openssl.rs

@@ -1,12 +1,14 @@
-use std::{
-    ptr::{self, NonNull},
-    sync::{Mutex, MutexGuard},
-};
+use std::ptr::{self, NonNull};
+use std::sync::Mutex;
 
+use arrayvec::ArrayVec;
 use openssl_sys::*;
+use zeroize::Zeroizing;
 
 use crate::crypto::*;
 
+/// A wrapper for a `EVP_CIPHER_CTX` that will free itself on drop.
+/// Users are encouraged to not use one of these directly.
 pub struct OpenSSLCtx(NonNull<openssl_sys::EVP_CIPHER_CTX>);
 impl Drop for OpenSSLCtx {
     fn drop(&mut self) {
@@ -21,6 +23,8 @@ impl OpenSSLCtx {
         unsafe { Some(OpenSSLCtx(NonNull::new(EVP_CIPHER_CTX_new())?)) }
     }
 
+    /// Initialize a cipher context for encryption or decryption using the specified `key` and `iv`.
+    /// If `key` is null then the previous key assigned to this context will be used.
     pub unsafe fn cipher_init<const ENCRYPT: bool>(
         &self,
         t: *const openssl_sys::EVP_CIPHER,
@@ -36,7 +40,15 @@ impl OpenSSLCtx {
         // OpenSSL will usually leak a static amount of memory per cipher given here.
         evp_f(self.0.as_ptr(), t, ptr::null_mut(), key, iv) > 0
     }
-
+    /// Stream a portion of text to be encrypted or decrypted.
+    /// `input` will be the input to the cipher stream, and the resulting plaintext or ciphertext
+    /// will be written to output. Both buffers must be of size `len`.
+    ///
+    /// If `output == input`, then the operation will be performed "in-place".
+    /// `output` and `input` must not overlap otherwise.
+    ///
+    /// If `output` is null, then `input` will be treated as AAD rather than plaintext or ciphertext.
+    /// `input` must not be null.
     pub unsafe fn update<const ENCRYPT: bool>(&self, input: &[u8], output: *mut u8) -> bool {
         let evp_f = if ENCRYPT {
             EVP_EncryptUpdate
@@ -55,6 +67,8 @@ impl OpenSSLCtx {
         ) > 0
     }
 
+    /// Finish encryption or decryption.
+    /// If performing decryption this will return whether the set tag is correct.
     pub unsafe fn finalize<const ENCRYPT: bool>(&self) -> bool {
         let evp_f = if ENCRYPT {
             EVP_EncryptFinal_ex
@@ -66,6 +80,8 @@ impl OpenSSLCtx {
         evp_f(self.0.as_ptr(), ptr::null_mut(), &mut outl) > 0
     }
 
+    /// Retreive the authentication tag from this context.
+    /// This must be called after `finalize` is called.
     pub unsafe fn get_tag(&self, tag: &mut [u8]) -> bool {
         EVP_CIPHER_CTX_ctrl(
             self.0.as_ptr(),
@@ -74,6 +90,10 @@ impl OpenSSLCtx {
             tag.as_mut_ptr() as *mut _,
         ) > 0
     }
+
+    /// Set the authentication tag that was assigned to the input ciphertext.
+    /// Once set, OpenSLL will check whether it matches the expected authentication tag
+    /// produced by decryption.
     #[allow(unused)]
     pub unsafe fn set_tag(&self, tag: &[u8]) -> bool {
         EVP_CIPHER_CTX_ctrl(
@@ -83,11 +103,17 @@ impl OpenSSLCtx {
             tag.as_ptr() as *mut _,
         ) > 0
     }
+    /// Returns the raw pointer to the `EVP_CIPHER_CTX`
+    /// object used internally with OpenSSL.
+    ///
+    /// This function is guaranteed to return a non-null pointer.
     pub fn as_ptr(&self) -> *mut openssl_sys::EVP_CIPHER_CTX {
         self.0.as_ptr()
     }
 }
 
+/// An `OpenSSLCtx` wrapped in a mutex for thread-safety.
+/// This `OpenSSLCtx` struct only supports AES256 block operations, and implements `OpenSSLAes256Enc`.
 pub struct OpenSSLAes256Enc(Mutex<OpenSSLCtx>);
 unsafe impl Send for OpenSSLAes256Enc {}
 unsafe impl Sync for OpenSSLAes256Enc {}
@@ -118,6 +144,8 @@ impl Aes256Enc for OpenSSLAes256Enc {
         unsafe { assert!(ctx.update::<true>(block, ptr)) }
     }
 }
+/// An `OpenSSLCtx` wrapped in a mutex for thread-safety.
+/// This `OpenSSLCtx` struct only supports AES256 block operations, and implements `OpenSSLAes256Dec`.
 pub struct OpenSSLAes256Dec(Mutex<OpenSSLCtx>);
 unsafe impl Send for OpenSSLAes256Dec {}
 unsafe impl Sync for OpenSSLAes256Dec {}
@@ -149,87 +177,87 @@ impl Aes256Dec for OpenSSLAes256Dec {
     }
 }
 
-pub struct OpenSSLAesGcmEnc<'a>(MutexGuard<'a, OpenSSLCtx>);
-impl<'a> AesGcmEncContext for OpenSSLAesGcmEnc<'a> {
-    fn encrypt(&mut self, input: &[u8], output: &mut [u8]) {
-        unsafe { assert!(self.0.update::<true>(input, output.as_mut_ptr())) };
-    }
-
-    fn finish(self) -> [u8; AES_GCM_TAG_SIZE] {
-        let mut output = [0u8; AES_GCM_TAG_SIZE];
-        unsafe {
-            assert!(self.0.finalize::<true>());
-            assert!(self.0.get_tag(&mut output));
-        }
-        output
-    }
-}
-
-pub struct OpenSSLAesGcmDec<'a>(MutexGuard<'a, OpenSSLCtx>);
-impl<'a> AesGcmDecContext for OpenSSLAesGcmDec<'a> {
-    fn decrypt_in_place(&mut self, data: &mut [u8]) {
-        let p = data.as_mut_ptr();
-        unsafe { assert!(self.0.update::<false>(data, p)) };
-    }
-
-    fn finish(self, tag: &[u8; AES_GCM_TAG_SIZE]) -> bool {
-        unsafe { self.0.set_tag(tag) && self.0.finalize::<false>() }
-    }
-}
-
 /// A pool of OpenSSL AES-GCM ciphers.
 pub struct OpenSSLAesGcmPool {
-    enc: [Mutex<OpenSSLCtx>; 8],
-    dec: [Mutex<OpenSSLCtx>; 8],
+    enc: Mutex<ArrayVec<OpenSSLCtx, 8>>,
+    dec: Mutex<ArrayVec<OpenSSLCtx, 8>>,
+    enc_key: Zeroizing<[u8; AES_256_KEY_SIZE]>,
+    dec_key: Zeroizing<[u8; AES_256_KEY_SIZE]>,
 }
 unsafe impl Send for OpenSSLAesGcmPool {}
 unsafe impl Sync for OpenSSLAesGcmPool {}
 
 impl HighThroughputAesGcmPool for OpenSSLAesGcmPool {
-    type EncContext<'a> = OpenSSLAesGcmEnc<'a>;
+    type EncContext<'a> = OpenSSLCtx;
 
-    type DecContext<'a> = OpenSSLAesGcmDec<'a>;
+    type DecContext<'a> = OpenSSLCtx;
 
     fn new(encrypt_key: &[u8; AES_256_KEY_SIZE], decrypt_key: &[u8; AES_256_KEY_SIZE]) -> Self {
+        Self {
+            enc: Default::default(),
+            dec: Default::default(),
+            enc_key: Zeroizing::new(*encrypt_key),
+            dec_key: Zeroizing::new(*decrypt_key),
+        }
+    }
+
+    fn start_enc(&self, nonce: &[u8; AES_GCM_NONCE_SIZE]) -> OpenSSLCtx {
+        let ctx = self.enc.lock().unwrap().pop();
         unsafe {
-            OpenSSLAesGcmPool {
-                enc: std::array::from_fn(|_| {
-                    let ctx = OpenSSLCtx::new().unwrap();
-                    let t = openssl_sys::EVP_aes_256_gcm();
-                    assert!(ctx.cipher_init::<true>(t, encrypt_key.as_ptr(), ptr::null()));
-                    openssl_sys::EVP_CIPHER_CTX_set_padding(ctx.as_ptr(), 0);
-                    Mutex::new(ctx)
-                }),
-                dec: std::array::from_fn(|_| {
-                    let ctx = OpenSSLCtx::new().unwrap();
-                    let t = openssl_sys::EVP_aes_256_gcm();
-                    assert!(ctx.cipher_init::<false>(t, decrypt_key.as_ptr(), ptr::null()));
-                    openssl_sys::EVP_CIPHER_CTX_set_padding(ctx.as_ptr(), 0);
-                    Mutex::new(ctx)
-                }),
+            if let Some(ctx) = ctx {
+                assert!(ctx.cipher_init::<true>(ptr::null(), ptr::null(), nonce.as_ptr()));
+                ctx
+            } else {
+                let ctx = OpenSSLCtx::new().unwrap();
+                let t = openssl_sys::EVP_aes_256_gcm();
+                assert!(ctx.cipher_init::<true>(t, self.enc_key.as_ptr(), nonce.as_ptr()));
+                openssl_sys::EVP_CIPHER_CTX_set_padding(ctx.as_ptr(), 0);
+                ctx
+            }
+        }
+    }
+    fn start_dec(&self, nonce: &[u8; AES_GCM_NONCE_SIZE]) -> OpenSSLCtx {
+        let ctx = self.dec.lock().unwrap().pop();
+        unsafe {
+            if let Some(ctx) = ctx {
+                assert!(ctx.cipher_init::<false>(ptr::null(), ptr::null(), nonce.as_ptr()));
+                ctx
+            } else {
+                let ctx = OpenSSLCtx::new().unwrap();
+                let t = openssl_sys::EVP_aes_256_gcm();
+                assert!(ctx.cipher_init::<false>(t, self.dec_key.as_ptr(), nonce.as_ptr()));
+                openssl_sys::EVP_CIPHER_CTX_set_padding(ctx.as_ptr(), 0);
+                ctx
             }
         }
     }
 
-    fn start_enc<'a>(&'a self, nonce: &[u8; AES_GCM_NONCE_SIZE]) -> OpenSSLAesGcmEnc<'a> {
-        let i = u64::from_be_bytes(nonce[4..].try_into().unwrap());
-        let g = self.enc[(i as usize) % self.enc.len()].lock().unwrap();
-        unsafe {
-            assert!(g.cipher_init::<true>(ptr::null(), ptr::null(), nonce.as_ptr()));
-        }
-        OpenSSLAesGcmEnc(g)
+    fn encrypt(&self, ctx: &mut OpenSSLCtx, input: &[u8], output: &mut [u8]) {
+        unsafe { assert!(ctx.update::<true>(input, output.as_mut_ptr())) };
+    }
+    fn decrypt_in_place(&self, ctx: &mut OpenSSLCtx, data: &mut [u8]) {
+        let p = data.as_mut_ptr();
+        unsafe { assert!(ctx.update::<false>(data, p)) };
     }
 
-    fn start_dec<'a>(&'a self, nonce: &[u8; AES_GCM_NONCE_SIZE]) -> OpenSSLAesGcmDec<'a> {
-        let i = u64::from_be_bytes(nonce[4..].try_into().unwrap());
-        let g = self.dec[(i as usize) % self.enc.len()].lock().unwrap();
+    fn finish_enc(&self, ctx: OpenSSLCtx) -> [u8; AES_GCM_TAG_SIZE] {
+        let mut output = [0u8; AES_GCM_TAG_SIZE];
         unsafe {
-            assert!(g.cipher_init::<false>(ptr::null(), ptr::null(), nonce.as_ptr()));
+            assert!(ctx.finalize::<true>());
+            assert!(ctx.get_tag(&mut output));
         }
-        OpenSSLAesGcmDec(g)
+        let _ = self.enc.lock().unwrap().try_push(ctx);
+        output
+    }
+    fn finish_dec(&self, ctx: OpenSSLCtx, tag: &[u8; AES_GCM_TAG_SIZE]) -> bool {
+        let output = unsafe { ctx.set_tag(tag) && ctx.finalize::<false>() };
+        let _ = self.dec.lock().unwrap().try_push(ctx);
+        output
     }
 }
-
+/// An empty struct which implements `LowThroughputAesGcm` using OpenSSL.
+///
+/// It is just a namespace and wrapper for OpenSSL.
 pub struct OpenSSLAesGcm;
 impl LowThroughputAesGcm for OpenSSLAesGcm {
     fn encrypt_in_place(

performance/src/lib.rs

@@ -36,9 +36,34 @@
 //!  - **KBKDF**: Key mixing, sub-key derivation
 //!  - **AES-256**: Single block encryption of header to harden packet fragmentation protocol
 //!  - **AES-256-GCM**: Authenticated encryption
-//#![warn(missing_docs, rust_2018_idioms)]
+#![warn(missing_docs, rust_2018_idioms)]
 #![allow(clippy::too_many_arguments, clippy::type_complexity, clippy::assertions_on_constants)]
+
+/// A collection of implementation-independent traits for the various specific cryptographic
+/// algorithms ZSSP depends on.
+///
+/// Each trait is very specific about the semantics of the algorithms and the lengths of their
+/// inputs and outputs.
+/// This is to enforced a basic sanity-check upon anyone trying to connect this library to
+/// someone else's implementation.
+///
+/// You should not use this module to implement your own crypto from scratch.
+///
+/// The `crypto_impl` module contains implementations of these traits in terms of popular Rust
+/// crates.
 pub mod crypto;
+/// A module containing optional implementations of the ZSSP `crypto` traits in terms of popular
+/// Rust crates. Some of these crates are not thoroughly audited, so use at your own risk.
+///
+/// The current version of the AES crate does not support stream encryption. For this reason
+/// ZSSP's AES trait is implemented with OpenSSL instead. It has been optimized for
+/// hardware accelerated, parrallel encryption and decryption.
+///
+/// Note that none of these crates are FIPS certified, meaning a build of ZSSP using them will not
+/// be FIPS compliant. However lack of FIPS compliance by no means implies lack of security or lack
+/// of confidence.
+///
+/// This module contains the trait implementations as well as re-exports of those crates.
 pub mod crypto_impl;
 
 mod antireplay;
@@ -58,8 +83,14 @@ mod symmetric_state;
 mod zeta;
 mod zssp;
 
+/// An abstraction over OS and use-case specific resources and queries.
+/// This allows this library to be platform independent.
+/// A user of this library will need to implement the `ApplicationLayer` trait.
 pub mod application;
+/// This module contains several ZSSP constants that a user might want to know for key management,
+/// or in order to avoid "data too large" or "mtu too small" errors.
 pub mod proto;
+/// The collection of the major return types for ZSSP.
 pub mod result;
 
 pub use crate::log_event::*;

performance/src/proto.rs

@@ -124,14 +124,12 @@ pub(crate) const PACKET_TYPE_DATA: u8 = 8;
 pub(crate) const PACKET_TYPE_CHALLENGE: u8 = 9;
 pub(crate) const PACKET_TYPE_USES_COUNTER_RANGE: std::ops::Range<u8> = 3..9;
 
-pub(crate) const HANDSHAKE_HELLO_MIN_SIZE: usize =
-    KID_SIZE + P384_PUBLIC_KEY_SIZE + KYBER_PUBLIC_KEY_SIZE + AES_GCM_TAG_SIZE + AES_GCM_TAG_SIZE;
-pub(crate) const HANDSHAKE_HELLO_MAX_SIZE: usize = HANDSHAKE_HELLO_MIN_SIZE + RATCHET_SIZE + RATCHET_SIZE;
+pub(crate) const HANDSHAKE_HELLO_SIZE: usize =
+    KID_SIZE + P384_PUBLIC_KEY_SIZE + KYBER_PUBLIC_KEY_SIZE + AES_GCM_TAG_SIZE + 2 * RATCHET_SIZE + AES_GCM_TAG_SIZE;
 
-pub(crate) const HANDSHAKE_HELLO_CHALLENGE_MIN_SIZE: usize = HANDSHAKE_HELLO_MIN_SIZE + CHALLENGE_SIZE;
-pub(crate) const HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE: usize = HANDSHAKE_HELLO_MAX_SIZE + CHALLENGE_SIZE;
+pub(crate) const HANDSHAKE_HELLO_CHALLENGE_SIZE: usize = HANDSHAKE_HELLO_SIZE + CHALLENGE_SIZE;
 
-pub(crate) const HEADERED_HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE: usize = HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE + HEADER_SIZE;
+pub(crate) const HEADERED_HANDSHAKE_HELLO_CHALLENGE_SIZE: usize = HANDSHAKE_HELLO_CHALLENGE_SIZE + HEADER_SIZE;
 
 pub(crate) const HANDSHAKE_RESPONSE_SIZE: usize =
     P384_PUBLIC_KEY_SIZE + KYBER_CIPHERTEXT_SIZE + AES_GCM_TAG_SIZE + KID_SIZE + AES_GCM_TAG_SIZE;
@@ -184,6 +182,8 @@ pub(crate) const MAX_UNASSOCIATED_FRAGMENTS: usize = 32 * 32;
 
 /// The maximum size a packet that is not associated to a session may be.
 /// Excludes the size of headers for fragmentation.
-pub(crate) const MAX_UNASSOCIATED_PACKET_SIZE: usize = HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE;
+pub(crate) const MAX_UNASSOCIATED_PACKET_SIZE: usize = HANDSHAKE_HELLO_CHALLENGE_SIZE;
 
+/// This number determines how many defragmentation buffers are created per session.
+/// Each defragmentation buffer handles one packet at a time.
 pub(crate) const SESSION_MAX_FRAGMENTS_OOO: usize = 64;

performance/src/ratchet_state.rs

@@ -14,30 +14,23 @@ use crate::proto::*;
 #[derive(Clone, Eq)]
 pub struct RatchetState {
     pub(crate) key: Zeroizing<[u8; RATCHET_SIZE]>,
-    pub(crate) fingerprint: Option<Zeroizing<[u8; RATCHET_SIZE]>>,
+    pub(crate) fingerprint: Zeroizing<[u8; RATCHET_SIZE]>,
     pub(crate) chain_len: u64,
 }
 impl PartialEq for RatchetState {
     fn eq(&self, other: &Self) -> bool {
-        let ret = match (self.fingerprint.as_ref(), other.fingerprint.as_ref()) {
-            (Some(rf1), Some(rf2)) => secure_eq(rf1, rf2),
-            (None, None) => true,
-            _ => false,
-        };
-        ret & secure_eq(&self.key, &other.key) & (self.chain_len == other.chain_len)
+        self.fingerprint.eq(&other.fingerprint) & (self.chain_len == other.chain_len)
     }
 }
 impl std::hash::Hash for RatchetState {
     fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
-        if let Some(rf) = &self.fingerprint {
-            state.write_u64(u64::from_ne_bytes(rf[..8].try_into().unwrap()))
-        }
+        state.write(self.fingerprint.as_ref())
     }
 }
 impl RatchetState {
     /// Creates a new ratchet state from the given ratchet key, ratchet fingerprint and chain length.
     pub fn new(key: Zeroizing<[u8; RATCHET_SIZE]>, fingerprint: Zeroizing<[u8; RATCHET_SIZE]>, chain_len: u64) -> Self {
-        RatchetState { key, fingerprint: Some(fingerprint), chain_len }
+        RatchetState { key, fingerprint, chain_len }
     }
     /// Creates a new ratchet state from the given ratchet key, ratchet fingerprint and chain length.
     ///
@@ -46,28 +39,10 @@ impl RatchetState {
     pub fn new_raw(key: [u8; RATCHET_SIZE], fingerprint: [u8; RATCHET_SIZE], chain_len: u64) -> Self {
         RatchetState {
             key: Zeroizing::new(key),
-            fingerprint: Some(Zeroizing::new(fingerprint)),
+            fingerprint: Zeroizing::new(fingerprint),
             chain_len,
         }
     }
-    /// The ratchet key for this ratchet state. This is directly mixed into the master secret of a
-    /// session and so is very sensitive. All operations upon a ratchet key must be implemented
-    /// in constant time. The user should prefer to do nothing with the ratchet key besides copying
-    /// it to or from a storage device.
-    ///
-    /// If `fingerprint` returns `None` then this is the "empty" ratchet state and the key will be
-    /// all zeros.
-    pub fn key(&self) -> &[u8; RATCHET_SIZE] {
-        &self.key
-    }
-    /// Ratchet keys and fingerprints are "chained together", where each set is derived from the
-    /// previous set.
-    ///
-    /// This function outputs the total length of that chain, as in the total number of previous
-    /// ratchet states that this ratchet state was derived from.
-    pub fn chain_len(&self) -> u64 {
-        self.chain_len
-    }
     /// Creates a new "empty" ratchet state, where the ratchet fingerprint is the
     /// empty string, the ratchet key is all zeros, and the chain length is 0.
     ///
@@ -75,7 +50,7 @@ impl RatchetState {
     pub fn empty() -> Self {
         RatchetState {
             key: Zeroizing::new([0u8; RATCHET_SIZE]),
-            fingerprint: None,
+            fingerprint: Zeroizing::new([0u8; RATCHET_SIZE]),
             chain_len: 0,
         }
     }
@@ -99,15 +74,15 @@ impl RatchetState {
 
         Self::new(rk, rf, 1)
     }
-    /// Returns true if this is the "empty" ratchet state, where the ratchet fingerprint is the
-    /// empty string, the ratchet key is all zeros, and the chain length is 0.
-    pub fn is_empty(&self) -> bool {
-        self.fingerprint.is_none()
-    }
-    /// Checks if the fingerprint of this ratchet state equals the fingerprint contained in argument
-    /// `rf`. Uses constant time equality.
-    pub fn fingerprint_eq(&self, rf: &[u8; RATCHET_SIZE]) -> bool {
-        self.fingerprint.as_ref().map_or(false, |rf0| secure_eq(rf0, rf))
+    /// The ratchet key for this ratchet state. This is directly mixed into the master secret of a
+    /// session and so is very sensitive. All operations upon a ratchet key must be implemented
+    /// in constant time. The user should prefer to do nothing with the ratchet key besides copying
+    /// it to or from a storage device.
+    ///
+    /// If `fingerprint` returns `None` then this is the "empty" ratchet state and the key will be
+    /// all zeros.
+    pub fn key(&self) -> &[u8; RATCHET_SIZE] {
+        &self.key
     }
     /// The ratchet fingerprint for this ratchet state.
     ///
@@ -118,8 +93,43 @@ impl RatchetState {
     /// but the security of ZSSP can survive having this value leaked.
     /// Operations on a ratchet fingerprint should be implemented in constant time,
     /// but it is ok if they are not.
-    pub fn fingerprint(&self) -> Option<&[u8; RATCHET_SIZE]> {
-        self.fingerprint.as_deref()
+    pub fn fingerprint(&self) -> &[u8; RATCHET_SIZE] {
+        &self.fingerprint
+    }
+    /// Ratchet keys and fingerprints are "chained together", where each set is derived from the
+    /// previous set.
+    ///
+    /// This function outputs the total length of that chain, as in the total number of previous
+    /// ratchet states that this ratchet state was derived from.
+    pub fn chain_len(&self) -> u64 {
+        self.chain_len
+    }
+    /// Returns true if this is the "empty" ratchet state, where the ratchet fingerprint is the
+    /// empty string, the ratchet key is all zeros, and the chain length is 0.
+    pub fn is_empty(&self) -> bool {
+        secure_eq(self.fingerprint(), &[0u8; RATCHET_SIZE])
+    }
+    /// Checks if the fingerprint of this ratchet state equals the
+    ///  fingerprint contained in argument `rf`.
+    ///
+    /// Uses constant time equality.
+    pub fn fingerprint_eq(&self, rf: &[u8; RATCHET_SIZE]) -> bool {
+        secure_eq(self.fingerprint(), rf)
+    }
+    /// Returns true if the fingerprint of argument `this` equals the fingerprint represented by
+    /// argument `rf`.
+    /// If `this` is `None`, `this` is considered to be "null", as in a ratchet state that does not
+    /// exist or has been deleted.
+    /// If `rf` is `None`, `rf` is considered to be "null" as well.
+    /// If `rf` is `Some(None)`, `rf` is considered to be the empty ratchet fingerprint.
+    ///
+    /// Uses constant time equality.
+    pub fn fingerprint_eq_nullable(this: Option<&Self>, rf: Option<&[u8; RATCHET_SIZE]>) -> bool {
+        match (this, rf) {
+            (Some(this), Some(rf)) => this.fingerprint_eq(rf),
+            (None, None) => true,
+            _ => false,
+        }
     }
 }
 impl Default for RatchetState {
@@ -137,7 +147,8 @@ impl Default for RatchetState {
 pub struct RatchetStates {
     /// The first ratchet state from the pair.
     pub state1: RatchetState,
-    /// The second ratchet state from the pair. It can, and usually will be `None`.
+    /// The second ratchet state from the pair.
+    /// It can, and usually will be `None`, which means that this ratchet state is "null".
     pub state2: Option<RatchetState>,
 }
 impl RatchetStates {
@@ -172,7 +183,6 @@ impl Default for RatchetStates {
         Self::new_initial_states()
     }
 }
-
 /// A set of references to ratchet states specifying how a remote peer's persistent storage should
 /// be updated. This struct is designed to provide any and all potentially needed data for
 /// maintaining a store of these ratchet states. It should be straightforward to commit these updates
@@ -181,69 +191,136 @@ impl Default for RatchetStates {
 /// There will only be up to two ratchet states saved to storage at a time per peer.
 /// Every time a third ratchet state is generated, a previous ratchet state will be deleted.
 ///
-/// These are sensitive values should they ought to be securely stored, with restricted read-write
-/// permissions if stored on disk.
+/// As the name implies, these updates should be applied as one atomic compare-and-swap operation.
+/// If the two ratchet states currently in storage equal `cur_state1` and `cur_state2`,
+/// then `new_state1` and `new_state2` should overwrite them.
+/// If not, then storage must not be modified.
 ///
-/// To prevent desync or resource leakage, these updates should be committed atomically. If that is
-/// not possible, then new ratchet states should be written before old ratchet states are deleted
+/// Ratchet keys must never be checked for equality. Instead, if two ratchet states have equal
+/// ratchet fingerprints, it should be assumed that they also have equal ratchet keys.
+/// This policy exists to reduce the security impact of timing and other side-channel attacks.
+///
+/// These are sensitive values and ought to be securely stored, with restricted read-write
+/// permissions if stored on disk.
 #[derive(Clone)]
-pub struct RatchetUpdate<'a> {
-    /// The ratchet key and fingerprint to store in the first slot.
-    pub state1: &'a RatchetState,
-    /// The ratchet key and fingerprint to store in the second slot.
-    pub state2: Option<&'a RatchetState>,
-    /// Whether `state1` is a brand new ratchet state, or if it was previously saved.
-    pub state1_was_just_added: bool,
-    /// A previous ratchet key and fingerprint that now must be deleted from storage.
-    /// This will have been a previously given value of `state1` or `state2`.
-    pub deleted_state1: Option<&'a RatchetState>,
-    /// A previous ratchet key and fingerprint that now must be deleted from storage.
-    /// It is extremely rare that this field is occupied.
-    pub deleted_state2: Option<&'a RatchetState>,
+pub struct CompareAndSwap<'a> {
+    /// The ratchet state to store in the first slot.
+    pub new_state1: &'a RatchetState,
+    /// The ratchet state to store in the second slot.
+    /// A value of `None` implies that the second slot should be set to null.
+    pub new_state2: Option<&'a RatchetState>,
+    /// This field is `true` if and only if `new_state1 != cur_state1` and `new_state1 != cur_state2`.
+    ///
+    /// This implies whether `state1` is a brand new ratchet state, or if it was previously saved.
+    pub new_state1_was_just_added: bool,
+    /// The ratchet state that we expect to see in the first slot.
+    /// If this value is not currently stored in the first slot, the entire update must be aborted.
+    pub cur_state1: &'a RatchetState,
+    /// The ratchet state that we expect to see in the second slot.
+    /// A value of `None` implies we expect the second slot to be set to "null".
+    ///
+    /// If this value is not currently stored in the second slot, the entire update must be aborted.
+    pub cur_state2: Option<&'a RatchetState>,
+    /// This field is `true` if and only if `cur_state1 != new_state1` and `cur_state1 != new_state2`.
+    pub cur_state1_was_just_deleted: bool,
+    /// This field is `true` if and only if `cur_state2 != new_state1` and `cur_state2 != new_state2`.
+    pub cur_state2_was_just_deleted: bool,
 }
-impl<'a> RatchetUpdate<'a> {
-    /// Returns the final `RatchetStates` that should be the only thing saved after this update is
-    /// fully committed. Future calls to `ApplicationLayer::restore_by_identity` should return this struct.
-    pub fn to_states(&self) -> RatchetStates {
-        RatchetStates::new(self.state1.clone(), self.state2.cloned())
+impl<'a> CompareAndSwap<'a> {
+    pub(crate) fn new(
+        new_state1: &'a RatchetState,
+        new_state2: Option<&'a RatchetState>,
+        new_state1_was_just_added: bool,
+        cur_state1: &'a RatchetState,
+        cur_state2: Option<&'a RatchetState>,
+        cur_state1_was_just_deleted: bool,
+        cur_state2_was_just_deleted: bool,
+    ) -> Self {
+        Self {
+            new_state1,
+            new_state2,
+            new_state1_was_just_added,
+            cur_state1,
+            cur_state2,
+            cur_state1_was_just_deleted,
+            cur_state2_was_just_deleted,
+        }
+    }
+    /// Returns the final `RatchetStates` that must be swapped into storage if this update is
+    /// fully committed.
+    /// Future calls to `ApplicationLayer::restore_by_identity` should return this struct.
+    pub fn to_new_states(&self) -> RatchetStates {
+        RatchetStates::new(self.new_state1.clone(), self.new_state2.cloned())
+    }
+    /// Returns the `RatchetStates` that is expected to currently be in storage for this peer.
+    /// This value must be compared with the current value in storage, and if they are equal,
+    /// the return value of `CompareAndSwap::to_new_states` must overwrite it.
+    pub fn to_cur_states(&self) -> RatchetStates {
+        RatchetStates::new(self.cur_state1.clone(), self.cur_state2.cloned())
     }
     /// If this update specifies adding a brand new ratchet fingerprint, this function will return it.
     /// The returned ratchet fingerprint will always be the ratchet fingerprint of field `state1`.
     ///
-    /// If a fingerprint is returned then it is guaranteed that `state1_was_just_added` will be true.
+    /// If a fingerprint is returned then it is guaranteed that `state1_was_just_added` will be `true`.
     pub fn added_fingerprint(&self) -> Option<&[u8; RATCHET_SIZE]> {
-        if self.state1_was_just_added {
-            self.state1.fingerprint()
-        } else {
-            None
-        }
+        self.new_state1_was_just_added.then_some(self.new_state1.fingerprint())
     }
-    /// If this updated specifies to delete an old ratchet fingerprint, this function will return it.
-    /// The returned ratchet fingerprint will always be the ratchet fingerprint of field
-    /// `deleted_state1`.
+    /// If this updated specifies to delete `cur_state1`, and `cur_state1` has a non-zero ratchet
+    /// fingerprint, this function will return the ratchet fingerprint.
+    ///
+    /// If `cur_state1` is not being deleted (i.e. `cur_state1 == new_state1` or
+    /// `cur_state1 == new_state2`), then this will return `None`.
     ///
     /// There may be a second ratchet fingerprint to be deleted, which function
     /// `deleted_fingerprint2` will return.
     pub fn deleted_fingerprint1(&self) -> Option<&[u8; RATCHET_SIZE]> {
-        if let Some(rs) = &self.deleted_state1 {
-            rs.fingerprint()
-        } else {
-            None
+        if self.cur_state1_was_just_deleted && !self.cur_state1.is_empty() {
+            return Some(self.cur_state1.fingerprint());
         }
+        None
     }
-    /// If this updated specifies to delete two ratchet fingerprints, this function will return the
-    /// second one. `deleted_fingerprint1` will return the first one.
-    /// The returned ratchet fingerprint will always be the ratchet fingerprint of field
-    /// `deleted_state2`.
+    /// If this updated specifies to delete `cur_state2`, and `cur_state2` has a non-zero ratchet
+    /// fingerprint, this function will return the ratchet fingerprint.
     ///
-    /// It is exceptionally rare that there will be more than one ratchet fingerprint to be deleted.
-    /// Care should be taken to make sure that this update will still be correctly committed in the
-    /// rare event that this returns `Some`.
+    /// If `cur_state2` is not being deleted (i.e. `cur_state2 == new_state1` or
+    /// `cur_state2 == new_state2`), then this will return `None`.
+    ///
+    /// There may be a second ratchet fingerprint to be deleted, which function
+    /// `deleted_fingerprint1` will return.
     pub fn deleted_fingerprint2(&self) -> Option<&[u8; RATCHET_SIZE]> {
-        if let Some(rs) = &self.deleted_state2 {
-            rs.fingerprint()
-        } else {
-            None
+        if self.cur_state2_was_just_deleted {
+            if let Some(rf) = self.cur_state2 {
+                if !rf.is_empty() {
+                    return Some(rf.fingerprint());
+                }
+            }
         }
+        None
+    }
+    /// Returns true if the currently stored ratchet states are expected to be the initial ratchet
+    /// states. This is the default value for a peer's ratchet states in the event they could not
+    /// be found in storage.
+    ///
+    /// If a peer could not be found in storage, and this function returns true,
+    /// then the peer should be added to storage with the new ratchet states specified by this
+    /// `CompareAndSwap` struct.
+    pub fn cur_is_initial_states(&self) -> bool {
+        self.cur_state1.is_empty() && self.cur_state2.is_none()
+    }
+    /// Compares the ratchet fingerprints of `cur_state1` and `cur_state2` with `rf1` and `rf2`.
+    /// If they are equal this function will return `true`.
+    ///
+    /// If this function returns `true`, then the implementation may proceed to swap out the ratchet
+    /// states these fingerprints come from with `new_state1` and `new_state2`.
+    pub fn compare_fingerprints(&self, rf1: &[u8; RATCHET_SIZE], rf2: Option<&[u8; RATCHET_SIZE]>) -> bool {
+        self.cur_state1.fingerprint_eq(rf1) & RatchetState::fingerprint_eq_nullable(self.cur_state2, rf2)
+    }
+    /// Compares `cur_state1` and `cur_state2` with `other.state1` and `other.state2`.
+    /// If they are equal this function will return `true`.
+    ///
+    /// If this function returns `true`, then the implementation may proceed to swap `other.state1`
+    /// and `other.state2` with `new_state1` and `new_state2`.
+    pub fn compare(&self, other: &RatchetStates) -> bool {
+        self.cur_state1.eq(&other.state1) & self.cur_state2.eq(&other.state2.as_ref())
     }
 }

performance/src/result.rs

@@ -190,7 +190,7 @@ pub enum ReceiveOk<C: CryptoLayer> {
     Unassociated,
     /// The received packet was authentic and belongs to this specific session.
     Associated(Arc<Session<C>>, SessionEvent),
-    /// The received packet was a fragment of a larger packet.
+    /// The received packet was an incomplete fragment of a larger packet.
     ///
     /// ***The authenticity of this fragment cannot be fully known yet.***
     /// This return value should only be used for debugging and tracing purposes.

performance/src/zeta.rs

@@ -97,7 +97,7 @@ pub(crate) struct StateA1<C: CryptoLayer> {
     e_secret: C::KeyPair,
     e1_secret: C::Kem,
     identity: ArrayVec<u8, IDENTITY_MAX_SIZE>,
-    x1: ArrayVec<u8, HEADERED_HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE>,
+    x1: ArrayVec<u8, HEADERED_HANDSHAKE_HELLO_CHALLENGE_SIZE>,
 }
 
 pub(crate) struct StateA3 {
@@ -280,7 +280,7 @@ fn create_a1_state<C: CryptoLayer>(
     //    ...
     //    -> e, es, e1
     let mut noise = SymmetricState::<C>::initialize(PROTOCOL_NAME_NOISE_XK);
-    let mut x1 = ArrayVec::<u8, HEADERED_HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE>::new();
+    let mut x1 = ArrayVec::<u8, HEADERED_HANDSHAKE_HELLO_CHALLENGE_SIZE>::new();
     x1.extend([0u8; HEADER_SIZE]);
     // Noise process prologue.
     let kid = kid_recv.get().to_ne_bytes();
@@ -299,12 +299,9 @@ fn create_a1_state<C: CryptoLayer>(
     x1.extend(tag);
     // Process message pattern 1 payload.
     let i = x1.len();
-    if let Some(rf) = ratchet_state1.fingerprint() {
-        x1.try_extend_from_slice(rf).unwrap();
-    }
-    if let Some(Some(rf)) = ratchet_state2.map(|rs| rs.fingerprint()) {
-        x1.try_extend_from_slice(rf).unwrap();
-    }
+    x1.try_extend_from_slice(ratchet_state1.fingerprint()).unwrap();
+    x1.try_extend_from_slice(ratchet_state2.map_or(&[0u8; RATCHET_SIZE], |r| r.fingerprint()))
+        .unwrap();
     let tag = noise.encrypt_and_hash_in_place(hash, to_nonce(PACKET_TYPE_HANDSHAKE_HELLO, 1), &mut x1[i..]);
     x1.extend(tag);
 
@@ -325,12 +322,10 @@ pub(crate) fn trans_to_a1<C: CryptoLayer, App: ApplicationLayer<C>>(
     s_remote: C::PublicKey,
     session_data: C::SessionData,
     identity: &[u8],
+    ratchet_states: RatchetStates,
     send: impl FnOnce(&mut [u8], Option<&C::PrpEnc>),
 ) -> Result<(Arc<Session<C>>, Option<i64>), OpenError> {
-    let RatchetStates { state1, state2 } = app
-        .restore_by_identity(&s_remote, &session_data, None)
-        .map_err(OpenError::StorageError)?
-        .unwrap_or_default();
+    let RatchetStates { state1, state2 } = ratchet_states;
 
     let mut session_queue = ctx.session_queue.lock().unwrap();
     let mut session_map = ctx.session_map.write().unwrap();
@@ -428,11 +423,11 @@ pub(crate) fn received_x1_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
     //    ...
     //    -> e, es, e1
     //    <- e, ee, ekem1, psk
-    if !(HANDSHAKE_HELLO_MIN_SIZE..=HANDSHAKE_HELLO_MAX_SIZE).contains(&x1.len()) {
+    if x1.len() != HANDSHAKE_HELLO_SIZE {
         return Err(fault!(InvalidPacket, true));
     }
 
-    if n[AES_GCM_NONCE_SIZE - 8..] != x1[x1.len() - 8..] {
+    if !secure_eq(&n[AES_GCM_NONCE_SIZE - 8..], &x1[x1.len() - 8..]) {
         return Err(fault!(FailedAuth, true));
     }
     let hmac = &mut C::Hmac::new();
@@ -462,35 +457,46 @@ pub(crate) fn received_x1_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
     let e1_end = j;
     i = k;
     // Process message pattern 1 payload.
-    let k = x1.len();
-    let j = k - AES_GCM_TAG_SIZE;
+    let j = i + RATCHET_SIZE + RATCHET_SIZE;
+    let k = j + AES_GCM_TAG_SIZE;
     let tag = x1[j..k].try_into().unwrap();
     if !noise.decrypt_and_hash_in_place(hash, to_nonce(PACKET_TYPE_HANDSHAKE_HELLO, 1), &mut x1[i..j], tag) {
         return Err(fault!(FailedAuth, true));
     }
+    debug_assert_eq!(k, x1.len());
 
+    let rf1 = &x1[i..i + RATCHET_SIZE];
+    let rf2 = &x1[i + RATCHET_SIZE..j];
     let mut lookup_data = None;
     let mut ratchet_state = None;
-    while i + RATCHET_SIZE <= j {
-        match app.restore_by_fingerprint((&x1[i..i + RATCHET_SIZE]).try_into().unwrap()) {
+    if !secure_eq(rf1, &[0u8; RATCHET_SIZE]) {
+        match app.restore_by_fingerprint(rf1.try_into().unwrap()) {
             Ok(None) => {}
             Ok(Some((rs, data))) => {
                 lookup_data = Some(data);
                 ratchet_state = Some(rs);
-                break;
             }
             Err(e) => return Err(ReceiveError::StorageError(e)),
         }
-        i += RATCHET_SIZE;
     }
-    let ratchet_state = if let Some(rs) = ratchet_state {
-        rs
-    } else {
-        if app.hello_requires_recognized_ratchet() {
+    if ratchet_state.is_none() {
+        if !secure_eq(rf2, &[0u8; RATCHET_SIZE]) {
+            match app.restore_by_fingerprint(rf1.try_into().unwrap()) {
+                Ok(None) => {}
+                Ok(Some((rs, data))) => {
+                    lookup_data = Some(data);
+                    ratchet_state = Some(rs);
+                }
+                Err(e) => return Err(ReceiveError::StorageError(e)),
+            }
+        }
+        if ratchet_state.is_none() && app.hello_requires_recognized_ratchet() {
             return Err(fault!(FailedAuth, true));
         }
-        RatchetState::empty()
-    };
+    }
+    // If we get to this point and haven't found a full ratchet state,
+    // set it to the empty ratchet state.
+    let ratchet_state = ratchet_state.unwrap_or_default();
 
     let mut hk_recv = Zeroizing::new([0u8; HASHLEN]);
     let mut hk_send = Zeroizing::new([0u8; HASHLEN]);
@@ -591,7 +597,7 @@ pub(crate) fn received_x2_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
         return Err(fault!(UnknownLocalKeyId, true, session));
     }
     let (_, c) = from_nonce(n);
-    if c >= COUNTER_WINDOW_MAX_SKIP_AHEAD || n[AES_GCM_NONCE_SIZE - 3..] != x2[x2.len() - 3..] {
+    if c >= COUNTER_WINDOW_MAX_SKIP_AHEAD || !secure_eq(&n[AES_GCM_NONCE_SIZE - 3..], &x2[x2.len() - 3..]) {
         return Err(fault!(FailedAuth, true, session));
     }
 
@@ -692,23 +698,27 @@ pub(crate) fn received_x2_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
 
         let new_ratchet_state = create_ratchet_state(hmac, &noise, chain_len);
 
-        let (ratchet_to_preserve, ratchet_to_delete) = if ratchet_i == 1 {
-            (Some(&state.ratchet_state1), state.ratchet_state2.as_ref())
+        let ratchet_to_preserve = if ratchet_i == 1 {
+            Some(&state.ratchet_state1)
         } else {
-            (state.ratchet_state2.as_ref(), Some(&state.ratchet_state1))
+            state.ratchet_state2.as_ref()
         };
-        app.save_ratchet_state(
+        let result = app.save_ratchet_state(
             &session.s_remote,
             &session.session_data,
-            RatchetUpdate {
-                state1: &new_ratchet_state,
-                state2: ratchet_to_preserve,
-                state1_was_just_added: true,
-                deleted_state1: ratchet_to_delete,
-                deleted_state2: None,
-            },
-        )
-        .map_err(|e| ReceiveError::StorageError(e))?;
+            CompareAndSwap::new(
+                &new_ratchet_state,
+                ratchet_to_preserve,
+                true,
+                &state.ratchet_state1,
+                state.ratchet_state2.as_ref(),
+                ratchet_i == 2,
+                ratchet_i == 1,
+            ),
+        );
+        if !result.map_err(ReceiveError::StorageError)? {
+            return Err(fault!(OutOfSequence, true, session, true));
+        }
 
         let mut kek_recv = Zeroizing::new([0u8; HASHLEN]);
         let mut kek_send = Zeroizing::new([0u8; HASHLEN]);
@@ -738,7 +748,7 @@ pub(crate) fn received_x2_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
                 a1
             } else {
                 // This return is unreachable.
-                return Err(fault!(FailedAuth, true, session));
+                return Err(fault!(FailedAuth, true, session, true));
             };
             state.beta = ZetaAutomata::A3(Box::new(StateA3 { identity: a1.identity.clone(), x3: x3.clone() }));
             resend_timer
@@ -860,7 +870,7 @@ pub(crate) fn received_x3_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
                 let mut should_warn_missing_ratchet = false;
 
                 if (zeta.ratchet_state != state1) & (Some(&zeta.ratchet_state) != state2.as_ref()) {
-                    if !responder_disallows_downgrade && zeta.ratchet_state.fingerprint().is_none() {
+                    if !responder_disallows_downgrade && zeta.ratchet_state.is_empty() {
                         should_warn_missing_ratchet = true;
                     } else {
                         if !responder_silently_rejects {
@@ -879,18 +889,14 @@ pub(crate) fn received_x3_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
                 noise.split(hmac, &mut nk_send, &mut nk_recv);
 
                 // We must make sure the ratchet key is saved before we transition.
-                app.save_ratchet_state(
+                let result = app.save_ratchet_state(
                     &s_remote,
                     &session_data,
-                    RatchetUpdate {
-                        state1: &new_ratchet_state,
-                        state2: None,
-                        state1_was_just_added: true,
-                        deleted_state1: Some(&state1),
-                        deleted_state2: state2.as_ref(),
-                    },
-                )
-                .map_err(|e| ReceiveError::StorageError(e))?;
+                    CompareAndSwap::new(&new_ratchet_state, None, true, &state1, state2.as_ref(), true, true),
+                );
+                if !result.map_err(ReceiveError::StorageError)? {
+                    return Err(fault!(OutOfSequence, true));
+                }
 
                 let (session, reduced_service_time) = {
                     let mut session_map = ctx.session_map.write().unwrap();
@@ -1006,18 +1012,25 @@ pub(crate) fn received_c1_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
     if is_other {
         if let ZetaAutomata::A3 { .. } | ZetaAutomata::R2 { .. } = &state.beta {
             if state.ratchet_state2.is_some() {
-                app.save_ratchet_state(
+                let result = app.save_ratchet_state(
                     &session.s_remote,
                     &session.session_data,
-                    RatchetUpdate {
-                        state1: &state.ratchet_state1,
-                        state2: None,
-                        state1_was_just_added: false,
-                        deleted_state1: state.ratchet_state2.as_ref(),
-                        deleted_state2: None,
-                    },
-                )
-                .map_err(|e| ReceiveError::StorageError(e))?;
+                    CompareAndSwap::new(
+                        &state.ratchet_state1,
+                        None,
+                        false,
+                        &state.ratchet_state1,
+                        state.ratchet_state2.as_ref(),
+                        false,
+                        true,
+                    ),
+                );
+                if !result.map_err(ReceiveError::StorageError)? {
+                    drop(state);
+                    drop(kex_lock);
+                    session.expire();
+                    return Err(fault!(OutOfSequence, true, session, true));
+                }
             }
             drop(state);
             let timeout_timer = {
@@ -1417,18 +1430,22 @@ pub(crate) fn received_k1_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
         k2.extend(tag);
 
         let new_ratchet_state = create_ratchet_state(hmac, &noise, state.ratchet_state1.chain_len);
-        app.save_ratchet_state(
+        let result = app.save_ratchet_state(
             &session.s_remote,
             &session.session_data,
-            RatchetUpdate {
-                state1: &new_ratchet_state,
-                state2: Some(&state.ratchet_state1),
-                state1_was_just_added: true,
-                deleted_state1: state.ratchet_state2.as_ref(),
-                deleted_state2: None,
-            },
-        )
-        .map_err(|e| ReceiveError::StorageError(e))?;
+            CompareAndSwap::new(
+                &new_ratchet_state,
+                Some(&state.ratchet_state1),
+                true,
+                &state.ratchet_state1,
+                state.ratchet_state2.as_ref(),
+                false,
+                true,
+            ),
+        );
+        if !result.map_err(ReceiveError::StorageError)? {
+            return Err(fault!(OutOfSequence, true, session, true));
+        }
 
         let mut kek_recv = Zeroizing::new([0u8; HASHLEN]);
         let mut kek_send = Zeroizing::new([0u8; HASHLEN]);
@@ -1536,18 +1553,22 @@ pub(crate) fn received_k2_trans<C: CryptoLayer, App: ApplicationLayer<C>>(
                 .ok_or_else(|| fault!(InvalidPacket, true, session, true))?;
 
             let new_ratchet_state = create_ratchet_state(hmac, &noise, state.ratchet_state1.chain_len);
-            app.save_ratchet_state(
+            let result = app.save_ratchet_state(
                 &session.s_remote,
                 &session.session_data,
-                RatchetUpdate {
-                    state1: &new_ratchet_state,
-                    state2: None,
-                    state1_was_just_added: true,
-                    deleted_state1: Some(&state.ratchet_state1),
-                    deleted_state2: state.ratchet_state2.as_ref(),
-                },
-            )
-            .map_err(|e| ReceiveError::StorageError(e))?;
+                CompareAndSwap::new(
+                    &new_ratchet_state,
+                    None,
+                    true,
+                    &state.ratchet_state1,
+                    state.ratchet_state2.as_ref(),
+                    true,
+                    true,
+                ),
+            );
+            if !result.map_err(ReceiveError::StorageError)? {
+                return Err(fault!(OutOfSequence, true, session, true));
+            }
 
             let mut kek_recv = Zeroizing::new([0u8; HASHLEN]);
             let mut kek_send = Zeroizing::new([0u8; HASHLEN]);
@@ -1630,15 +1651,6 @@ pub(crate) fn send_payload<C: CryptoLayer>(
     };
     let nonce = to_nonce(PACKET_TYPE_DATA, c);
 
-    let key = state.key_ref(false);
-    let kid_send = key.send.kid.ok_or(SessionNotEstablished)?.get().to_ne_bytes();
-    let mut cipher = key.nk.as_ref().ok_or(SessionNotEstablished)?.start_enc(&nonce);
-
-    debug_assert!(matches!(
-        &state.beta,
-        ZetaAutomata::S1 | ZetaAutomata::S2 | ZetaAutomata::R1 { .. } | ZetaAutomata::R2 { .. }
-    ));
-
     let payload_mtu = mtu - HEADER_SIZE;
     debug_assert!(payload_mtu >= 4);
     let tagged_payload_len = payload.len() + AES_GCM_TAG_SIZE;
@@ -1649,6 +1661,16 @@ pub(crate) fn send_payload<C: CryptoLayer>(
         return Err(DataTooLarge);
     }
 
+    debug_assert!(matches!(
+        &state.beta,
+        ZetaAutomata::S1 | ZetaAutomata::S2 | ZetaAutomata::R1 { .. } | ZetaAutomata::R2 { .. }
+    ));
+
+    let key = state.key_ref(false);
+    let kid_send = key.send.kid.ok_or(SessionNotEstablished)?.get().to_ne_bytes();
+    let cipher_pool = key.nk.as_ref().ok_or(SessionNotEstablished)?;
+    let mut cipher = cipher_pool.start_enc(&nonce);
+
     let mut header = [0u8; HEADER_SIZE];
     header[..KID_SIZE].copy_from_slice(&kid_send);
     header[FRAGMENT_COUNT_IDX] = fragment_count as u8;
@@ -1662,12 +1684,15 @@ pub(crate) fn send_payload<C: CryptoLayer>(
         mtu_sized_buffer[..HEADER_SIZE].copy_from_slice(&header);
         mtu_sized_buffer[FRAGMENT_NO_IDX] = fragment_no as u8;
         let fragment_start = &mut mtu_sized_buffer[HEADER_SIZE..HEADER_SIZE + fragment_len];
-        cipher.encrypt(&payload[i..j], fragment_start);
+        cipher_pool.encrypt(&mut cipher, &payload[i..j], fragment_start);
 
         let header_auth = &mut mtu_sized_buffer[HEADER_AUTH_START..HEADER_AUTH_END];
         state.hk_send.encrypt_in_place(header_auth.try_into().unwrap());
 
         if !send.send_frag(&mut mtu_sized_buffer[..HEADER_SIZE + fragment_len]) {
+            // We need to give the cipher back to the pool instead of dropping it,
+            // so it can do memory cleanup.
+            cipher_pool.finish_enc(cipher);
             return Ok(false);
         }
         i = j;
@@ -1680,8 +1705,9 @@ pub(crate) fn send_payload<C: CryptoLayer>(
     mtu_sized_buffer[..HEADER_SIZE].copy_from_slice(&header);
     mtu_sized_buffer[FRAGMENT_NO_IDX] = fragment_no as u8;
     let fragment_start = &mut mtu_sized_buffer[HEADER_SIZE..HEADER_SIZE + payload_rem];
-    cipher.encrypt(&payload[i..], fragment_start);
-    mtu_sized_buffer[HEADER_SIZE + payload_rem..HEADER_SIZE + fragment_len].copy_from_slice(&cipher.finish());
+    cipher_pool.encrypt(&mut cipher, &payload[i..], fragment_start);
+    mtu_sized_buffer[HEADER_SIZE + payload_rem..HEADER_SIZE + fragment_len]
+        .copy_from_slice(&cipher_pool.finish_enc(cipher));
 
     let header_auth = &mut mtu_sized_buffer[HEADER_AUTH_START..HEADER_AUTH_END];
     state.hk_send.encrypt_in_place(header_auth.try_into().unwrap());
@@ -1728,9 +1754,8 @@ pub(crate) fn receive_payload_in_place<C: CryptoLayer>(
         return Err(fault!(UnknownLocalKeyId, true, session));
     };
 
-    let mut cipher = specified_key
-        .ok_or_else(|| fault!(OutOfSequence, true, session))?
-        .start_dec(nonce);
+    let cipher_pool = specified_key.ok_or_else(|| fault!(OutOfSequence, true, session))?;
+    let mut cipher = cipher_pool.start_dec(nonce);
     let (_, c) = from_nonce(nonce);
 
     // NOTE: This only works because we check the size of every received fragment in the receive
@@ -1738,14 +1763,14 @@ pub(crate) fn receive_payload_in_place<C: CryptoLayer>(
     for i in 0..fragments.len() - 1 {
         let fragment = &mut fragments[i].as_mut()[HEADER_SIZE..];
         debug_assert!(fragment.len() >= AES_GCM_TAG_SIZE);
-        cipher.decrypt_in_place(fragment);
+        cipher_pool.decrypt_in_place(&mut cipher, fragment);
     }
     let fragment = &mut fragments[fragments.len() - 1].as_mut()[HEADER_SIZE..];
     debug_assert!(fragment.len() >= AES_GCM_TAG_SIZE);
     let tag_idx = fragment.len() - AES_GCM_TAG_SIZE;
-    cipher.decrypt_in_place(&mut fragment[..tag_idx]);
+    cipher_pool.decrypt_in_place(&mut cipher, &mut fragment[..tag_idx]);
 
-    if !cipher.finish((&fragment[tag_idx..]).try_into().unwrap()) {
+    if !cipher_pool.finish_dec(cipher, (&fragment[tag_idx..]).try_into().unwrap()) {
         return Err(fault!(FailedAuth, true, session));
     }
 

performance/src/zssp.rs

@@ -161,6 +161,63 @@ impl<C: CryptoLayer> Context<C> {
     /// * `identity` - Payload to be sent to Bob that contains the information necessary
     ///   for the upper protocol to authenticate and approve of Alice's identity.
     pub fn open<App: ApplicationLayer<C>>(
+        &self,
+        mut app: App,
+        send: impl Sender,
+        mtu: usize,
+        static_remote_key: C::PublicKey,
+        session_data: C::SessionData,
+        identity: &[u8],
+    ) -> Result<(Arc<Session<C>>, Option<i64>), OpenError> {
+        let ratchet_states = app
+            .restore_by_identity(&static_remote_key, &session_data, None)
+            .map_err(OpenError::StorageError)?
+            .unwrap_or_default();
+        self.open_with_ratchet(
+            app,
+            send,
+            mtu,
+            static_remote_key,
+            session_data,
+            identity,
+            ratchet_states,
+        )
+    }
+    /// Create a new session and send initialization packets to Bob, our remote peer.
+    /// This function will use the specified `ratchet_states` to connect to Bob, as opposed to
+    /// calling `app.restore_by_identity`. This can be used to open a session with a specific
+    /// one-time-password using `RatchetStates::new_otp_states()`, or in situations where it is
+    /// desireable to avoid having Alice call `app.restore_by_identity`. Keep in mind that
+    /// `save_ratchet_state` likely will eventually be called to delete this ratchet state.
+    ///
+    /// If using a one-time-password, a call to `app.initiator_disallows_downgrade` on the created
+    /// session **must** return `true`. Otherwise the remote peer would be able to avoid having
+    /// to demonstrate knowledge of the otp by requesting a ratchet downgrade.
+    /// Only after the created session is established may a call to
+    /// `app.initiator_disallows_downgrade` return false.
+    ///
+    /// The session will not be "established" right away, and so will not be able to send data to
+    /// the remote peer until they respond and finish the handshake. A `SessionEvent` variant of
+    /// `Established` will be returned by `Context::receive` when this session is able to send data.
+    ///
+    /// This function returns an `Option<i64>`, which can safely be ignored if not using
+    /// `Context::service_scheduled`. `Context::service_scheduled` contains documentation on how to
+    /// handle the return value.
+    ///
+    /// To prevent desync, when this function is called, no other open session with the same remote
+    /// peer must exist. Drop or call expire on any pre-existing sessions before calling.
+    ///
+    /// * `app` - Application layer instance
+    /// * `send` - Function to be called to send one or more initial packets to the remote being
+    ///   contacted
+    /// * `mtu` - MTU for initial packets
+    /// * `static_remote_key` - Remote side's static public NIST P-384 key
+    /// * `session_data` - Arbitrary data meaningful to the application to include with session
+    ///   object
+    /// * `identity` - Payload to be sent to Bob that contains the information necessary
+    ///   for the upper protocol to authenticate and approve of Alice's identity.
+    /// * `ratchet_states` - The set of ratchet states that Alice should use to connect to Bob.
+    pub fn open_with_ratchet<App: ApplicationLayer<C>>(
         &self,
         app: App,
         send: impl Sender,
@@ -168,6 +225,7 @@ impl<C: CryptoLayer> Context<C> {
         static_remote_key: C::PublicKey,
         session_data: C::SessionData,
         identity: &[u8],
+        ratchet_states: RatchetStates,
     ) -> Result<(Arc<Session<C>>, Option<i64>), OpenError> {
         mtu = mtu.max(MIN_TRANSPORT_MTU);
         if identity.len() > IDENTITY_MAX_SIZE {
@@ -180,6 +238,7 @@ impl<C: CryptoLayer> Context<C> {
             static_remote_key,
             session_data,
             identity,
+            ratchet_states,
             |packet, hk_send| {
                 send_with_fragmentation(send, mtu, packet, hk_send);
             },
@@ -496,7 +555,7 @@ impl<C: CryptoLayer> Context<C> {
                 return Err(fault!(InvalidPacket, true));
             }
 
-            let mut buffer = ArrayVec::<u8, HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE>::new();
+            let mut buffer = ArrayVec::<u8, HANDSHAKE_HELLO_CHALLENGE_SIZE>::new();
             let assembled_packet = if fragment_count > 1 {
                 let mut next_service_time = self.0.unassociated_defrag_cache.lock().unwrap().assemble(
                     &nonce,
@@ -528,9 +587,7 @@ impl<C: CryptoLayer> Context<C> {
             if packet_type == PACKET_TYPE_HANDSHAKE_HELLO {
                 log!(app, ReceivedRawX1);
 
-                if !(HANDSHAKE_HELLO_CHALLENGE_MIN_SIZE..=HANDSHAKE_HELLO_CHALLENGE_MAX_SIZE)
-                    .contains(&assembled_packet.len())
-                {
+                if HANDSHAKE_HELLO_CHALLENGE_SIZE != assembled_packet.len() {
                     return Err(fault!(InvalidPacket, true));
                 }
                 // Process recv challenge layer.

reference/Cargo.toml

@@ -3,7 +3,7 @@ authors = ["ZeroTier, Inc. <contact@zerotier.com>", "Adam Ierymenko <adam.ieryme
 edition = "2021"
 license = "MPL-2.0"
 name = "zssp-proto"
-version = "0.3.0"
+version = "0.4.0"
 
 [lib]
 name = "zssp_proto"

reference/examples/extensive_test.rs

@@ -6,6 +6,7 @@
  * https://www.zerotier.com/
  */
 
+use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::iter::ExactSizeIterator;
 use std::str::FromStr;
@@ -18,7 +19,7 @@ use rand_core::OsRng;
 use rand_core::RngCore;
 
 use zssp_proto::application::{
-    AcceptAction, ApplicationLayer, CryptoLayer, RatchetState, RatchetStates, RatchetUpdate, Settings, RATCHET_SIZE,
+    AcceptAction, ApplicationLayer, CompareAndSwap, CryptoLayer, RatchetState, RatchetStates, Settings, RATCHET_SIZE,
 };
 use zssp_proto::crypto::P384KeyPair;
 use zssp_proto::crypto_impl::{
@@ -105,21 +106,37 @@ impl ApplicationLayer<TestApplication> for &mut TestApplication {
         &mut self,
         remote_static_key: &P384CratePublicKey,
         session_data: &u128,
-        update_data: RatchetUpdate<'_>,
-    ) -> Result<(), std::io::Error> {
-        self.ratchets.peer_map.insert(*session_data, update_data.to_states());
+        update_data: CompareAndSwap<'_>,
+    ) -> Result<bool, std::io::Error> {
+        let mut ratchets = &mut self.ratchets;
+        match ratchets.peer_map.entry(*session_data) {
+            Entry::Occupied(mut entry) => {
+                if update_data.compare(&entry.get()) {
+                    entry.insert(update_data.to_new_states());
+                } else {
+                    return Ok(false);
+                }
+            }
+            Entry::Vacant(entry) => {
+                if update_data.cur_is_initial_states() {
+                    entry.insert(update_data.to_new_states());
+                } else {
+                    return Ok(false);
+                }
+            }
+        }
 
         if let Some(rf) = update_data.added_fingerprint() {
-            self.ratchets.rf_map.insert(*rf, update_data.state1.clone());
-            println!("[{}] new ratchet #{}", self.name, update_data.state1.chain_len());
+            ratchets.rf_map.insert(*rf, update_data.new_state1.clone());
+            println!("[{}] new ratchet #{}", self.name, update_data.new_state1.chain_len());
         }
         if let Some(rf) = update_data.deleted_fingerprint1() {
-            self.ratchets.rf_map.remove(rf);
+            ratchets.rf_map.remove(rf);
         }
         if let Some(rf) = update_data.deleted_fingerprint2() {
-            self.ratchets.rf_map.remove(rf);
+            ratchets.rf_map.remove(rf);
         }
-        Ok(())
+        Ok(true)
     }
 
     fn time(&mut self) -> i64 {
@@ -202,7 +219,8 @@ fn alice_main(
                             }
                         }
                     }
-                } else if OsRng.next_u32() | 1 > 0 {
+                } else if (OsRng.next_u32() & 1) > 0 {
+                    // There is a 50% chance that a packet will be re-ordered instead of dropped.
                     let _ = recursive_out.send(pkt);
                 }
             } else {
@@ -290,7 +308,8 @@ fn bob_main(
                         }
                     }
                 }
-            } else if OsRng.next_u32() | 1 > 0 {
+            } else if (OsRng.next_u32() & 1) > 0 {
+                // There is a 50% chance that a packet will be re-ordered instead of dropped.
                 let _ = recursive_out.try_send(pkt);
             }
         }
@@ -377,6 +396,16 @@ fn main() {
 }
 
 #[test]
-fn test_main() {
-    core(2, u32::MAX / 2)
+fn test_50() {
+    core(10, u32::MAX / 2)
+}
+
+#[test]
+fn test_75() {
+    core(10, u32::MAX / 4 * 3)
+}
+
+#[test]
+fn test_99() {
+    core(10, u32::MAX / 100 * 99)
 }

reference/examples/ping_pong.rs

@@ -12,7 +12,7 @@ use std::sync::Arc;
 use std::time::Instant;
 
 use zssp_proto::application::{
-    AcceptAction, ApplicationLayer, CryptoLayer, RatchetState, RatchetStates, RatchetUpdate, RATCHET_SIZE,
+    AcceptAction, ApplicationLayer, CompareAndSwap, CryptoLayer, RatchetState, RatchetStates, RATCHET_SIZE,
 };
 use zssp_proto::crypto::{rand_core::OsRng, P384KeyPair};
 use zssp_proto::crypto_impl::{
@@ -81,9 +81,9 @@ impl ApplicationLayer<MyApp> for &mut MyApp {
         &mut self,
         remote_static_key: &P384CratePublicKey,
         session_data: &(),
-        update_data: RatchetUpdate<'_>,
-    ) -> Result<(), std::io::Error> {
-        Ok(())
+        update_data: CompareAndSwap<'_>,
+    ) -> Result<bool, std::io::Error> {
+        Ok(true)
     }
 
     fn time(&mut self) -> i64 {

reference/src/application.rs

@@ -213,11 +213,16 @@ pub trait ApplicationLayer<C: CryptoLayer>: Sized {
         remote_static_key: &<C as CryptoLayer>::PublicKey,
         session_data: &<C as CryptoLayer>::SessionData,
     ) -> Result<Option<RatchetStates>, std::io::Error>;
-    /// Atomically commit the update specified by `update_data` to storage, or return an error if
-    /// the update could not be made.
-    /// The implementor is free to choose how to apply these updates to storage.
+    /// Atomically compare-and-swap (a.k.a. compare-exchange) `update` to storage.
     ///
-    /// If this returns `Err(IoError)`, the packet which triggered this function to be called will be
+    /// If `update.cur_state1` and `update.cur_state2` are currently in storage, they must
+    /// be swapped with `update.new_state1` and `update.new_state2`.
+    /// Otherwise, storage must remain unchanged.
+    ///
+    /// `Ok(true)` must be returned if the compare-and-swap was successful.
+    /// `Ok(false)` must be returned if comparison failed and the swap was cancelled.
+    ///
+    /// If this returns `Err`, the packet which triggered this function to be called will be
     /// dropped, and no session state will be mutated, preserving synchronization. The remote peer
     /// will eventually resend that packet and so this function will be called again.
     ///
@@ -229,12 +234,20 @@ pub trait ApplicationLayer<C: CryptoLayer>: Sized {
     /// This function may also save state to volatile storage, in which case all peers which connect
     /// to us will have to allow downgrade across the board.
     /// Otherwise, when we restart, we will not be allowed to reconnect.
+    ///
+    /// # Security
+    /// Implementations must not perform comparison operations (equals, less than, etc.) directly
+    /// on two ratchet keys. Instead, comparison operations must be performed indirectly
+    /// upon their ratchet fingerprints. If two ratchet states have the same ratchet fingerprint,
+    /// it should be assumed that they also have the same ratchet key.
+    ///
+    /// The implementations of `PartialEq` for `RatchetState` and `RatchetStates` do this by default.
     fn save_ratchet_state(
         &mut self,
         remote_static_key: &<C as CryptoLayer>::PublicKey,
         session_data: &<C as CryptoLayer>::SessionData,
-        update_data: RatchetUpdate<'_>,
-    ) -> Result<(), std::io::Error>;
+        update_data: CompareAndSwap<'_>,
+    ) -> Result<bool, std::io::Error>;
 
     /// Receives a stream of events that occur during an execution of ZSSP.
     /// These are provided for debugging, logging or metrics purposes, and must be used for

reference/src/context.rs

@@ -11,7 +11,7 @@ use crate::challenge::ChallengeContext;
 use crate::crypto::{AES_256_KEY_SIZE, AES_GCM_NONCE_SIZE};
 use crate::fragmentation::{send_with_fragmentation, DefragBuffer};
 use crate::proto::*;
-use crate::result::{byzantine_fault, OpenError, ReceiveError, ReceiveOk, SendError, SessionEvent};
+use crate::result::{fault, OpenError, ReceiveError, ReceiveOk, SendError, SessionEvent};
 use crate::zeta::*;
 #[cfg(feature = "logging")]
 use crate::LogEvent::*;
@@ -161,24 +161,25 @@ impl<C: CryptoLayer> Context<C> {
                                 if !matches!(&zeta.beta, ZetaAutomata::A1(_)) {
                                     // A resent handshake response from Bob may have arrived out of order,
                                     // after we already received one.
-                                    return Err(byzantine_fault!(OutOfSequence, false));
+                                    return Err(fault!(OutOfSequence, false));
                                 }
                                 if c >= COUNTER_WINDOW_MAX_SKIP_AHEAD {
-                                    return Err(byzantine_fault!(ExpiredCounter, true));
+                                    return Err(fault!(ExpiredCounter, true));
                                 }
                                 Ok(())
                             } else if PACKET_TYPE_USES_COUNTER_RANGE.contains(&p) {
                                 if !zeta.check_counter_window(c) {
+                                    println!("{} {} {}", c, p, kid_recv);
                                     // The counter window has finite memory and so will occasionally give
                                     // false positives on very out-of-order packets.
-                                    return Err(byzantine_fault!(ExpiredCounter, false));
+                                    return Err(fault!(ExpiredCounter, false));
                                 }
                                 Ok(())
                             } else if p == PACKET_TYPE_HANDSHAKE_COMPLETION {
                                 // The handshake completion packet could have been resent.
-                                return Err(byzantine_fault!(InvalidPacket, false));
+                                return Err(fault!(InvalidPacket, false));
                             } else {
-                                return Err(byzantine_fault!(InvalidPacket, true));
+                                return Err(fault!(InvalidPacket, true));
                             }
                         })?;
                 if let Some((pn, mut assembled_packet)) = result {
@@ -295,7 +296,7 @@ impl<C: CryptoLayer> Context<C> {
                             log!(app, DIsAuthClosedSession(&session));
                             SessionEvent::Rejected
                         }
-                        _ => return Err(byzantine_fault!(InvalidPacket, true)), // This is unreachable.
+                        _ => return Err(fault!(InvalidPacket, true)), // This is unreachable.
                     };
                     drop(zeta);
                     Ok(ReceiveOk::Session(session, ret))
@@ -315,7 +316,7 @@ impl<C: CryptoLayer> Context<C> {
                                 if p == PACKET_TYPE_HANDSHAKE_COMPLETION && c == 0 {
                                     Ok(())
                                 } else {
-                                    Err(byzantine_fault!(InvalidPacket, true))
+                                    Err(fault!(InvalidPacket, true))
                                 }
                             })?;
                     if let Some((_, assembled_packet)) = result {
@@ -351,7 +352,7 @@ impl<C: CryptoLayer> Context<C> {
                 } else {
                     // When sessions are added or dropped or packets arrive extremely delayed it is
                     // possible to receive no longer recognized key ids.
-                    Err(byzantine_fault!(UnknownLocalKeyId, false))
+                    Err(fault!(UnknownLocalKeyId, false))
                 }
             }
         } else {
@@ -365,7 +366,7 @@ impl<C: CryptoLayer> Context<C> {
                     if p == PACKET_TYPE_HANDSHAKE_HELLO || p == PACKET_TYPE_CHALLENGE {
                         Ok(())
                     } else {
-                        Err(byzantine_fault!(InvalidPacket, true))
+                        Err(fault!(InvalidPacket, true))
                     }
                 },
             )?;
@@ -394,7 +395,7 @@ impl<C: CryptoLayer> Context<C> {
                             None,
                         );
                         // If we issue a challenge the first hello packet will always fail.
-                        return Err(byzantine_fault!(FailedAuth, false));
+                        return Err(fault!(FailedAuth, false));
                     } else if let Ok(true) = result {
                         log!(app, X1SucceededChallenge);
                     }
@@ -423,7 +424,7 @@ impl<C: CryptoLayer> Context<C> {
                     log!(app, ReceivedRawChallenge);
                     // Process recv challenge layer.
                     if assembled_packet.len() != KID_SIZE + CHALLENGE_SIZE {
-                        return Err(byzantine_fault!(InvalidPacket, true));
+                        return Err(fault!(InvalidPacket, true));
                     }
                     if let Some(kid_recv) =
                         NonZeroU32::new(u32::from_be_bytes(assembled_packet[..KID_SIZE].try_into().unwrap()))
@@ -439,9 +440,9 @@ impl<C: CryptoLayer> Context<C> {
                             return Ok(ReceiveOk::Unassociated);
                         }
                     }
-                    Err(byzantine_fault!(UnknownLocalKeyId, true))
+                    Err(fault!(UnknownLocalKeyId, true))
                 } else {
-                    Err(byzantine_fault!(InvalidPacket, true))
+                    Err(fault!(InvalidPacket, true))
                 }
             } else {
                 Ok(ReceiveOk::Unassociated)

reference/src/fragmentation.rs

@@ -7,7 +7,7 @@ use zeroize::Zeroizing;
 use crate::application::CryptoLayer;
 use crate::crypto::{Aes256Prp, AES_256_KEY_SIZE};
 use crate::proto::*;
-use crate::result::{byzantine_fault, ReceiveError};
+use crate::result::{fault, ReceiveError};
 
 /// Corresponds to Figure 13 found in Section 6.
 fn create_fragment_header(
@@ -94,7 +94,7 @@ impl DefragBuffer {
     ) -> Result<Option<([u8; PACKET_NONCE_SIZE], Vec<u8>)>, ReceiveError> {
         use crate::result::FaultType::*;
         if raw_fragment.len() < MIN_PACKET_SIZE {
-            return Err(byzantine_fault!(InvalidPacket, true));
+            return Err(fault!(InvalidPacket, true));
         }
 
         if let Some(hk_recv) = self.hk_recv.as_ref() {
@@ -109,7 +109,7 @@ impl DefragBuffer {
         let fragment_no = raw_fragment[FRAGMENT_NO_IDX] as usize;
         let fragment_count = raw_fragment[FRAGMENT_COUNT_IDX] as usize;
         if fragment_no >= fragment_count || fragment_count > MAX_FRAGMENTS {
-            return Err(byzantine_fault!(InvalidPacket, true));
+            return Err(fault!(InvalidPacket, true));
         }
 
         let n = raw_fragment[PACKET_NONCE_START..HEADER_SIZE].try_into().unwrap();
@@ -125,7 +125,7 @@ impl DefragBuffer {
                     || raw_fragment.len() > buffer.fragment_max_size
                 {
                     // Some parts of the protocol can cause duplicate fragments to be sent.
-                    return Err(byzantine_fault!(InvalidPacket, false));
+                    return Err(fault!(InvalidPacket, false));
                 }
                 buffer.fragments[fragment_no] = Some(raw_fragment);
                 buffer.total += 1;