From 598e5acb158d9e66882fd0f64767c2bcbd7df83f Mon Sep 17 00:00:00 2001
From: Monica Moniot <mamoniot@protonmail.com>
Date: Fri, 13 Oct 2023 15:20:30 -0400
Subject: [PATCH] added helper functions

---
 performance/src/application.rs |  2 +-
 performance/src/zeta.rs        | 26 ++++++++++++++++++--------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/performance/src/application.rs b/performance/src/application.rs
index 81c80ff..8c595fb 100644
--- a/performance/src/application.rs
+++ b/performance/src/application.rs
@@ -186,7 +186,7 @@ pub trait ApplicationLayer<Crypto: CryptoLayer>: Sized {
     /// Function to accept sessions after final negotiation.
     /// The second argument is the identity that the remote peer sent us. The application
     /// must verify this identity is associated with the remote peer's static key.
-    /// 
+    ///
     /// To prevent desync, if this function specifies that we should connect, no other open session
     /// with the same remote peer must exist. Drop or call expire on any pre-existing sessions
     /// before returning.
diff --git a/performance/src/zeta.rs b/performance/src/zeta.rs
index 36d75ac..82083ad 100644
--- a/performance/src/zeta.rs
+++ b/performance/src/zeta.rs
@@ -869,7 +869,6 @@ pub(crate) fn received_x3_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
     let action = app.check_accept_session(&s_remote, &x3[identity_start..identity_end]);
     let responder_disallows_downgrade = action.responder_disallows_downgrade;
     let responder_silently_rejects = action.responder_silently_rejects;
-    let session_data = action.session_data;
     let create_reject = || {
         // We just used a counter with this key, but we are not storing
         // the fact we used it in memory. This is currently ok because the
@@ -882,7 +881,7 @@ pub(crate) fn received_x3_trans<Crypto: CryptoLayer, App: ApplicationLayer<Crypt
         set_header(&mut d, zeta.kid_send.get(), &nonce);
         d
     };
-    if let Some(session_data) = session_data {
+    if let Some(session_data) = action.session_data {
         let result = app.restore_by_identity(&s_remote, &session_data);
         match result {
             Ok(rss) => {
@@ -1832,9 +1831,8 @@ impl<Crypto: CryptoLayer> Session<Crypto> {
             }
         }
     }
-    ///
-    ///// The current ratchet state of this session.
-    ///// The returned values are sensitive and should be securely erased before being dropped.
+    /// The current ratchet state of this session.
+    /// The returned values are sensitive and should be securely erased before being dropped.
     pub fn ratchet_states(&self) -> RatchetStates {
         let state = self.state.read().unwrap();
         RatchetStates::new(state.ratchet_state1.clone(), state.ratchet_state2.clone())
@@ -1843,14 +1841,26 @@ impl<Crypto: CryptoLayer> Session<Crypto> {
     pub fn ratchet_count(&self) -> u64 {
         self.state.read().unwrap().ratchet_state1.chain_len
     }
-    /// Check whether this session is established.
+    /// Check whether this session is established and can send data.
+    /// Expired sessions will return false.
     pub fn established(&self) -> bool {
-        let state = self.state.read().unwrap();
         !matches!(
-            &state.beta,
+            &self.state.read().unwrap().beta,
             ZetaAutomata::A1(_) | ZetaAutomata::A3 { .. } | ZetaAutomata::Null
         )
     }
+    /// Check whether this session is still in the establishing phase of the handshake.
+    /// Expired sessions will return false.
+    pub fn establishing(&self) -> bool {
+        matches!(
+            &self.state.read().unwrap().beta,
+            ZetaAutomata::A1(_) | ZetaAutomata::A3 { .. }
+        )
+    }
+    /// Check whether this session is expired and can no longer be used.
+    pub fn expired(&self) -> bool {
+        matches!(&self.state.read().unwrap().beta, ZetaAutomata::Null)
+    }
     /// The static public key of the remote peer.
     pub fn remote_static_key(&self) -> &Crypto::PublicKey {
         &self.s_remote