From df6e09872a847723adb230aade808a77e5a3e34a Mon Sep 17 00:00:00 2001
From: Sachymetsu <sachymetsu@tutamail.com>
Date: Sat, 4 Apr 2026 11:39:09 +0200
Subject: [PATCH] Clearer context naming, remove redundant trait

---
 sachy-crypto/src/lib.rs | 88 ++++++++++++++---------------------------
 1 file changed, 30 insertions(+), 58 deletions(-)

diff --git a/sachy-crypto/src/lib.rs b/sachy-crypto/src/lib.rs
index 3cff8f2..72c8175 100644
--- a/sachy-crypto/src/lib.rs
+++ b/sachy-crypto/src/lib.rs
@@ -1,6 +1,6 @@
 #![no_std]
 
-use core::ops::{AddAssign, BitXor};
+use core::ops::BitXor;
 
 use chacha20poly1305::{AeadInOut, ChaCha20Poly1305, KeyInit, aead};
 use dhkem::{
@@ -120,24 +120,6 @@ impl ServerHandshake {
     }
 }
 
-/// Low-level Transport implementation.
-///
-/// This trait provides a particular "flavor" of transport, as there are
-/// different ways the specifics of the construction can be implemented.
-trait TransportPrimitive<A>
-where
-    A: AeadInOut,
-{
-    /// Type used as the Trasnport counter.
-    type Counter: AddAssign + Copy + Default + Eq;
-
-    /// Value to use when incrementing the Transport counter (i.e. one)
-    const COUNTER_INCR: Self::Counter;
-
-    /// Maximum number of messages allowed to be sent via Transport
-    const COUNTER_MAX: Self::Counter;
-}
-
 pub struct SendingState<'a> {
     transport: &'a TransportState,
     counter: u64,
@@ -149,19 +131,19 @@ impl SendingState<'_> {
         msg: &mut dyn aead::Buffer,
         associated_data: &[u8],
     ) -> Result<(), ProtoError> {
-        if self.counter.ct_eq(&TransportState::COUNTER_MAX).into() {
+        if self.counter.ct_eq(&u64::MAX).into() {
             return Err(ProtoError);
         }
 
-        let counter = self.counter.to_be_bytes();
-
         self.transport.aead.encrypt_in_place(
-            &self.transport.mix_nonce(&counter, Role::Sender),
+            &self
+                .transport
+                .mix_nonce(&self.counter.to_be_bytes(), Role::Sender),
             associated_data,
             msg,
         )?;
 
-        self.counter = self.counter.wrapping_add(TransportState::COUNTER_INCR);
+        self.counter = self.counter.wrapping_add(1);
 
         Ok(())
     }
@@ -178,37 +160,29 @@ impl ReceivingState<'_> {
         msg: &mut dyn aead::Buffer,
         associated_data: &[u8],
     ) -> Result<(), ProtoError> {
-        if self.counter.ct_eq(&TransportState::COUNTER_MAX).into() {
+        if self.counter.ct_eq(&u64::MAX).into() {
             return Err(ProtoError);
         }
 
-        let counter = self.counter.to_be_bytes();
-
         self.transport.aead.decrypt_in_place(
-            &self.transport.mix_nonce(&counter, Role::Receiver),
+            &self
+                .transport
+                .mix_nonce(&self.counter.to_be_bytes(), Role::Receiver),
             associated_data,
             msg,
         )?;
 
-        self.counter = self.counter.wrapping_add(TransportState::COUNTER_INCR);
+        self.counter = self.counter.wrapping_add(1);
 
         Ok(())
     }
 }
 
-impl TransportPrimitive<ChaCha20Poly1305> for TransportState {
-    type Counter = u64;
-
-    const COUNTER_INCR: Self::Counter = 1;
-
-    const COUNTER_MAX: Self::Counter = u64::MAX;
-}
-
 #[repr(align(4))]
 pub struct TransportState {
     aead: ChaCha20Poly1305,
-    first: aead::Nonce<ChaCha20Poly1305>,
-    second: aead::Nonce<ChaCha20Poly1305>,
+    client: aead::Nonce<ChaCha20Poly1305>,
+    server: aead::Nonce<ChaCha20Poly1305>,
     role: Role,
 }
 
@@ -222,22 +196,20 @@ impl TransportState {
         let kdf = noncer.extract::<sha2::Sha256>(Some(psk));
 
         let mut key = [0u8; 32];
-
-        let mut first = aead::Nonce::<ChaCha20Poly1305>::default();
-        let mut second = aead::Nonce::<ChaCha20Poly1305>::default();
+        let mut client = aead::Nonce::<ChaCha20Poly1305>::default();
+        let mut server = aead::Nonce::<ChaCha20Poly1305>::default();
 
         kdf.expand(b"SachY-Crypt0", &mut key)
             .map_err(|_| ProtoError)?;
-
-        kdf.expand(b"N*nceOne", &mut first)
+        kdf.expand(b"NonceClient*", &mut client)
             .map_err(|_| ProtoError)?;
-        kdf.expand(b"N#nceTwo", &mut second)
+        kdf.expand(b"NonceServer#", &mut server)
             .map_err(|_| ProtoError)?;
 
         Ok(Self {
             aead: ChaCha20Poly1305::new(&key.into()),
-            first,
-            second,
+            client,
+            server,
             role,
         })
     }
@@ -262,14 +234,14 @@ impl TransportState {
 
         // Handshake ROLE XOR Transport ROLE selects either one or other nonce context,
         // (0) for first context, (1) for second context
-        // Sending:   Client ^ Sender   = 0 (select first)
-        // Receiving: Server ^ Receiver = 0 (select first)
-        // Sending:   Server ^ Sender   = 1 (select second)
-        // Receiving: Client ^ Receiver = 1 (select second)
+        // Sending:   Client ^ Sender   = 0 (select first/client context)
+        // Receiving: Server ^ Receiver = 0 (select first/client context)
+        // Sending:   Server ^ Sender   = 1 (select second/server context)
+        // Receiving: Client ^ Receiver = 1 (select second/server context)
         if context_select.ct_eq(&0).into() {
-            &self.first
+            &self.client
         } else {
-            &self.second
+            &self.server
         }
     }
 
@@ -403,11 +375,11 @@ mod tests {
         assert_eq!(&buffer1, &buffer2);
 
         // Both Transports have derived base nonces for each context.
-        // First context nonces will not match Second context nonces.
-        assert_eq!(alice.first, bob.first);
-        assert_eq!(alice.second, bob.second);
-        assert_ne!(alice.first, alice.second);
-        assert_ne!(bob.first, bob.second);
+        // Client context nonces will not match Server context nonces.
+        assert_eq!(alice.client, bob.client);
+        assert_eq!(alice.server, bob.server);
+        assert_ne!(alice.client, alice.server);
+        assert_ne!(bob.client, bob.server);
 
         Ok(())
     }
-- 
2.51.2