clean up sstream

Author: cowlicks

src/error.rs

@@ -2,4 +2,12 @@
 pub enum Error {
     #[error("Error from `snow`: {0}")]
     Snow(#[from] snow::Error),
+    #[error("Error from `crypto_secretstream`: {0}")]
+    SecretStream(crypto_secretstream::aead::Error),
+}
+
+impl From<crypto_secretstream::aead::Error> for Error {
+    fn from(value: crypto_secretstream::aead::Error) -> Self {
+        Error::SecretStream(value)
+    }
 }

src/lib.rs

@@ -125,7 +125,7 @@ mod message;
 mod mqueue;
 mod noise;
 mod protocol;
-mod sstream;
+pub mod sstream;
 #[cfg(test)]
 mod test_utils;
 mod util;

src/noise.rs

@@ -514,7 +514,7 @@ mod test {
         panic!()
     }
     #[tokio::test]
-    async fn encrypted() -> Result<()> {
+    async fn encrypted_one() -> Result<()> {
         let hello = b"hello".to_vec();
         let world = b"world".to_vec();
         let (lc, rc) = create_result_connected();

src/sstream.rs

@@ -1,34 +1,42 @@
-#![allow(dead_code)]
+//! Create a sodium secret stream using the IK pattern
 use crypto_secretstream::{Header, Key, PullStream, PushStream, Tag};
 use rand::rngs::OsRng;
 use snow::{params::NoiseParams, Builder, HandshakeState};
+use std::fmt::Debug;
 
 use crate::{crypto::write_stream_id, Error};
 
 const PARAMS: &str = "Noise_IK_25519_ChaChaPoly_BLAKE2b";
 const STREAM_ID_LENGTH: usize = 32;
 const RAW_HEADER_MSG_LEN: usize = STREAM_ID_LENGTH + Header::BYTES;
 const SNOW_CIPHERKEYLEN: usize = 32;
+const PUBLIC_KEYLEN: usize = 32;
 
-struct InitiatorConfig {
-    remote_public_key: [u8; 32],
+#[derive(Debug)]
+/// Data for creating an initiator
+pub struct InitiatorConfig {
+    remote_public_key: [u8; PUBLIC_KEYLEN],
 }
+
 impl InitiatorConfig {
-    fn new(remote_public_key: [u8; 32]) -> Self {
+    /// Create a new [`InitiatorConfig`]
+    pub fn new(remote_public_key: [u8; PUBLIC_KEYLEN]) -> Self {
         Self { remote_public_key }
     }
 }
 
-struct IkSecretStream<Step> {
+#[derive(Debug)]
+/// Secret Stream protocol state
+pub struct IkSecretStream<Step> {
     is_initiator: bool,
     state: HandshakeState,
     msg_buf: [u8; 1024],
     _step: Step,
 }
 
 impl<Step> IkSecretStream<Step> {
-    // split handshake into (tx, rx)
-    fn split_handshake(&mut self) -> ([u8; SNOW_CIPHERKEYLEN], [u8; SNOW_CIPHERKEYLEN]) {
+    /// split handshake into (tx, rx)
+    pub fn split_handshake(&mut self) -> ([u8; SNOW_CIPHERKEYLEN], [u8; SNOW_CIPHERKEYLEN]) {
         let (a, b) = self.state.dangerously_get_raw_split();
         if self.is_initiator {
             (a, b)
@@ -38,27 +46,51 @@ impl<Step> IkSecretStream<Step> {
     }
 }
 
-struct InitiatorInitial;
-struct ResponderInitial;
-struct InitiatorInitialSent;
-struct ResponderReplied {
+/// Initial initiator state
+#[derive(Debug)]
+pub struct InitiatorInitial;
+
+/// Initial responder state
+#[derive(Debug)]
+pub struct ResponderInitial;
+
+/// Initiator has sent the first message
+#[derive(Debug)]
+pub struct InitiatorInitialSent;
+
+/// No decryptor yet
+pub struct EncryptorReady {
     rx: Key,
     pusher: PushStream,
     handshake_hash: Vec<u8>,
 }
 
-struct InitiatorEnc {
-    pusher: PushStream,
-    rx: Key,
-    handshake_hash: Vec<u8>,
-}
-struct Ready {
+/// Encryptor and decryptor
+pub struct Ready {
     puller: PullStream,
     pusher: PushStream,
 }
+impl Debug for EncryptorReady {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("InitiatorEnc")
+            .field("rx", &"Key(..)")
+            .field("pusher", &"PushStream(..)")
+            .field("handshake_hash", &self.handshake_hash)
+            .finish()
+    }
+}
+impl Debug for Ready {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("Ready")
+            .field("pusher", &"PushStream(..)")
+            .field("puller", &"PullStream(..)")
+            .finish()
+    }
+}
 
 impl IkSecretStream<InitiatorInitial> {
-    fn new_initiator(config: InitiatorConfig) -> Result<Self, Error> {
+    /// Create an initiator of a secret stream
+    pub fn new_initiator(config: InitiatorConfig) -> Result<Self, Error> {
         let params: NoiseParams = PARAMS.parse().expect("known to work");
         let kp = Builder::new(params.clone()).generate_keypair()?;
         let state = Builder::new(params.clone())
@@ -74,7 +106,8 @@ impl IkSecretStream<InitiatorInitial> {
         })
     }
 
-    fn make_first_msg(
+    /// Create the first message the initiator sends to the responder
+    pub fn make_first_msg(
         mut self,
         prologue: &[u8],
     ) -> Result<(IkSecretStream<InitiatorInitialSent>, Vec<u8>), Error> {
@@ -99,7 +132,8 @@ impl IkSecretStream<InitiatorInitial> {
 }
 
 impl IkSecretStream<ResponderInitial> {
-    fn new_responder(private: &[u8]) -> Result<Self, Error> {
+    /// Create a responder of a secret stream
+    pub fn new_responder(private: &[u8]) -> Result<Self, Error> {
         let params: NoiseParams = PARAMS.parse().expect("known to work");
         let state = Builder::new(params.clone())
             .local_private_key(&private)?
@@ -112,10 +146,11 @@ impl IkSecretStream<ResponderInitial> {
         })
     }
 
-    fn read_first_msg(
+    /// Read the first message of the protocol, create the next two messages to send to the initiator.
+    pub fn read_first_msg(
         mut self,
         msg: &[u8],
-    ) -> Result<(IkSecretStream<ResponderReplied>, [Vec<u8>; 2]), Error> {
+    ) -> Result<(IkSecretStream<EncryptorReady>, [Vec<u8>; 2]), Error> {
         self.state.read_message(msg, &mut self.msg_buf)?;
         let len = self.state.write_message(&[], &mut self.msg_buf)?;
         let hs_msg = self.msg_buf[..len].to_vec();
@@ -147,7 +182,7 @@ impl IkSecretStream<ResponderInitial> {
                 is_initiator,
                 state,
                 msg_buf,
-                _step: ResponderReplied {
+                _step: EncryptorReady {
                     rx: Key::from(rx),
                     pusher,
                     handshake_hash,
@@ -159,7 +194,11 @@ impl IkSecretStream<ResponderInitial> {
 }
 
 impl IkSecretStream<InitiatorInitialSent> {
-    fn receive_msg(mut self, msg: &[u8]) -> Result<(IkSecretStream<InitiatorEnc>, Vec<u8>), Error> {
+    /// Recieve the last message to complet the handsake
+    pub fn receive_msg(
+        mut self,
+        msg: &[u8],
+    ) -> Result<(IkSecretStream<EncryptorReady>, Vec<u8>), Error> {
         self.state.read_message(msg, &mut self.msg_buf)?;
 
         let (tx, rx) = self.split_handshake();
@@ -191,7 +230,7 @@ impl IkSecretStream<InitiatorInitialSent> {
                 is_initiator,
                 state,
                 msg_buf,
-                _step: InitiatorEnc {
+                _step: EncryptorReady {
                     pusher,
                     rx: Key::from(rx),
                     handshake_hash,
@@ -201,12 +240,13 @@ impl IkSecretStream<InitiatorInitialSent> {
         ))
     }
 }
-impl IkSecretStream<InitiatorEnc> {
-    fn receive_msg(self, msg: &[u8]) -> Result<IkSecretStream<Ready>, Error> {
+impl IkSecretStream<EncryptorReady> {
+    /// Recieve message the last message, used to set up the decryption stream
+    pub fn receive_msg(self, msg: &[u8]) -> Result<IkSecretStream<Ready>, Error> {
         let Self {
             is_initiator,
             _step:
-                InitiatorEnc {
+                EncryptorReady {
                     pusher,
                     rx,
                     handshake_hash,
@@ -217,11 +257,12 @@ impl IkSecretStream<InitiatorEnc> {
         // Read the received message from the other peer
         let mut expected_stream_id: [u8; STREAM_ID_LENGTH] = [0; STREAM_ID_LENGTH];
         write_stream_id(&handshake_hash, !is_initiator, &mut expected_stream_id);
-        if expected_stream_id != msg[..32] {
+        if expected_stream_id != msg[..STREAM_ID_LENGTH] {
             panic!()
         }
 
-        let header: [u8; 24] = msg[32..].try_into().expect("TODO wrong size");
+        let header: [u8; Header::BYTES] =
+            msg[STREAM_ID_LENGTH..].try_into().expect("TODO wrong size");
         let puller = PullStream::init(header.into(), &rx);
         Ok(IkSecretStream {
             is_initiator,
@@ -232,42 +273,18 @@ impl IkSecretStream<InitiatorEnc> {
     }
 }
 
-impl IkSecretStream<ResponderReplied> {
-    fn receive_msg(self, msg: &[u8]) -> Result<IkSecretStream<Ready>, Error> {
-        let Self {
-            is_initiator,
-            _step:
-                ResponderReplied {
-                    pusher,
-                    rx,
-                    handshake_hash,
-                },
-            state,
-            msg_buf,
-        } = self;
-        // Read the received message from the other peer
-        let mut expected_stream_id: [u8; STREAM_ID_LENGTH] = [0; STREAM_ID_LENGTH];
-        write_stream_id(&handshake_hash, !is_initiator, &mut expected_stream_id);
-        if expected_stream_id != msg[..32] {
-            panic!()
-        }
-
-        let header: [u8; 24] = msg[32..].try_into().expect("TODO wrong size");
-        let puller = PullStream::init(header.into(), &rx);
-        Ok(IkSecretStream {
-            is_initiator,
-            state,
-            msg_buf,
-            _step: Ready { puller, pusher },
-        })
-    }
-}
-
 impl IkSecretStream<Ready> {
-    fn push(&mut self, msg: &mut Vec<u8>, associated_data: &[u8], tag: Tag) -> Result<(), Error> {
+    /// Encrypt a message in place
+    pub fn push(
+        &mut self,
+        msg: &mut Vec<u8>,
+        associated_data: &[u8],
+        tag: Tag,
+    ) -> Result<(), Error> {
         Ok(self._step.pusher.push(msg, associated_data, tag)?)
     }
-    fn pull(&mut self, msg: &mut Vec<u8>, associated_data: &[u8]) -> Result<Tag, Error> {
+    /// Decrypt a message in place
+    pub fn pull(&mut self, msg: &mut Vec<u8>, associated_data: &[u8]) -> Result<Tag, Error> {
         Ok(self._step.puller.pull(msg, associated_data)?)
     }
 }

src/test_utils.rs

@@ -87,7 +87,7 @@ pub(crate) fn log() {
             .with_bracketed_fields(true)
             .with_indent_lines(true)
             .with_thread_ids(false)
-            .with_thread_names(true)
+            //.with_thread_names(true)
             //.with_span_modes(true)
             ;