From e8c4aa7d25fa38e75f45e0133dc85f8855127be7 Mon Sep 17 00:00:00 2001
From: c0repwn3r <core@coredoes.dev>
Date: Thu, 15 Dec 2022 08:05:29 -0500
Subject: [PATCH] [noise] handshake initiator working better now ( forgot about
 i_i )

---
 .../src/noise/{handshake.rs => handshake_init.rs}   | 13 +++++++++++--
 quicktap/src/noise/mod.rs                           |  2 +-
 quicktap/src/qcrypto/tests.rs                       |  2 +-
 3 files changed, 13 insertions(+), 4 deletions(-)
 rename quicktap/src/noise/{handshake.rs => handshake_init.rs} (96%)

diff --git a/quicktap/src/noise/handshake.rs b/quicktap/src/noise/handshake_init.rs
similarity index 96%
rename from quicktap/src/noise/handshake.rs
rename to quicktap/src/noise/handshake_init.rs
index 9e1014f..1dea804 100644
--- a/quicktap/src/noise/handshake.rs
+++ b/quicktap/src/noise/handshake_init.rs
@@ -1,6 +1,8 @@
-//! `Noise_IKpsk2` handshake packets
+//! `Noise_IKpsk2` handshake initiator packets
 
 use std::fmt::{Debug, Formatter};
+use std::mem;
+use rand::{Rng, thread_rng};
 use tai64::Tai64N;
 use x25519_dalek::{EphemeralSecret, PublicKey, StaticSecret};
 use crate::noise::error::NoiseError;
@@ -83,6 +85,8 @@ pub fn handshake_init_to(session: &mut HandshakeState) -> Result<[u8; 148], Nois
     session.s_pub_i = PublicKey::from(&session.s_priv_me);
     session.s_pub_r = session.s_pub_them;
 
+    session.i_i = thread_rng().gen();
+
     let mut msg = HandshakeInitiatorRaw {
         sender: session.i_i.to_le_bytes(),
         ephemeral: [0u8; 32],
@@ -187,12 +191,15 @@ impl HandshakeInitiatorRaw {
 /// # Panics
 /// While containing unwraps, this function will never panic.
 #[allow(clippy::module_name_repetitions)]
-pub fn handshake_init_from(session: &mut HandshakeState, packet: [u8; 148]) -> Result<(), NoiseError> {
+pub fn handshake_init_from(session_orig: &mut HandshakeState, packet: [u8; 148]) -> Result<(), NoiseError> {
+    let mut session = session_orig.clone();
     session.s_pub_i = session.s_pub_them;
     session.s_pub_r = PublicKey::from(&session.s_priv_me);
 
     let mut msg = HandshakeInitiatorRaw::from_bytes(packet);
 
+    session.i_i = u32::from_le_bytes(msg.sender);
+
     session.c_i = HANDSHAKE_INITIATOR_CHAIN_KEY;
     session.h_i = HANDSHAKE_INITIATOR_CHAIN_KEY_HASH;
     session.h_i = qcrypto_hash_twice(&session.h_i, session.s_pub_r.as_bytes());
@@ -248,5 +255,7 @@ pub fn handshake_init_from(session: &mut HandshakeState, packet: [u8; 148]) -> R
         return Err(NoiseError::PacketUnauthenticated)
     }
 
+    mem::swap(session_orig, &mut session);
+
     Ok(())
 }
\ No newline at end of file
diff --git a/quicktap/src/noise/mod.rs b/quicktap/src/noise/mod.rs
index ecac1d8..26b8674 100644
--- a/quicktap/src/noise/mod.rs
+++ b/quicktap/src/noise/mod.rs
@@ -1,4 +1,4 @@
 //! Contains structs and functions for serializing and deserializing different packets in the Noise_IKpsk2 handshake and data frames
 
-pub mod handshake;
+pub mod handshake_init;
 pub mod error;
\ No newline at end of file
diff --git a/quicktap/src/qcrypto/tests.rs b/quicktap/src/qcrypto/tests.rs
index cc91489..620c8f4 100644
--- a/quicktap/src/qcrypto/tests.rs
+++ b/quicktap/src/qcrypto/tests.rs
@@ -1,7 +1,7 @@
 use hex_lit::hex;
 use rand::rngs::OsRng;
 use x25519_dalek::{EphemeralSecret, PublicKey, StaticSecret};
-use crate::noise::handshake::{handshake_init_from, handshake_init_to, HandshakeState};
+use crate::noise::handshake_init::{handshake_init_from, handshake_init_to, HandshakeState};
 use crate::qcrypto::aead::{qcrypto_aead, qcrypto_aead_decrypt, qcrypto_xaead, qcrypto_xaead_decrypt};
 use crate::qcrypto::{CONSTURCTION, IDENTIFIER};
 use crate::qcrypto::hashes::{qcrypto_hash, qcrypto_hash_twice, qcrypto_hmac, qcrypto_mac};