From 4360e0b48756ef6614c0e54af551fd06444b021e Mon Sep 17 00:00:00 2001
From: Adam Ierymenko <adam.ierymenko@zerotier.com>
Date: Fri, 3 Mar 2023 11:08:12 -0500
Subject: [PATCH] ZSSP API updates.

---
 zssp/src/main.rs |  6 +++---
 zssp/src/zssp.rs | 28 +++++++++++++++++-----------
 2 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/zssp/src/main.rs b/zssp/src/main.rs
index 52f84824b..3d223c5a1 100644
--- a/zssp/src/main.rs
+++ b/zssp/src/main.rs
@@ -18,7 +18,7 @@ struct TestApplication {
 }
 
 impl zssp::ApplicationLayer for TestApplication {
-    const REKEY_AFTER_USES: u64 = 350000;
+    const REKEY_AFTER_USES: u64 = 300000;
     const EXPIRE_AFTER_USES: u64 = 2147483648;
     const REKEY_AFTER_TIME_MS: i64 = 1000 * 60 * 60 * 2;
     const REKEY_AFTER_TIME_MS_MAX_JITTER: u32 = 1000 * 60 * 10;
@@ -90,7 +90,7 @@ fn alice_main(
                         TEST_MTU,
                         current_time,
                     ) {
-                        Ok(zssp::ReceiveResult::Ok) => {
+                        Ok(zssp::ReceiveResult::Ok(_)) => {
                             //println!("[alice] ok");
                         }
                         Ok(zssp::ReceiveResult::OkData(_, _)) => {
@@ -188,7 +188,7 @@ fn bob_main(
                     TEST_MTU,
                     current_time,
                 ) {
-                    Ok(zssp::ReceiveResult::Ok) => {
+                    Ok(zssp::ReceiveResult::Ok(_)) => {
                         //println!("[bob] ok");
                     }
                     Ok(zssp::ReceiveResult::OkData(s, data)) => {
diff --git a/zssp/src/zssp.rs b/zssp/src/zssp.rs
index 01985c22b..406cdaba4 100644
--- a/zssp/src/zssp.rs
+++ b/zssp/src/zssp.rs
@@ -57,8 +57,8 @@ struct SessionsById<Application: ApplicationLayer> {
 
 /// Result generated by the context packet receive function, with possible payloads.
 pub enum ReceiveResult<'b, Application: ApplicationLayer> {
-    /// Packet was valid, but no action needs to be taken.
-    Ok,
+    /// Packet was valid, but no action needs to be taken and no payload was delivered.
+    Ok(Option<Arc<Session<Application>>>),
 
     /// Packet was valid and a data payload was decoded and authenticated.
     OkData(Arc<Session<Application>>, &'b mut [u8]),
@@ -173,7 +173,6 @@ impl<Application: ApplicationLayer> Context<Application> {
 
         {
             let sessions = self.sessions.read().unwrap();
-
             for (id, s) in sessions.active.iter() {
                 if let Some(session) = s.upgrade() {
                     let state = session.state.read().unwrap();
@@ -390,7 +389,11 @@ impl<Application: ApplicationLayer> Context<Application> {
     ///
     /// Note that if check_accept_session accepts and returns Some() the session could still fail with
     /// receive() returning an error. A Some() return from check_accept_sesion doesn't guarantee
-    /// successful new session init.
+    /// successful new session init, only that the application has authorized it.
+    ///
+    /// Finally, note that the check_X() functions can end up getting called more than once for a given
+    /// incoming attempt from a given node if the network quality is poor. That's because the caller may
+    /// have to retransmit init packets causing repetition of parts of the exchange.
     ///
     /// * `app` - Interface to application using ZSSP
     /// * `check_allow_incoming_session` - Function to call to check whether an unidentified new session should be accepted
@@ -453,7 +456,8 @@ impl<Application: ApplicationLayer> Context<Application> {
                                 current_time,
                             );
                         } else {
-                            return Ok(ReceiveResult::Ok);
+                            drop(fragged);
+                            return Ok(ReceiveResult::Ok(Some(session)));
                         }
                     } else {
                         return self.process_complete_incoming_packet(
@@ -556,7 +560,7 @@ impl<Application: ApplicationLayer> Context<Application> {
             );
         }
 
-        return Ok(ReceiveResult::Ok);
+        return Ok(ReceiveResult::Ok(None));
     }
 
     fn process_complete_incoming_packet<
@@ -661,7 +665,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                             if packet_type == PACKET_TYPE_DATA {
                                 return Ok(ReceiveResult::OkData(session, &mut data_buf[..data_len]));
                             } else {
-                                return Ok(ReceiveResult::Ok);
+                                return Ok(ReceiveResult::Ok(Some(session)));
                             }
                         } else {
                             return Err(Error::OutOfSequence);
@@ -835,7 +839,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                         Some(&Aes::new(header_protection_key.as_bytes())),
                     )?;
 
-                    return Ok(ReceiveResult::Ok);
+                    return Ok(ReceiveResult::Ok(session));
                 }
 
                 PACKET_TYPE_BOB_NOISE_XK_ACK => {
@@ -995,7 +999,7 @@ impl<Application: ApplicationLayer> Context<Application> {
                                     Some(&session.header_protection_cipher),
                                 )?;
 
-                                return Ok(ReceiveResult::Ok);
+                                return Ok(ReceiveResult::Ok(Some(session)));
                             } else {
                                 return Err(Error::InvalidPacket);
                             }
@@ -1220,7 +1224,8 @@ impl<Application: ApplicationLayer> Context<Application> {
                                                 false,
                                             ));
 
-                                            return Ok(ReceiveResult::Ok);
+                                            drop(state);
+                                            return Ok(ReceiveResult::Ok(Some(session)));
                                         }
                                     }
                                     return Err(Error::FailedAuthentication);
@@ -1280,7 +1285,8 @@ impl<Application: ApplicationLayer> Context<Application> {
                                                 state.current_key = next_key_index; // this is an ACK so it's confirmed
                                                 state.current_offer = Offer::None;
 
-                                                return Ok(ReceiveResult::Ok);
+                                                drop(state);
+                                                return Ok(ReceiveResult::Ok(Some(session)));
                                             }
                                         }
                                     }