Repository: commonwarexyz/monorepo
Pull request #624: [stream] Use `codec` crate
Author: login: BrendanChou
Created: 2025-03-17T22:09:03Z | Merged: 2025-03-18T19:13:28Z

Changes:

- [general]: Use `commonware_codec` instead of `prost`
- [`lib.rs`]: Remove unnecessary `Error` types
- [`connection.rs / handshake.rs`]: Combine `IncomingConnection` and `IncomingHandshake` into a single struct `IncomingConnection` for less type proliferation
- [`handshake.rs`]: `fn create_handshake(...)` changed to `handshake::Signed::sign(...)`, creating a `handshake::Signed` object (that can later be serialized) rather than a serialized handshake message
- [`handshake.rs`]: removed `handshake::Handshake`. Just exposed some fields on `handshake::Signed` to reduce the amount of type proliferation and reduce confusion
- [`handshake.rs`]: added `handshake::Signed` (the message that gets fully serialized to the wire) and `handshake::Info` (the information signed over) that both implement `Codec` and `SizedCodec`


Diff:
diff --git a/Cargo.lock b/Cargo.lock
index 13b5e0b7af..03d14561ef 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1114,14 +1114,13 @@ version = "0.0.43"
 dependencies = [
  "bytes",
  "chacha20poly1305",
+ "commonware-codec",
  "commonware-cryptography",
  "commonware-macros",
  "commonware-runtime",
  "commonware-utils",
  "criterion",
  "futures",
- "prost",
- "prost-build",
  "rand",
  "thiserror 2.0.12",
  "x25519-dalek",
diff --git a/p2p/src/simulated/network.rs b/p2p/src/simulated/network.rs
index 30fa3531ca..825aeaf39d 100644
--- a/p2p/src/simulated/network.rs
+++ b/p2p/src/simulated/network.rs
@@ -663,7 +663,7 @@ impl Link {
                 // Dial the peer and handshake by sending it the dialer's public key
                 let (mut sink, _) = context.dial(socket).await.unwrap();
                 if let Err(err) = send_frame(&mut sink, &dialer, max_size).await {
-                    error!(?err, "failed to send public key to dialee");
+                    error!(?err, "failed to send public key to listener");
                     return;
                 }
 
diff --git a/runtime/src/deterministic.rs b/runtime/src/deterministic.rs
index 326515620f..42447055bb 100644
--- a/runtime/src/deterministic.rs
+++ b/runtime/src/deterministic.rs
@@ -200,22 +200,22 @@ impl Auditor {
         *hash = hasher.finalize().to_vec();
     }
 
-    fn dial(&self, dialer: SocketAddr, dialee: SocketAddr) {
+    fn dial(&self, dialer: SocketAddr, listener: SocketAddr) {
         let mut hash = self.hash.lock().unwrap();
         let mut hasher = Sha256::new();
         hasher.update(&*hash);
         hasher.update(b"dial");
         hasher.update(dialer.to_string().as_bytes());
-        hasher.update(dialee.to_string().as_bytes());
+        hasher.update(listener.to_string().as_bytes());
         *hash = hasher.finalize().to_vec();
     }
 
-    fn accept(&self, dialee: SocketAddr, dialer: SocketAddr) {
+    fn accept(&self, listener: SocketAddr, dialer: SocketAddr) {
         let mut hash = self.hash.lock().unwrap();
         let mut hasher = Sha256::new();
         hasher.update(&*hash);
         hasher.update(b"accept");
-        hasher.update(dialee.to_string().as_bytes());
+        hasher.update(listener.to_string().as_bytes());
         hasher.update(dialer.to_string().as_bytes());
         *hash = hasher.finalize().to_vec();
     }
@@ -1085,13 +1085,13 @@ impl ReasonablyRealtime for Context {}
 
 type Dialable = mpsc::UnboundedSender<(
     SocketAddr,
-    mocks::Sink,   // Dialee -> Dialer
-    mocks::Stream, // Dialer -> Dialee
+    mocks::Sink,   // Listener -> Dialer
+    mocks::Stream, // Dialer -> Listener
 )>;
 
 /// Implementation of [`crate::Network`] for the `deterministic` runtime.
 ///
-/// When a dialer connects to a dialee, the dialee is given a new ephemeral port
+/// When a dialer connects to a listener, the listener is given a new ephemeral port
 /// from the range `32768..61000`. To keep things simple, it is not possible to
 /// bind to an ephemeral port. Likewise, if ports are not reused and when exhausted,
 /// the runtime will panic.
@@ -1152,7 +1152,7 @@ impl Networking {
         };
         self.auditor.dial(dialer, socket);
 
-        // Get dialee
+        // Get listener
         let mut sender = {
             let listeners = self.listeners.lock().unwrap();
             let sender = listeners.get(&socket).ok_or(Error::ConnectionFailed)?;
@@ -1161,9 +1161,9 @@ impl Networking {
 
         // Construct connection
         let (dialer_sender, dialer_receiver) = mocks::Channel::init();
-        let (dialee_sender, dialee_receiver) = mocks::Channel::init();
+        let (listener_sender, listener_receiver) = mocks::Channel::init();
         sender
-            .send((dialer, dialer_sender, dialee_receiver))
+            .send((dialer, dialer_sender, listener_receiver))
             .await
             .map_err(|_| Error::ConnectionFailed)?;
         Ok((
@@ -1172,7 +1172,7 @@ impl Networking {
                 auditor: self.auditor.clone(),
                 me: dialer,
                 peer: socket,
-                sender: dialee_sender,
+                sender: listener_sender,
             },
             Stream {
                 auditor: self.auditor.clone(),
diff --git a/stream/Cargo.toml b/stream/Cargo.toml
index c1b0fc6723..f2f8984381 100644
--- a/stream/Cargo.toml
+++ b/stream/Cargo.toml
@@ -11,6 +11,7 @@ repository = "https://github.com/commonwarexyz/monorepo/tree/main/stream"
 documentation = "https://docs.rs/commonware-stream"
 
 [dependencies]
+commonware-codec = { workspace = true }
 commonware-cryptography = { workspace = true }
 commonware-macros = { workspace = true }
 commonware-runtime = { workspace = true }
@@ -18,7 +19,6 @@ commonware-utils = { workspace = true }
 thiserror = { workspace = true }
 bytes = { workspace = true }
 futures = { workspace = true }
-prost = { workspace = true }
 rand = { workspace = true }
 chacha20poly1305 = "0.10"
 x25519-dalek = "2"
@@ -26,9 +26,6 @@ x25519-dalek = "2"
 [dev-dependencies]
 criterion = { workspace = true }
 
-[build-dependencies]
-prost-build = { workspace = true }
-
 [lib]
 bench = false
 
diff --git a/stream/build.rs b/stream/build.rs
deleted file mode 100644
index 201c49bedb..0000000000
--- a/stream/build.rs
+++ /dev/null
@@ -1,7 +0,0 @@
-use std::io::Result;
-fn main() -> Result<()> {
-    // Proto compilation rules for `public_key` dialect
-    let mut config = prost_build::Config::new();
-    config.compile_protos(&["src/public_key/wire.proto"], &["src/public_key/"])?;
-    Ok(())
-}
diff --git a/stream/src/lib.rs b/stream/src/lib.rs
index 83290df4f1..b3ba0d50ec 100644
--- a/stream/src/lib.rs
+++ b/stream/src/lib.rs
@@ -4,39 +4,30 @@ pub mod public_key;
 pub mod utils;
 
 use bytes::Bytes;
-use prost::DecodeError;
+use commonware_codec::Error as CodecError;
+use commonware_runtime::Error as RuntimeError;
 use std::future::Future;
 use thiserror::Error;
 
 /// Errors that can occur when interacting with a stream.
 #[derive(Error, Debug)]
 pub enum Error {
-    #[error("unexpected message")]
-    UnexpectedMessage,
     #[error("unable to decode: {0}")]
-    UnableToDecode(DecodeError),
-    #[error("invalid ephemeral public key")]
-    InvalidEphemeralPublicKey,
-    #[error("invalid channel public key")]
-    InvalidChannelPublicKey,
-    #[error("invalid peer public key")]
-    InvalidPeerPublicKey,
+    UnableToDecode(CodecError),
     #[error("handshake not for us")]
     HandshakeNotForUs,
     #[error("handshake timeout")]
     HandshakeTimeout,
-    #[error("missing signature")]
-    MissingSignature,
     #[error("invalid signature")]
     InvalidSignature,
     #[error("wrong peer")]
     WrongPeer,
     #[error("recv failed")]
-    RecvFailed,
+    RecvFailed(RuntimeError),
     #[error("recv too large: {0} bytes")]
     RecvTooLarge(usize),
     #[error("send failed")]
-    SendFailed,
+    SendFailed(RuntimeError),
     #[error("send zero size")]
     SendZeroSize,
     #[error("send too large: {0} bytes")]
diff --git a/stream/src/public_key/connection.rs b/stream/src/public_key/connection.rs
index e4bb2458df..0f39c06452 100644
--- a/stream/src/public_key/connection.rs
+++ b/stream/src/public_key/connection.rs
@@ -1,7 +1,4 @@
-use super::{
-    handshake::{create_handshake, Handshake, IncomingHandshake},
-    nonce, x25519, Config,
-};
+use super::{handshake, nonce, x25519, Config};
 use crate::{
     utils::codec::{recv_frame, send_frame},
     Error,
@@ -11,11 +8,13 @@ use chacha20poly1305::{
     aead::{Aead, KeyInit},
     ChaCha20Poly1305,
 };
+use commonware_codec::Codec;
 use commonware_cryptography::Scheme;
 use commonware_macros::select;
 use commonware_runtime::{Clock, Sink, Spawner, Stream};
 use commonware_utils::SystemTimeExt as _;
 use rand::{CryptoRng, Rng};
+use std::time::SystemTime;
 
 // When encrypting data, an encryption tag is appended to the ciphertext.
 // This constant represents the size of the encryption tag in bytes.
@@ -24,35 +23,57 @@ const ENCRYPTION_TAG_LENGTH: usize = 16;
 /// An incoming connection with a verified peer handshake.
 pub struct IncomingConnection<C: Scheme, Si: Sink, St: Stream> {
     config: Config<C>,
-    handshake: IncomingHandshake<Si, St, C>,
+    sink: Si,
+    stream: St,
+    deadline: SystemTime,
+    ephemeral_public_key: x25519::PublicKey,
+    peer_public_key: C::PublicKey,
 }
 
 impl<C: Scheme, Si: Sink, St: Stream> IncomingConnection<C, Si, St> {
-    /// Verify the handshake of an incoming connection.
-    pub async fn verify<R: Rng + CryptoRng + Spawner + Clock>(
-        context: &R,
+    pub async fn verify<E: Clock + Spawner>(
+        context: &E,
         config: Config<C>,
         sink: Si,
-        stream: St,
+        mut stream: St,
     ) -> Result<Self, Error> {
-        let handshake = IncomingHandshake::verify(
+        // Set handshake deadline
+        let deadline = context.current() + config.handshake_timeout;
+
+        // Wait for up to handshake timeout for response
+        let msg = select! {
+            _ = context.sleep_until(deadline) => { return Err(Error::HandshakeTimeout) },
+            result = recv_frame(&mut stream, config.max_message_size) => { result? },
+        };
+
+        // Verify handshake message from peer
+        let signed_handshake =
+            handshake::Signed::<C>::decode(msg).map_err(Error::UnableToDecode)?;
+        signed_handshake.verify(
             context,
             &config.crypto,
             &config.namespace,
-            config.max_message_size,
             config.synchrony_bound,
             config.max_handshake_age,
-            config.handshake_timeout,
+        )?;
+        Ok(Self {
+            config,
             sink,
             stream,
-        )
-        .await?;
-        Ok(Self { config, handshake })
+            deadline,
+            ephemeral_public_key: signed_handshake.ephemeral(),
+            peer_public_key: signed_handshake.signer(),
+        })
     }
 
     /// The public key of the peer attempting to connect.
     pub fn peer(&self) -> C::PublicKey {
-        self.handshake.peer_public_key.clone()
+        self.peer_public_key.clone()
+    }
+
+    /// The ephemeral public key of the peer attempting to connect.
+    pub fn ephemeral(&self) -> x25519::PublicKey {
+        self.ephemeral_public_key
     }
 }
 
@@ -101,17 +122,15 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
 
         // Generate shared secret
         let secret = x25519::new(&mut context);
-        let ephemeral = x25519_dalek::PublicKey::from(&secret);
 
         // Send handshake
         let timestamp = context.current().epoch_millis();
-        let msg = create_handshake(
+        let msg = handshake::Signed::sign(
             &mut config.crypto,
             &config.namespace,
-            timestamp,
-            peer.clone(),
-            ephemeral,
-        )?;
+            handshake::Info::<C>::new(peer.clone(), &secret, timestamp),
+        )
+        .encode();
 
         // Wait for up to handshake timeout to send
         select! {
@@ -119,7 +138,7 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
                 return Err(Error::HandshakeTimeout)
             },
             result = send_frame(&mut sink, &msg, config.max_message_size) => {
-                result.map_err(|_| Error::SendFailed)?;
+                result?;
             },
         }
 
@@ -129,27 +148,28 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
                 return Err(Error::HandshakeTimeout)
             },
             result = recv_frame(&mut stream, config.max_message_size) => {
-                result.map_err(|_| Error::RecvFailed)?
+                result?
             },
         };
 
         // Verify handshake message from peer
-        let handshake = Handshake::verify(
+        let signed_handshake =
+            handshake::Signed::<C>::decode(msg).map_err(Error::UnableToDecode)?;
+        signed_handshake.verify(
             &context,
             &config.crypto,
             &config.namespace,
             config.synchrony_bound,
             config.max_handshake_age,
-            msg,
         )?;
 
         // Ensure we connected to the right peer
-        if peer != handshake.peer_public_key {
+        if peer != signed_handshake.signer() {
             return Err(Error::WrongPeer);
         }
 
         // Create cipher
-        let shared_secret = secret.diffie_hellman(&handshake.ephemeral_public_key);
+        let shared_secret = secret.diffie_hellman(signed_handshake.ephemeral().as_ref());
         let cipher = ChaCha20Poly1305::new_from_slice(shared_secret.as_bytes())
             .map_err(|_| Error::CipherCreationFailed)?;
 
@@ -172,43 +192,47 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
         mut context: R,
         incoming: IncomingConnection<C, Si, St>,
     ) -> Result<Self, Error> {
-        // Generate shared secret
+        // Extract fields
+        let max_message_size = incoming.config.max_message_size;
+        let mut crypto = incoming.config.crypto;
+        let namespace = incoming.config.namespace;
+        let mut sink = incoming.sink;
+        let stream = incoming.stream;
+
+        // Generate personal secret
         let secret = x25519::new(&mut context);
-        let ephemeral = x25519_dalek::PublicKey::from(&secret);
 
         // Send handshake
-        let (mut handshake, mut config) = (incoming.handshake, incoming.config);
         let timestamp = context.current().epoch_millis();
-        let msg = create_handshake(
-            &mut config.crypto,
-            &config.namespace,
-            timestamp,
-            handshake.peer_public_key,
-            ephemeral,
-        )?;
+        let msg = handshake::Signed::sign(
+            &mut crypto,
+            &namespace,
+            handshake::Info::<C>::new(incoming.peer_public_key, &secret, timestamp),
+        )
+        .encode();
 
         // Wait for up to handshake timeout
         select! {
-            _ = context.sleep_until(handshake.deadline) => {
+            _ = context.sleep_until(incoming.deadline) => {
                 return Err(Error::HandshakeTimeout)
             },
-            result = send_frame(&mut handshake.sink, &msg, config.max_message_size) => {
-                result.map_err(|_| Error::SendFailed)?;
+            result = send_frame(&mut sink, &msg, max_message_size) => {
+                result?;
             },
         }
 
-        // Create cipher
-        let shared_secret = secret.diffie_hellman(&handshake.ephemeral_public_key);
+        // Create cipher based on the shared secret
+        let shared_secret = secret.diffie_hellman(incoming.ephemeral_public_key.as_ref());
         let cipher = ChaCha20Poly1305::new_from_slice(shared_secret.as_bytes())
             .map_err(|_| Error::CipherCreationFailed)?;
 
         // Track whether or not we are the dialer to ensure we send correctly formatted nonces.
         Ok(Connection {
             dialer: false,
-            sink: handshake.sink,
-            stream: handshake.stream,
+            sink,
+            stream,
             cipher,
-            max_message_size: config.max_message_size,
+            max_message_size,
         })
     }
 
@@ -294,9 +318,12 @@ impl<St: Stream> crate::Receiver for Receiver<St> {
 
 #[cfg(test)]
 mod tests {
+    use std::time::Duration;
+
     use super::*;
     use crate::{Receiver as _, Sender as _};
-    use commonware_runtime::{deterministic::Executor, mocks, Runner};
+    use commonware_cryptography::Ed25519;
+    use commonware_runtime::{deterministic::Executor, mocks, Metrics, Runner};
 
     #[test]
     fn test_decryption_failure() {
@@ -374,28 +401,204 @@ mod tests {
         let (executor, _, _) = Executor::default();
         executor.start(async move {
             let cipher = ChaCha20Poly1305::new(&[0u8; 32].into());
-            let message = b"hello world";
-            let max_message_size = message.len();
-
-            let (sink, stream) = mocks::Channel::init();
-            let is_dialer = false;
-            let mut sender = Sender {
-                cipher: cipher.clone(),
-                sink,
+            let max_message_size = 1024;
+
+            // Create channels
+            let (dialer_sink, listener_stream) = mocks::Channel::init();
+            let (listener_sink, dialer_stream) = mocks::Channel::init();
+
+            // Create dialer connection
+            let connection_dialer = Connection::from_preestablished(
+                true, // dialer
+                dialer_sink,
+                dialer_stream,
+                cipher.clone(),
                 max_message_size,
-                nonce: nonce::Info::new(is_dialer),
-            };
-            let mut receiver = Receiver {
+            );
+
+            // Create listener connection
+            let connection_listener = Connection::from_preestablished(
+                false, // listener
+                listener_sink,
+                listener_stream,
                 cipher,
-                stream,
                 max_message_size,
-                nonce: nonce::Info::new(is_dialer),
+            );
+
+            // Split into sender and receiver for both connections
+            let (mut dialer_sender, mut dialer_receiver) = connection_dialer.split();
+            let (mut listener_sender, mut listener_receiver) = connection_listener.split();
+
+            // Test 1: Send from dialer to listener
+            let msg1 = b"hello from dialer";
+            dialer_sender.send(msg1).await.unwrap();
+            let received1 = listener_receiver.receive().await.unwrap();
+            assert_eq!(received1, &msg1[..]);
+
+            // Test 2: Send from listener to dialer
+            let msg2 = b"hello from listener";
+            listener_sender.send(msg2).await.unwrap();
+            let received2 = dialer_receiver.receive().await.unwrap();
+            assert_eq!(received2, &msg2[..]);
+
+            // Test 3: Send multiple messages both ways
+            let messages_to_listener = vec![b"msg1", b"msg2", b"msg3"];
+            for msg in &messages_to_listener {
+                dialer_sender.send(*msg).await.unwrap();
+                let received = listener_receiver.receive().await.unwrap();
+                assert_eq!(received, &msg[..]);
+            }
+            let messages_to_dialer = vec![b"reply1", b"reply2", b"reply3"];
+            for msg in &messages_to_dialer {
+                listener_sender.send(*msg).await.unwrap();
+                let received = dialer_receiver.receive().await.unwrap();
+                assert_eq!(received, &msg[..]);
+            }
+        });
+    }
+    #[test]
+    fn test_full_connection_establishment_and_exchange() {
+        let (executor, context, _) = Executor::default();
+        executor.start(async move {
+            // Create cryptographic identities
+            let dialer_crypto = Ed25519::from_seed(0);
+            let listener_crypto = Ed25519::from_seed(1);
+
+            // Set up mock channels for transport simulation
+            let (dialer_sink, listener_stream) = mocks::Channel::init();
+            let (listener_sink, dialer_stream) = mocks::Channel::init();
+
+            // Configuration for dialer
+            let dialer_config = Config {
+                crypto: dialer_crypto.clone(),
+                namespace: b"test_namespace".to_vec(),
+                max_message_size: 1024,
+                synchrony_bound: Duration::from_secs(5),
+                max_handshake_age: Duration::from_secs(5),
+                handshake_timeout: Duration::from_secs(5),
             };
 
-            // Send data
-            sender.send(message).await.unwrap();
-            let data = receiver.receive().await.unwrap();
-            assert_eq!(data, &message[..]);
+            // Configuration for listener
+            let listener_config = Config {
+                crypto: listener_crypto.clone(),
+                namespace: b"test_namespace".to_vec(),
+                max_message_size: 1024,
+                synchrony_bound: Duration::from_secs(5),
+                max_handshake_age: Duration::from_secs(5),
+                handshake_timeout: Duration::from_secs(5),
+            };
+
+            // Spawn listener to handle incoming connection
+            let listener_handle = context.with_label("listener").spawn({
+                move |context| async move {
+                    let incoming = IncomingConnection::verify(
+                        &context,
+                        listener_config,
+                        listener_sink,
+                        listener_stream,
+                    )
+                    .await
+                    .unwrap();
+                    Connection::upgrade_listener(context, incoming)
+                        .await
+                        .unwrap()
+                }
+            });
+
+            // Dialer initiates the connection
+            let dialer_connection = Connection::upgrade_dialer(
+                context.clone(),
+                dialer_config,
+                dialer_sink,
+                dialer_stream,
+                listener_crypto.public_key(),
+            )
+            .await
+            .unwrap();
+
+            // Wait for listener connection to be established
+            let listener_connection = listener_handle.await.unwrap();
+
+            // Split connections into sender and receiver halves
+            let (mut dialer_sender, mut dialer_receiver) = dialer_connection.split();
+            let (mut listener_sender, mut listener_receiver) = listener_connection.split();
+
+            // Dialer sends to listener twice
+            let message1 = b"Hello from dialer";
+            dialer_sender.send(message1).await.unwrap();
+            dialer_sender.send(message1).await.unwrap();
+            let received = listener_receiver.receive().await.unwrap();
+            assert_eq!(&received[..], &message1[..]);
+            let received = listener_receiver.receive().await.unwrap();
+            assert_eq!(&received[..], &message1[..]);
+
+            // Listener sends to dialer twice
+            let message2 = b"Hello from listener";
+            listener_sender.send(message2).await.unwrap();
+            listener_sender.send(message2).await.unwrap();
+            let received = dialer_receiver.receive().await.unwrap();
+            assert_eq!(&received[..], &message2[..]);
+            let received = dialer_receiver.receive().await.unwrap();
+            assert_eq!(&received[..], &message2[..]);
+        });
+    }
+
+    #[test]
+    fn test_upgrade_dialer_wrong_peer() {
+        let (executor, context, _) = Executor::default();
+        executor.start(async move {
+            // Create cryptographic identities
+            let dialer_crypto = Ed25519::from_seed(0);
+            let expected_peer = Ed25519::from_seed(1).public_key();
+            let mut actual_peer = Ed25519::from_seed(2);
+
+            // Set up mock channels
+            let (dialer_sink, mut peer_stream) = mocks::Channel::init();
+            let (mut peer_sink, dialer_stream) = mocks::Channel::init();
+
+            // Dialer configuration
+            let dialer_config = Config {
+                crypto: dialer_crypto,
+                namespace: b"test_namespace".to_vec(),
+                max_message_size: 1024,
+                synchrony_bound: Duration::from_secs(5),
+                max_handshake_age: Duration::from_secs(5),
+                handshake_timeout: Duration::from_secs(5),
+            };
+            let peer_config = dialer_config.clone();
+
+            // Spawn a mock peer that responds with its own handshake without checking recipient
+            context.with_label("mock_peer").spawn({
+                move |mut context| async move {
+                    // Read the handshake from dialer
+                    let msg = recv_frame(&mut peer_stream, 1024).await.unwrap();
+                    let _ = handshake::Signed::<Ed25519>::decode(msg).unwrap(); // Simulate reading
+
+                    // Create and send own handshake
+                    let secret = x25519::new(&mut context);
+                    let timestamp = context.current().epoch_millis();
+                    let info =
+                        handshake::Info::new(peer_config.crypto.public_key(), &secret, timestamp);
+                    let signed_handshake =
+                        handshake::Signed::sign(&mut actual_peer, &peer_config.namespace, info);
+                    send_frame(&mut peer_sink, &signed_handshake.encode(), 1024)
+                        .await
+                        .unwrap();
+                }
+            });
+
+            // Attempt connection with expected peer key
+            let result = Connection::upgrade_dialer(
+                context,
+                dialer_config,
+                dialer_sink,
+                dialer_stream,
+                expected_peer,
+            )
+            .await;
+
+            // Verify the error
+            assert!(matches!(result, Err(Error::WrongPeer)));
         });
     }
 }
diff --git a/stream/src/public_key/handshake.rs b/stream/src/public_key/handshake.rs
index 3b53735d11..8449dced4d 100644
--- a/stream/src/public_key/handshake.rs
+++ b/stream/src/public_key/handshake.rs
@@ -1,73 +1,113 @@
-use super::{wire, x25519};
-use crate::{utils::codec::recv_frame, Error};
-use bytes::{BufMut, Bytes};
+use super::x25519;
+use crate::Error;
+use commonware_codec::{Codec, Error as CodecError, Reader, SizedCodec, Writer};
 use commonware_cryptography::Scheme;
-use commonware_macros::select;
-use commonware_runtime::{Clock, Sink, Spawner, Stream};
-use commonware_utils::{SizedSerialize, SystemTimeExt};
-use prost::Message;
-use std::time::{Duration, SystemTime};
-
-pub fn create_handshake<C: Scheme>(
-    crypto: &mut C,
-    namespace: &[u8],
+use commonware_runtime::Clock;
+use commonware_utils::SystemTimeExt;
+use std::time::Duration;
+
+/// Handshake information that is signed over by the sender.
+pub struct Info<C: Scheme> {
+    recipient: C::PublicKey,
+    ephemeral_public_key: x25519::PublicKey,
     timestamp: u64,
-    recipient_public_key: C::PublicKey,
-    ephemeral_public_key: x25519_dalek::PublicKey,
-) -> Result<Bytes, Error> {
-    // Sign their public key
-    let ephemeral_public_key_bytes = ephemeral_public_key.as_bytes();
-    let payload_len =
-        C::PublicKey::SERIALIZED_LEN + ephemeral_public_key_bytes.len() + u64::SERIALIZED_LEN;
-    let mut payload = Vec::with_capacity(payload_len);
-    payload.extend_from_slice(&recipient_public_key);
-    payload.extend_from_slice(ephemeral_public_key_bytes);
-    payload.put_u64(timestamp);
-    let signature = crypto.sign(Some(namespace), &payload);
-
-    // Send handshake
-    Ok(wire::Handshake {
-        recipient_public_key: recipient_public_key.to_vec(),
-        ephemeral_public_key: x25519::encode_public_key(ephemeral_public_key).to_vec(),
-        timestamp,
-        signature: Some(wire::Signature {
-            public_key: crypto.public_key().to_vec(),
-            signature: signature.to_vec(),
-        }),
+}
+
+impl<C: Scheme> Info<C> {
+    pub fn new(
+        recipient: C::PublicKey,
+        secret: &x25519_dalek::EphemeralSecret,
+        timestamp: u64,
+    ) -> Self {
+        Self {
+            recipient,
+            ephemeral_public_key: x25519::PublicKey::from_secret(secret),
+            timestamp,
+        }
     }
-    .encode_to_vec()
-    .into())
 }
 
-pub struct Handshake<C: Scheme> {
-    pub ephemeral_public_key: x25519_dalek::PublicKey,
-    pub peer_public_key: C::PublicKey,
+impl<C: Scheme> Codec for Info<C> {
+    fn write(&self, writer: &mut impl Writer) {
+        self.recipient.write(writer);
+        self.ephemeral_public_key.write(writer);
+        self.timestamp.write(writer);
+    }
+
+    fn read(reader: &mut impl Reader) -> Result<Self, CodecError> {
+        let recipient = C::PublicKey::read(reader)?;
+        let ephemeral_public_key = x25519::PublicKey::read(reader)?;
+        let timestamp = u64::read(reader)?;
+        Ok(Info {
+            recipient,
+            ephemeral_public_key,
+            timestamp,
+        })
+    }
+
+    fn len_encoded(&self) -> usize {
+        Self::LEN_ENCODED
+    }
+}
+
+impl<C: Scheme> SizedCodec for Info<C> {
+    const LEN_ENCODED: usize =
+        C::PublicKey::LEN_ENCODED + x25519::PublicKey::LEN_ENCODED + u64::LEN_ENCODED;
+}
+
+// Allows recipient to verify that the sender has the private key
+// of public key before sending any data.
+//
+// By requiring the server to have their public key signed, they prevent
+// a malicious peer from forwarding a handshake message from a previous
+// connection with public key (which could be used to convince the server
+// to start a useless handshake). Alternatively, we could require the
+// dialer to sign some random bytes provided by the server but this would
+// require the server to send a message to a peer before authorizing that
+// it should connect to them.
+pub struct Signed<C: Scheme> {
+    // The handshake info that was signed over
+    info: Info<C>,
+
+    // The public key of the sender
+    signer: C::PublicKey,
+
+    // The signature of the sender
+    signature: C::Signature,
 }
 
-impl<C: Scheme> Handshake<C> {
+impl<C: Scheme> Signed<C> {
+    pub fn sign(crypto: &mut C, namespace: &[u8], info: Info<C>) -> Self {
+        let signature = crypto.sign(Some(namespace), &info.encode());
+        Self {
+            info,
+            signer: crypto.public_key(),
+            signature,
+        }
+    }
+
+    pub fn signer(&self) -> C::PublicKey {
+        self.signer.clone()
+    }
+
+    pub fn ephemeral(&self) -> x25519::PublicKey {
+        self.info.ephemeral_public_key
+    }
+
     pub fn verify<E: Clock>(
+        &self,
         context: &E,
         crypto: &C,
         namespace: &[u8],
         synchrony_bound: Duration,
         max_handshake_age: Duration,
-        msg: Bytes,
-    ) -> Result<Self, Error> {
-        // Parse handshake message
-        let handshake = wire::Handshake::decode(msg).map_err(Error::UnableToDecode)?;
-
-        // Verify that ephemeral public key is valid
-        let ephemeral_public_key = x25519::decode_public_key(&handshake.ephemeral_public_key)
-            .map_err(|_| Error::InvalidEphemeralPublicKey)?;
-
+    ) -> Result<(), Error> {
         // Verify that the signature is for us
         //
         // If we didn't verify this, it would be trivial for any peer to impersonate another peer (even though
         // they would not be able to decrypt any messages from the shared secret). This would prevent us
         // from making a legitimate connection to the intended peer.
-        let our_public_key = C::PublicKey::try_from(handshake.recipient_public_key)
-            .map_err(|_| Error::InvalidChannelPublicKey)?;
-        if crypto.public_key() != our_public_key {
+        if crypto.public_key() != self.info.recipient {
             return Err(Error::HandshakeNotForUs);
         }
 
@@ -79,101 +119,65 @@ impl<C: Scheme> Handshake<C> {
         // to others (if an adversary recovered a handshake message could open a
         // connection to a peer first, peers only maintain one connection per peer).
         let current_timestamp = context.current().epoch();
-        let handshake_timestamp = Duration::from_millis(handshake.timestamp);
+        let handshake_timestamp = Duration::from_millis(self.info.timestamp);
         if handshake_timestamp + max_handshake_age < current_timestamp {
-            return Err(Error::InvalidTimestampOld(handshake.timestamp));
+            return Err(Error::InvalidTimestampOld(self.info.timestamp));
         }
         if handshake_timestamp > current_timestamp + synchrony_bound {
-            return Err(Error::InvalidTimestampFuture(handshake.timestamp));
+            return Err(Error::InvalidTimestampFuture(self.info.timestamp));
         }
 
-        // Get signature from peer
-        let signature = handshake.signature.ok_or(Error::MissingSignature)?;
-        let public_key: C::PublicKey = C::PublicKey::try_from(signature.public_key)
-            .map_err(|_| Error::InvalidPeerPublicKey)?;
-
-        // Construct signing payload (ephemeral public key + my public key + timestamp)
-        let payload_len = C::PublicKey::SERIALIZED_LEN
-            + handshake.ephemeral_public_key.len()
-            + u64::SERIALIZED_LEN;
-        let mut payload = Vec::with_capacity(payload_len);
-        payload.extend_from_slice(&our_public_key);
-        payload.extend_from_slice(&handshake.ephemeral_public_key);
-        payload.put_u64(handshake.timestamp);
-
         // Verify signature
-        let signature =
-            C::Signature::try_from(signature.signature).map_err(|_| Error::InvalidSignature)?;
-        if !C::verify(Some(namespace), &payload, &public_key, &signature) {
+        if !C::verify(
+            Some(namespace),
+            &self.info.encode(),
+            &self.signer,
+            &self.signature,
+        ) {
             return Err(Error::InvalidSignature);
         }
-        Ok(Self {
-            ephemeral_public_key,
-            peer_public_key: public_key,
-        })
+        Ok(())
     }
 }
 
-pub struct IncomingHandshake<Si: Sink, St: Stream, C: Scheme> {
-    pub sink: Si,
-    pub stream: St,
-    pub deadline: SystemTime,
-    pub ephemeral_public_key: x25519_dalek::PublicKey,
-    pub peer_public_key: C::PublicKey,
-}
+impl<C: Scheme> Codec for Signed<C> {
+    fn write(&self, writer: &mut impl Writer) {
+        self.info.write(writer);
+        self.signer.write(writer);
+        self.signature.write(writer);
+    }
 
-impl<Si: Sink, St: Stream, C: Scheme> IncomingHandshake<Si, St, C> {
-    #[allow(clippy::too_many_arguments)]
-    pub async fn verify<E: Clock + Spawner>(
-        context: &E,
-        crypto: &C,
-        namespace: &[u8],
-        max_message_size: usize,
-        synchrony_bound: Duration,
-        max_handshake_age: Duration,
-        handshake_timeout: Duration,
-        sink: Si,
-        mut stream: St,
-    ) -> Result<Self, Error> {
-        // Set handshake deadline
-        let deadline = context.current() + handshake_timeout;
-
-        // Wait for up to handshake timeout for response
-        let msg = select! {
-            _ = context.sleep_until(deadline) => {
-                return Err(Error::HandshakeTimeout);
-            },
-            result = recv_frame(&mut stream, max_message_size) => {
-                result.map_err(|_| Error::RecvFailed)?
-            },
-        };
-
-        // Verify handshake message from peer
-        let handshake = Handshake::verify(
-            context,
-            crypto,
-            namespace,
-            synchrony_bound,
-            max_handshake_age,
-            msg,
-        )?;
+    fn read(reader: &mut impl Reader) -> Result<Self, CodecError> {
+        let info = Info::<C>::read(reader)?;
+        let signer = C::PublicKey::read(reader)?;
+        let signature = C::Signature::read(reader)?;
         Ok(Self {
-            sink,
-            stream,
-            deadline,
-            ephemeral_public_key: handshake.ephemeral_public_key,
-            peer_public_key: handshake.peer_public_key,
+            info,
+            signer,
+            signature,
         })
     }
+
+    fn len_encoded(&self) -> usize {
+        Self::LEN_ENCODED
+    }
+}
+
+impl<C: Scheme> SizedCodec for Signed<C> {
+    const LEN_ENCODED: usize =
+        Info::<C>::LEN_ENCODED + C::PublicKey::LEN_ENCODED + C::Signature::LEN_ENCODED;
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::utils::codec::send_frame;
+    use crate::{
+        public_key::{Config, IncomingConnection},
+        utils::codec::send_frame,
+    };
     use commonware_cryptography::{Ed25519, Scheme};
-    use commonware_runtime::{deterministic::Executor, mocks, Metrics, Runner};
-    use x25519_dalek::PublicKey;
+    use commonware_runtime::{deterministic::Executor, mocks, Metrics, Runner, Spawner};
+    use x25519::PublicKey;
 
     const TEST_NAMESPACE: &[u8] = b"test_namespace";
     const ONE_MEGABYTE: usize = 1024 * 1024;
@@ -186,67 +190,56 @@ mod tests {
             // Create participants
             let mut sender = Ed25519::from_seed(0);
             let recipient = Ed25519::from_seed(1);
-            let ephemeral_public_key = PublicKey::from([3u8; 32]);
+            let ephemeral_public_key = PublicKey::from_bytes([3u8; 32]);
 
             // Create handshake message
-            let epoch_millis = context.current().epoch_millis();
-            let handshake_bytes = create_handshake(
+            let timestamp = context.current().epoch_millis();
+            let handshake = Signed::sign(
                 &mut sender,
                 TEST_NAMESPACE,
-                epoch_millis,
-                recipient.public_key(),
-                ephemeral_public_key,
-            )
-            .unwrap();
+                Info {
+                    timestamp,
+                    recipient: recipient.public_key(),
+                    ephemeral_public_key,
+                },
+            );
 
             // Decode the handshake message
-            let handshake = wire::Handshake::decode(handshake_bytes.clone())
-                .expect("failed to decode handshake");
+            let handshake =
+                Signed::<Ed25519>::decode(handshake.encode()).expect("failed to decode handshake");
 
             // Verify the timestamp
             let synchrony_bound = Duration::from_secs(5);
             let max_handshake_age = Duration::from_secs(5);
-            let handshake_timestamp = Duration::from_millis(handshake.timestamp);
-            let current_timestamp = Duration::from_millis(epoch_millis);
+            let handshake_timestamp = Duration::from_millis(handshake.info.timestamp);
+            let current_timestamp = Duration::from_millis(timestamp);
             assert!(handshake_timestamp <= current_timestamp + synchrony_bound);
             assert!(handshake_timestamp + max_handshake_age >= current_timestamp);
-            let handshake_recipient_public_key =
-                <Ed25519 as Scheme>::PublicKey::try_from(&handshake.recipient_public_key).unwrap();
-            let handshake_signature =
-                <Ed25519 as Scheme>::Signature::try_from(&handshake.signature.unwrap().signature)
-                    .unwrap();
 
             // Verify the signature
-            assert_eq!(handshake_recipient_public_key, recipient.public_key());
-            assert_eq!(
-                handshake.ephemeral_public_key,
-                x25519::encode_public_key(ephemeral_public_key)
-            );
-            let mut payload = Vec::new();
-            payload.extend_from_slice(&handshake.recipient_public_key);
-            payload.extend_from_slice(&handshake.ephemeral_public_key);
-            payload.put_u64(handshake.timestamp);
+            assert_eq!(handshake.info.recipient, recipient.public_key());
+            assert_eq!(handshake.info.ephemeral_public_key, ephemeral_public_key,);
 
             // Verify signature
             assert!(Ed25519::verify(
                 Some(TEST_NAMESPACE),
-                &payload,
+                &handshake.info.encode(),
                 &sender.public_key(),
-                &handshake_signature,
+                &handshake.signature,
             ));
 
             // Verify using the handshake struct
-            let handshake = Handshake::verify(
-                &context,
-                &recipient,
-                TEST_NAMESPACE,
-                synchrony_bound,
-                max_handshake_age,
-                handshake_bytes,
-            )
-            .unwrap();
-            assert_eq!(handshake.peer_public_key, sender.public_key());
-            assert_eq!(handshake.ephemeral_public_key, ephemeral_public_key);
+            handshake
+                .verify(
+                    &context,
+                    &recipient,
+                    TEST_NAMESPACE,
+                    synchrony_bound,
+                    max_handshake_age,
+                )
+                .unwrap();
+            assert_eq!(handshake.signer, sender.public_key());
+            assert_eq!(handshake.info.ephemeral_public_key, ephemeral_public_key);
         });
     }
 
@@ -258,17 +251,18 @@ mod tests {
             // Create participants
             let mut sender = Ed25519::from_seed(0);
             let recipient = Ed25519::from_seed(1);
-            let ephemeral_public_key = PublicKey::from([3u8; 32]);
+            let ephemeral_public_key = PublicKey::from_bytes([3u8; 32]);
 
             // Create handshake message
-            let handshake_bytes = create_handshake(
+            let handshake = Signed::sign(
                 &mut sender,
                 TEST_NAMESPACE,
-                0, // timestamp
-                recipient.public_key(),
-                ephemeral_public_key,
-            )
-            .unwrap();
+                Info {
+                    timestamp: 0,
+                    recipient: recipient.public_key(),
+                    ephemeral_public_key,
+                },
+            );
 
             // Setup a mock sink and stream
             let (sink, _) = mocks::Channel::init();
@@ -276,29 +270,27 @@ mod tests {
 
             // Send message over stream
             context.with_label("stream_sender").spawn(|_| async move {
-                send_frame(&mut stream_sender, &handshake_bytes, ONE_MEGABYTE)
+                send_frame(&mut stream_sender, &handshake.encode(), ONE_MEGABYTE)
                     .await
                     .unwrap();
             });
 
             // Call the verify function
-            let result = IncomingHandshake::verify(
-                &context,
-                &recipient,
-                TEST_NAMESPACE,
-                ONE_MEGABYTE,
-                Duration::from_secs(5),
-                Duration::from_secs(5),
-                Duration::from_secs(5),
-                sink,
-                stream,
-            )
-            .await
-            .unwrap();
+            let config = Config {
+                crypto: recipient.clone(),
+                namespace: TEST_NAMESPACE.to_vec(),
+                max_message_size: ONE_MEGABYTE,
+                synchrony_bound: Duration::from_secs(5),
+                max_handshake_age: Duration::from_secs(5),
+                handshake_timeout: Duration::from_secs(5),
+            };
+            let result = IncomingConnection::verify(&context, config, sink, stream)
+                .await
+                .unwrap();
 
             // Assert that the result is expected
-            assert_eq!(result.peer_public_key, sender.public_key());
-            assert_eq!(result.ephemeral_public_key, ephemeral_public_key);
+            assert_eq!(result.peer(), sender.public_key());
+            assert_eq!(result.ephemeral(), ephemeral_public_key);
         });
     }
 
@@ -309,17 +301,18 @@ mod tests {
         executor.start(async move {
             // Create participants
             let mut sender = Ed25519::from_seed(0);
-            let ephemeral_public_key = PublicKey::from([3u8; 32]);
+            let ephemeral_public_key = PublicKey::from_bytes([3u8; 32]);
 
             // Create handshake message
-            let handshake_bytes = create_handshake(
+            let handshake = Signed::sign(
                 &mut sender,
                 TEST_NAMESPACE,
-                0, // timestamp
-                Ed25519::from_seed(1).public_key(),
-                ephemeral_public_key,
-            )
-            .unwrap();
+                Info {
+                    timestamp: 0,
+                    recipient: Ed25519::from_seed(1).public_key(),
+                    ephemeral_public_key,
+                },
+            );
 
             // Setup a mock sink and stream
             let (sink, _) = mocks::Channel::init();
@@ -327,24 +320,21 @@ mod tests {
 
             // Send message over stream
             context.with_label("stream_sender").spawn(|_| async move {
-                send_frame(&mut stream_sender, &handshake_bytes, ONE_MEGABYTE)
+                send_frame(&mut stream_sender, &handshake.encode(), ONE_MEGABYTE)
                     .await
                     .unwrap();
             });
 
             // Call the verify function
-            let result = IncomingHandshake::verify(
-                &context,
-                &Ed25519::from_seed(2),
-                TEST_NAMESPACE,
-                ONE_MEGABYTE,
-                Duration::from_secs(5),
-                Duration::from_secs(5),
-                Duration::from_secs(5),
-                sink,
-                stream,
-            )
-            .await;
+            let config = Config {
+                crypto: Ed25519::from_seed(2),
+                namespace: TEST_NAMESPACE.to_vec(),
+                max_message_size: ONE_MEGABYTE,
+                synchrony_bound: Duration::from_secs(5),
+                max_handshake_age: Duration::from_secs(5),
+                handshake_timeout: Duration::from_secs(5),
+            };
+            let result = IncomingConnection::verify(&context, config, sink, stream).await;
 
             // Assert that the result is an error
             assert!(matches!(result, Err(Error::HandshakeNotForUs)));
@@ -368,18 +358,15 @@ mod tests {
             });
 
             // Call the verify function
-            let result = IncomingHandshake::verify(
-                &context,
-                &Ed25519::from_seed(0),
-                TEST_NAMESPACE,
-                ONE_MEGABYTE,
-                Duration::from_secs(1),
-                Duration::from_secs(1),
-                Duration::from_secs(1),
-                sink,
-                stream,
-            )
-            .await;
+            let config = Config {
+                crypto: Ed25519::from_seed(0),
+                namespace: TEST_NAMESPACE.to_vec(),
+                max_message_size: ONE_MEGABYTE,
+                synchrony_bound: Duration::from_secs(1),
+                max_handshake_age: Duration::from_secs(1),
+                handshake_timeout: Duration::from_secs(1),
+            };
+            let result = IncomingConnection::verify(&context, config, sink, stream).await;
 
             // Assert that the result is an error
             assert!(matches!(result, Err(Error::UnableToDecode(_))));
@@ -388,183 +375,58 @@ mod tests {
 
     #[test]
     fn test_incoming_handshake_verify_timeout() {
-        // Initialize context
         let (executor, context, _) = Executor::default();
         executor.start(async move {
-            // Create participants
-            let mut sender = Ed25519::from_seed(0);
-            let recipient = Ed25519::from_seed(1);
-            let ephemeral_public_key = PublicKey::from([3u8; 32]);
-
             // Setup a mock sink and stream
             let (sink, _) = mocks::Channel::init();
-            let (mut stream_sender, stream) = mocks::Channel::init();
-
-            // Accept connections but do nothing
-            context.with_label("stream_sender").spawn({
-                let recipient = recipient.clone();
-                move |context| async move {
-                    context.sleep(Duration::from_secs(10)).await;
-                    let timestamp = context.current().epoch_millis();
-                    let handshake_bytes = create_handshake(
-                        &mut sender,
-                        TEST_NAMESPACE,
-                        timestamp,
-                        recipient.public_key(),
-                        ephemeral_public_key,
-                    )
-                    .unwrap();
-                    send_frame(&mut stream_sender, &handshake_bytes, ONE_MEGABYTE)
-                        .await
-                        .unwrap();
-                }
-            });
-
-            // Call the verify function
-            let result = IncomingHandshake::verify(
-                &context,
-                &recipient,
-                TEST_NAMESPACE,
-                ONE_MEGABYTE,
-                Duration::from_secs(1),
-                Duration::from_secs(1),
-                Duration::from_secs(1),
-                sink,
-                stream,
-            )
-            .await;
+            let (_, stream) = mocks::Channel::init();
+
+            // Call the verify function for one peer, but never send the handshake from the other
+            let config = Config {
+                crypto: Ed25519::from_seed(1),
+                namespace: TEST_NAMESPACE.to_vec(),
+                max_message_size: ONE_MEGABYTE,
+                synchrony_bound: Duration::from_secs(1),
+                max_handshake_age: Duration::from_secs(1),
+                handshake_timeout: Duration::from_secs(1),
+            };
+            let result = IncomingConnection::verify(&context, config, sink, stream).await;
 
             // Assert that the result is an Err of type Error::HandshakeTimeout
             assert!(matches!(result, Err(Error::HandshakeTimeout)));
         });
     }
 
-    #[test]
-    fn test_handshake_verify_invalid_public_key() {
-        let (executor, context, _) = Executor::default();
-        executor.start(async move {
-            let mut crypto = Ed25519::from_seed(0);
-            let recipient_public_key = crypto.public_key();
-            let ephemeral_public_key = x25519_dalek::PublicKey::from([0u8; 32]);
-
-            let handshake = create_handshake(
-                &mut crypto,
-                TEST_NAMESPACE,
-                0, // timestamp
-                recipient_public_key,
-                ephemeral_public_key,
-            )
-            .unwrap();
-
-            // Tamper with the handshake to make the signature invalid
-            let mut handshake =
-                wire::Handshake::decode(handshake).expect("failed to decode handshake");
-            let (public_key, signature) = match handshake.signature {
-                Some(wire::Signature {
-                    public_key,
-                    signature,
-                }) => (public_key, signature),
-                _ => panic!("signature missing"),
-            };
-            let mut public_key = public_key.to_vec();
-            public_key.truncate(28);
-            handshake.signature = Some(wire::Signature {
-                public_key,
-                signature,
-            });
-
-            // Verify the handshake
-            let result = Handshake::verify(
-                &context,
-                &crypto,
-                TEST_NAMESPACE,
-                Duration::from_secs(5),
-                Duration::from_secs(5),
-                handshake.encode_to_vec().into(),
-            );
-            assert!(matches!(result, Err(Error::InvalidPeerPublicKey)));
-        });
-    }
-
-    #[test]
-    fn test_handshake_verify_invalid_ephemeral_public_key() {
-        let (executor, context, _) = Executor::default();
-        executor.start(async move {
-            let mut crypto = Ed25519::from_seed(0);
-            let recipient_public_key = crypto.public_key();
-            let ephemeral_public_key = x25519_dalek::PublicKey::from([0u8; 32]);
-
-            let handshake = create_handshake(
-                &mut crypto,
-                TEST_NAMESPACE,
-                0, // timestamp
-                recipient_public_key,
-                ephemeral_public_key,
-            )
-            .unwrap();
-
-            // Tamper with the handshake to make the signature invalid
-            let mut handshake =
-                wire::Handshake::decode(handshake).expect("failed to decode handshake");
-            let mut ephemeral_public_key = handshake.ephemeral_public_key.to_vec();
-            ephemeral_public_key.truncate(28);
-            handshake.ephemeral_public_key = ephemeral_public_key;
-
-            // Verify the handshake
-            let result = Handshake::verify(
-                &context,
-                &crypto,
-                TEST_NAMESPACE,
-                Duration::from_secs(5),
-                Duration::from_secs(5),
-                handshake.encode_to_vec().into(),
-            );
-            assert!(matches!(result, Err(Error::InvalidEphemeralPublicKey)));
-        });
-    }
-
     #[test]
     fn test_handshake_verify_invalid_signature() {
         let (executor, context, _) = Executor::default();
         executor.start(async move {
             let mut crypto = Ed25519::from_seed(0);
-            let recipient_public_key = crypto.public_key();
-            let ephemeral_public_key = x25519_dalek::PublicKey::from([0u8; 32]);
+            let recipient = crypto.public_key();
+            let ephemeral_public_key = x25519::PublicKey::from_bytes([0u8; 32]);
 
-            let handshake = create_handshake(
+            let handshake = Signed::sign(
                 &mut crypto,
                 TEST_NAMESPACE,
-                0, // timestamp
-                recipient_public_key,
-                ephemeral_public_key,
-            )
-            .unwrap();
+                Info {
+                    timestamp: 0,
+                    recipient,
+                    ephemeral_public_key,
+                },
+            );
 
             // Tamper with the handshake to make the signature invalid
             let mut handshake =
-                wire::Handshake::decode(handshake).expect("failed to decode handshake");
-            let (public_key, signature) = match handshake.signature {
-                Some(wire::Signature {
-                    public_key,
-                    signature,
-                }) => (public_key, signature),
-                _ => panic!("signature missing"),
-            };
-            let mut signature = signature.to_vec();
-            signature[0] ^= 0xFF;
-            handshake.signature = Some(wire::Signature {
-                public_key,
-                signature,
-            });
+                Signed::decode(handshake.encode()).expect("failed to decode handshake");
+            handshake.info.timestamp += 1;
 
             // Verify the handshake
-            let result = Handshake::verify(
+            let result = handshake.verify(
                 &context,
                 &crypto,
                 TEST_NAMESPACE,
                 Duration::from_secs(5),
                 Duration::from_secs(5),
-                handshake.encode_to_vec().into(),
             );
             assert!(matches!(result, Err(Error::InvalidSignature)));
         });
@@ -575,48 +437,48 @@ mod tests {
         let (executor, context, _) = Executor::default();
         executor.start(async move {
             let mut crypto = Ed25519::from_seed(0);
-            let recipient_public_key = crypto.public_key();
-            let ephemeral_public_key = x25519_dalek::PublicKey::from([0u8; 32]);
+            let recipient = crypto.public_key();
+            let ephemeral_public_key = x25519::PublicKey::from_bytes([0u8; 32]);
 
             let timeout_duration = Duration::from_secs(5);
             let synchrony_bound = Duration::from_secs(0);
 
             // Create a handshake, setting the timestamp to 0.
-            let handshake = create_handshake(
+            let handshake = Signed::sign(
                 &mut crypto,
                 TEST_NAMESPACE,
-                0, // timestamp
-                recipient_public_key,
-                ephemeral_public_key,
-            )
-            .unwrap();
+                Info {
+                    timestamp: 0,
+                    recipient,
+                    ephemeral_public_key,
+                },
+            );
 
             // Time starts at 0 in deterministic executor.
             // Sleep for the exact timeout duration.
             context.sleep(timeout_duration).await;
 
             // Verify the handshake, it should be fine still.
-            Handshake::verify(
-                &context,
-                &crypto,
-                TEST_NAMESPACE,
-                synchrony_bound,
-                timeout_duration,
-                handshake.clone(),
-            )
-            .unwrap(); // no error
+            handshake
+                .verify(
+                    &context,
+                    &crypto,
+                    TEST_NAMESPACE,
+                    synchrony_bound,
+                    timeout_duration,
+                )
+                .unwrap();
 
             // Timeout by waiting 1 more millisecond.
             context.sleep(Duration::from_millis(1)).await;
 
             // Verify that a timeout error is returned.
-            let result = Handshake::verify(
+            let result = handshake.verify(
                 &context,
                 &crypto,
                 TEST_NAMESPACE,
                 synchrony_bound,
                 timeout_duration,
-                handshake,
             );
             assert!(matches!(result, Err(Error::InvalidTimestampOld(t)) if t == 0));
         });
@@ -627,49 +489,51 @@ mod tests {
         let (executor, context, _) = Executor::default();
         executor.start(async move {
             let mut crypto = Ed25519::from_seed(0);
-            let recipient_public_key = crypto.public_key();
-            let ephemeral_public_key = x25519_dalek::PublicKey::from([0u8; 32]);
+            let recipient = crypto.public_key();
+            let ephemeral_public_key = x25519::PublicKey::from_bytes([0u8; 32]);
 
             let timeout_duration = Duration::from_secs(0);
             const SYNCHRONY_BOUND_MILLIS: u64 = 5_000;
             let synchrony_bound = Duration::from_millis(SYNCHRONY_BOUND_MILLIS);
 
             // Create a handshake at the synchrony bound.
-            let handshake_ok = create_handshake(
+            let handshake_ok = Signed::sign(
                 &mut crypto,
                 TEST_NAMESPACE,
-                SYNCHRONY_BOUND_MILLIS,
-                recipient_public_key.clone(),
-                ephemeral_public_key,
-            ).unwrap();
+                Info{
+                    timestamp: SYNCHRONY_BOUND_MILLIS,
+                    recipient: recipient.clone(),
+                    ephemeral_public_key,
+                },
+            );
 
             // Create a handshake 1ms too far into the future.
-            let handshake_late = create_handshake(
+            let handshake_late = Signed::sign(
                 &mut crypto,
                 TEST_NAMESPACE,
-                SYNCHRONY_BOUND_MILLIS + 1,
-                recipient_public_key,
-                ephemeral_public_key,
-            ).unwrap();
+                Info{
+                    timestamp:SYNCHRONY_BOUND_MILLIS + 1,
+                    recipient,
+                    ephemeral_public_key,
+                },
+            );
 
             // Verify the okay handshake.
-            Handshake::verify(
+            handshake_ok.verify(
                 &context,
                 &crypto,
                 TEST_NAMESPACE,
                 synchrony_bound,
                 timeout_duration,
-                handshake_ok,
             ).unwrap(); // no error
 
             // Handshake too far into the future fails.
-            let result = Handshake::verify(
+            let result = handshake_late.verify(
                 &context,
                 &crypto,
                 TEST_NAMESPACE,
                 synchrony_bound,
                 timeout_duration,
-                handshake_late,
             );
             assert!(matches!(result, Err(Error::InvalidTimestampFuture(t)) if t == SYNCHRONY_BOUND_MILLIS + 1));
         });
diff --git a/stream/src/public_key/mod.rs b/stream/src/public_key/mod.rs
index 54d7ba0c3a..b124531db1 100644
--- a/stream/src/public_key/mod.rs
+++ b/stream/src/public_key/mod.rs
@@ -21,24 +21,20 @@
 //! it can be simulated deterministically.
 //!
 //! In any handshake, the dialer is the party that attempts to connect to some known address/identity (public key)
-//! and the recipient of this connection is the dialee. Upon forming a TCP connection, the dialer sends a signed
-//! handshake message to the dialee.
+//! and the recipient of this connection is the listener. Upon forming a TCP connection, the dialer sends a signed
+//! handshake message to the listener. Besides the signature and the public key of the dialer, the handshake message
+//! contains:
 //!
-//! ```protobuf
-//! message Handshake {
-//!     bytes recipient_public_key = 1;
-//!     bytes ephemeral_public_key = 2;
-//!     uint64 timestamp = 3;
-//!     Signature signature = 4;
-//! }
-//! ```
+//! - The receiver's public key.
+//! - The sender's ephemeral public key (used to establish a shared secret).
+//! - The current timestamp (used to prevent replay attacks).
 //!
-//! The dialee verifies the public keys are well-formatted, the timestamp is valid (not too old/not too far in the future),
-//! and that the signature is valid. If all these checks pass, the dialee checks to see if it is already connected or dialing
+//! The listener verifies the public keys are well-formatted, the timestamp is valid (not too old/not too far in the future),
+//! and that the signature is valid. If all these checks pass, the listener checks to see if it is already connected or dialing
 //! this peer. If it is, it drops the connection. If it isn't, it sends back its own signed handshake message (same as above)
 //! and considers the connection established.
 //!
-//! Upon receiving the dialee's handshake message, the dialer verifies the same data as the dialee and additionally verifies
+//! Upon receiving the listener's handshake message, the dialer verifies the same data as the listener and additionally verifies
 //! that the public key returned matches what they expected at the address. If all these checks pass, the dialer considers the
 //! connection established. If not, the dialer drops the connection.
 //!
@@ -54,7 +50,7 @@
 //!
 //! Each peer maintains a pair of ChaCha20-Poly1305 nonces (12 bytes), one for itself and one for
 //! the other. Each nonce is constructed using a counter that starts at either 1 for the dialer, or
-//! 0 for the dialee. For each message sent, the relevant counter is incremented by 2, ensuring that
+//! 0 for the listener. For each message sent, the relevant counter is incremented by 2, ensuring that
 //! the two counters have disjoint nonce spaces.
 //!
 //! The nonce is the least-significant 12 bytes of the counter, encoded big-endian. This provides 2^95 unique nonces
@@ -73,9 +69,6 @@ mod connection;
 pub use connection::{Connection, IncomingConnection, Receiver, Sender};
 mod handshake;
 mod nonce;
-mod wire {
-    include!(concat!(env!("OUT_DIR"), "/wire.rs"));
-}
 mod x25519;
 
 /// Configuration for a connection.
@@ -90,18 +83,21 @@ pub struct Config<C: Scheme> {
     /// Cryptographic primitives.
     pub crypto: C,
 
-    /// Prefix for all signed messages to avoid replay attacks.
+    /// Prefix for all signed messages. Should be unique to the application.
+    /// Used to avoid replay attacks across different applications
     pub namespace: Vec<u8>,
 
-    /// Maximum size allowed for messages over any connection.
+    /// Maximum size allowed for messages (in bytes).
+    /// Used to prevent DoS attacks.
     pub max_message_size: usize,
 
     /// Time into the future that a timestamp can be and still be considered valid.
+    /// Used to handle clock skew between peers.
     pub synchrony_bound: Duration,
 
-    /// Duration after which a handshake message is considered stale.
+    /// Maximum age of a handshake message before it is considered stale.
     pub max_handshake_age: Duration,
 
-    /// Timeout for the handshake process.
+    /// Timeout for completing the handshake process.
     pub handshake_timeout: Duration,
 }
diff --git a/stream/src/public_key/wire.proto b/stream/src/public_key/wire.proto
deleted file mode 100644
index b4b335b11b..0000000000
--- a/stream/src/public_key/wire.proto
+++ /dev/null
@@ -1,26 +0,0 @@
-syntax = "proto3";
-package wire;
-
-// Allows recipient to verify that the sender has the private key
-// of public key before sending any data.
-//
-// By requiring the server to have their public key signed, they prevent
-// a malicious peer from forwarding a handshake message from a previous
-// connection with public key (which could be used to convince the server
-// to start a useless handshake). Alternatively, we could require the
-// dialer to sign some random bytes provided by the server but this would
-// require the server to send a message to a peer before authorizing that
-// it should connect to them.
-message Handshake {
-    bytes recipient_public_key = 1;
-    bytes ephemeral_public_key = 2;
-    uint64 timestamp = 3;
-    Signature signature = 4;
-}
-
-// Signature is an arbitrary public key and signature over
-// either a Handshake or Peer message.
-message Signature {
-    bytes public_key = 1;
-    bytes signature = 2;
-}
diff --git a/stream/src/public_key/x25519.rs b/stream/src/public_key/x25519.rs
index 25ee44c57d..0f8a4e43dc 100644
--- a/stream/src/public_key/x25519.rs
+++ b/stream/src/public_key/x25519.rs
@@ -1,55 +1,88 @@
-use crate::Error;
-use bytes::Bytes;
+use commonware_codec::{Codec, Error as CodecError, Reader, SizedCodec, Writer};
 use rand::{CryptoRng, Rng};
-use x25519_dalek::{EphemeralSecret, PublicKey};
+use x25519_dalek::{EphemeralSecret, PublicKey as X25519PublicKey};
 
-const PUBLIC_KEY_LENGTH: usize = 32;
+#[derive(PartialEq, Eq, Hash, Copy, Clone, Debug)]
+pub struct PublicKey {
+    inner: X25519PublicKey,
+}
 
-pub fn new<R: Rng + CryptoRng>(rng: &mut R) -> EphemeralSecret {
-    EphemeralSecret::random_from_rng(rng)
+impl PublicKey {
+    pub fn from_secret(secret: &EphemeralSecret) -> Self {
+        PublicKey {
+            inner: X25519PublicKey::from(secret),
+        }
+    }
+
+    #[cfg(test)]
+    pub fn from_bytes(array: [u8; 32]) -> Self {
+        PublicKey {
+            inner: X25519PublicKey::from(array),
+        }
+    }
+}
+
+impl AsRef<x25519_dalek::PublicKey> for PublicKey {
+    fn as_ref(&self) -> &x25519_dalek::PublicKey {
+        &self.inner
+    }
 }
 
-pub fn decode_public_key(public_key: &[u8]) -> Result<PublicKey, Error> {
-    // Construct a public key array from the data
-    let public_key: [u8; PUBLIC_KEY_LENGTH] = match public_key.as_ref().try_into() {
-        Ok(key) => key,
-        Err(_) => return Err(Error::InvalidEphemeralPublicKey),
-    };
+impl Codec for PublicKey {
+    fn write(&self, writer: &mut impl Writer) {
+        self.inner.as_bytes().write(writer);
+    }
+    fn read(reader: &mut impl Reader) -> Result<Self, CodecError> {
+        let public_key: [u8; Self::LEN_ENCODED] = reader.read_fixed()?;
+        Ok(PublicKey {
+            inner: X25519PublicKey::from(public_key),
+        })
+    }
+    fn len_encoded(&self) -> usize {
+        Self::LEN_ENCODED
+    }
+}
 
-    // Create the public key from the array
-    Ok(PublicKey::from(public_key))
+impl SizedCodec for PublicKey {
+    const LEN_ENCODED: usize = 32;
 }
 
-pub fn encode_public_key(public_key: PublicKey) -> Bytes {
-    public_key.as_bytes().to_vec().into()
+pub fn new<R: Rng + CryptoRng>(rng: &mut R) -> EphemeralSecret {
+    EphemeralSecret::random_from_rng(rng)
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
+    use bytes::Bytes;
+    use commonware_runtime::deterministic::Executor;
 
     #[test]
-    fn test_encode_decode_public_key() {
-        let mut rng = rand::thread_rng();
-        let secret = new(&mut rng);
-        let public_key = PublicKey::from(&secret);
-
-        let encoded = encode_public_key(public_key);
-        let decoded = decode_public_key(&encoded).unwrap();
+    fn test_codec() {
+        // Create a random public key
+        let (_, mut context, _) = Executor::default();
+        let mut buf = [0u8; PublicKey::LEN_ENCODED];
+        context.fill(&mut buf);
 
-        assert_eq!(public_key, decoded);
+        let original = PublicKey {
+            inner: X25519PublicKey::from(buf),
+        };
+        let encoded = original.encode();
+        assert_eq!(encoded.len(), PublicKey::LEN_ENCODED);
+        let decoded = PublicKey::decode(encoded).unwrap();
+        assert_eq!(original, decoded);
     }
 
     #[test]
-    fn invalid_public_key() {
+    fn test_decode_invalid() {
         // Create a Bytes object that is too short
         let invalid_bytes = Bytes::from(vec![1, 2, 3]); // Length 3 instead of 32
-        let result = decode_public_key(&invalid_bytes);
-        assert!(matches!(result, Err(Error::InvalidEphemeralPublicKey)));
+        let result = PublicKey::decode(invalid_bytes);
+        assert!(result.is_err());
 
         // Create Bytes object that's too long
         let too_long_bytes = Bytes::from(vec![0u8; 33]); // Length 33
-        let result = decode_public_key(&too_long_bytes);
-        assert!(matches!(result, Err(Error::InvalidEphemeralPublicKey)));
+        let result = PublicKey::decode(too_long_bytes);
+        assert!(result.is_err());
     }
 }
diff --git a/stream/src/utils/codec.rs b/stream/src/utils/codec.rs
index 3bdcf34e16..387dcd4138 100644
--- a/stream/src/utils/codec.rs
+++ b/stream/src/utils/codec.rs
@@ -21,10 +21,10 @@ pub async fn send_frame<S: Sink>(
 
     // Send the length of the message
     let f: [u8; 4] = len.to_be_bytes();
-    sink.send(&f).await.map_err(|_| Error::SendFailed)?;
+    sink.send(&f).await.map_err(Error::SendFailed)?;
 
     // Send the rest of the message
-    sink.send(buf).await.map_err(|_| Error::SendFailed)?;
+    sink.send(buf).await.map_err(Error::SendFailed)?;
 
     Ok(())
 }
@@ -37,7 +37,7 @@ pub async fn recv_frame<T: Stream>(
 ) -> Result<Bytes, Error> {
     // Read the first 4 bytes to get the length of the message
     let mut buf = [0u8; 4];
-    stream.recv(&mut buf).await.map_err(|_| Error::RecvFailed)?;
+    stream.recv(&mut buf).await.map_err(Error::RecvFailed)?;
 
     // Validate frame size
     let len = u32::from_be_bytes(buf) as usize;
@@ -50,7 +50,7 @@ pub async fn recv_frame<T: Stream>(
 
     // Read the rest of the message
     let mut buf = vec![0u8; len];
-    stream.recv(&mut buf).await.map_err(|_| Error::RecvFailed)?;
+    stream.recv(&mut buf).await.map_err(Error::RecvFailed)?;
 
     Ok(Bytes::from(buf))
 }


Reviews:
user: login: patrick-ogrady
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-17T23:41:14Z
user: login: BrendanChou
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T02:05:02Z
user: login: patrick-ogrady
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T18:07:51Z
user: login: BrendanChou
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T18:25:22Z
user: login: patrick-ogrady
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T18:46:06Z
user: login: patrick-ogrady
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T18:47:29Z
user: login: BrendanChou
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T18:51:15Z
user: login: patrick-ogrady
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T18:52:17Z
user: login: patrick-ogrady
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T18:55:31Z
user: login: patrick-ogrady
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T19:01:47Z

Review comments:
diff_hunk: @@ -168,47 +187,54 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
     /// Because we already verified the peer's handshake, this function
     /// only needs to send our handshake message for the connection to be fully
     /// initialized.
-    pub async fn upgrade_listener<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
+    pub async fn upgrade_as_dialee<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
path: stream/src/public_key/connection.rs
user: login: patrick-ogrady
body: nit: I'm definitely biased, but I strongly prefer `upgrade_listener` and `upgrade_dialer` to the new names
created_at: 2025-03-17T23:41:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -168,47 +187,54 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
     /// Because we already verified the peer's handshake, this function
     /// only needs to send our handshake message for the connection to be fully
     /// initialized.
-    pub async fn upgrade_listener<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
+    pub async fn upgrade_as_dialee<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
path: stream/src/public_key/connection.rs
user: login: BrendanChou
body: I can remove the `as` part.

However, everywhere else we refer to the peer being dialed as the "dialee" and nowhere else do we call it the "listener". We should either call it `upgrade_dialee` or change all `dialee` references to `listener` IMO. Given that, it seems using `dialee` is better
created_at: 2025-03-18T02:05:02Z
author_association: COLLABORATOR
diff_hunk: @@ -168,47 +187,54 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
     /// Because we already verified the peer's handshake, this function
     /// only needs to send our handshake message for the connection to be fully
     /// initialized.
-    pub async fn upgrade_listener<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
+    pub async fn upgrade_as_dialee<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
path: stream/src/public_key/connection.rs
user: login: patrick-ogrady
body: I'd strongly prefer `dialer`/`listener` everywhere. I think that makes it much clearer what is going on.
created_at: 2025-03-18T18:07:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -168,47 +187,54 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
     /// Because we already verified the peer's handshake, this function
     /// only needs to send our handshake message for the connection to be fully
     /// initialized.
-    pub async fn upgrade_listener<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
+    pub async fn upgrade_as_dialee<R: Rng + CryptoRng + Spawner + Clock, C: Scheme>(
path: stream/src/public_key/connection.rs
user: login: BrendanChou
body: Changed
created_at: 2025-03-18T18:25:22Z
author_association: COLLABORATOR
diff_hunk: @@ -24,35 +23,57 @@ const ENCRYPTION_TAG_LENGTH: usize = 16;
 /// An incoming connection with a verified peer handshake.
 pub struct IncomingConnection<C: Scheme, Si: Sink, St: Stream> {
     config: Config<C>,
-    handshake: IncomingHandshake<Si, St, C>,
+    sink: Si,
+    stream: St,
+    deadline: SystemTime,
+    ephemeral_public_key: x25519::PublicKey,
+    peer_public_key: C::PublicKey,
 }
 
 impl<C: Scheme, Si: Sink, St: Stream> IncomingConnection<C, Si, St> {
-    /// Verify the handshake of an incoming connection.
-    pub async fn verify<R: Rng + CryptoRng + Spawner + Clock>(
-        context: &R,
+    pub async fn verify<E: Clock + Spawner>(
+        context: &E,
         config: Config<C>,
         sink: Si,
-        stream: St,
+        mut stream: St,
     ) -> Result<Self, Error> {
-        let handshake = IncomingHandshake::verify(
+        // Set handshake deadline
+        let deadline = context.current() + config.handshake_timeout;
+
+        // Wait for up to handshake timeout for response
+        let msg = select! {
+            _ = context.sleep_until(deadline) => { return Err(Error::HandshakeTimeout) },
+            result = recv_frame(&mut stream, config.max_message_size) => { result? },
+        };
+
+        // Verify handshake message from peer
+        let signed_handshake =
path: stream/src/public_key/connection.rs
user: login: patrick-ogrady
body: Does this assert that the handshake is exactly the size we expect (no trailing bytes)? Or is the assumption that we just rely on the network layer to ensure we don't get too much junk with any particular message?

(This was the rationale for adding a type to the codec everywhere so we could both parse different objects explicitly + assert those different objects have no trailing bytes).
created_at: 2025-03-18T18:44:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -21,24 +21,20 @@
 //! it can be simulated deterministically.
 //!
 //! In any handshake, the dialer is the party that attempts to connect to some known address/identity (public key)
-//! and the recipient of this connection is the dialee. Upon forming a TCP connection, the dialer sends a signed
-//! handshake message to the dialee.
+//! and the recipient of this connection is the listener. Upon forming a TCP connection, the dialer sends a signed
+//! handshake message to the listener. Besides the signature and the public key of the dialer, the handshake message
+//! contains:
path: stream/src/public_key/mod.rs
user: login: patrick-ogrady
body: It also includes a signature?
created_at: 2025-03-18T18:45:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,35 +23,57 @@ const ENCRYPTION_TAG_LENGTH: usize = 16;
 /// An incoming connection with a verified peer handshake.
 pub struct IncomingConnection<C: Scheme, Si: Sink, St: Stream> {
     config: Config<C>,
-    handshake: IncomingHandshake<Si, St, C>,
+    sink: Si,
+    stream: St,
+    deadline: SystemTime,
+    ephemeral_public_key: x25519::PublicKey,
+    peer_public_key: C::PublicKey,
 }
 
 impl<C: Scheme, Si: Sink, St: Stream> IncomingConnection<C, Si, St> {
-    /// Verify the handshake of an incoming connection.
-    pub async fn verify<R: Rng + CryptoRng + Spawner + Clock>(
-        context: &R,
+    pub async fn verify<E: Clock + Spawner>(
+        context: &E,
         config: Config<C>,
         sink: Si,
-        stream: St,
+        mut stream: St,
     ) -> Result<Self, Error> {
-        let handshake = IncomingHandshake::verify(
+        // Set handshake deadline
+        let deadline = context.current() + config.handshake_timeout;
+
+        // Wait for up to handshake timeout for response
+        let msg = select! {
+            _ = context.sleep_until(deadline) => { return Err(Error::HandshakeTimeout) },
+            result = recv_frame(&mut stream, config.max_message_size) => { result? },
+        };
+
+        // Verify handshake message from peer
+        let signed_handshake =
path: stream/src/public_key/connection.rs
user: login: patrick-ogrady
body: I think this could be accomplished by adding a `.done() -> Result<(), CodecError>` that checks remaining is zero and if not returns an error with bytes left?
created_at: 2025-03-18T18:47:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -24,35 +23,57 @@ const ENCRYPTION_TAG_LENGTH: usize = 16;
 /// An incoming connection with a verified peer handshake.
 pub struct IncomingConnection<C: Scheme, Si: Sink, St: Stream> {
     config: Config<C>,
-    handshake: IncomingHandshake<Si, St, C>,
+    sink: Si,
+    stream: St,
+    deadline: SystemTime,
+    ephemeral_public_key: x25519::PublicKey,
+    peer_public_key: C::PublicKey,
 }
 
 impl<C: Scheme, Si: Sink, St: Stream> IncomingConnection<C, Si, St> {
-    /// Verify the handshake of an incoming connection.
-    pub async fn verify<R: Rng + CryptoRng + Spawner + Clock>(
-        context: &R,
+    pub async fn verify<E: Clock + Spawner>(
+        context: &E,
         config: Config<C>,
         sink: Si,
-        stream: St,
+        mut stream: St,
     ) -> Result<Self, Error> {
-        let handshake = IncomingHandshake::verify(
+        // Set handshake deadline
+        let deadline = context.current() + config.handshake_timeout;
+
+        // Wait for up to handshake timeout for response
+        let msg = select! {
+            _ = context.sleep_until(deadline) => { return Err(Error::HandshakeTimeout) },
+            result = recv_frame(&mut stream, config.max_message_size) => { result? },
+        };
+
+        // Verify handshake message from peer
+        let signed_handshake =
path: stream/src/public_key/connection.rs
user: login: BrendanChou
body: Yes! `decode` [takes care of the trailing bytes](https://github.com/commonwarexyz/monorepo/blob/5b09315df818cd0c084cc2b8086b708940631eb1/codec/src/codec.rs#L35-L37)!
created_at: 2025-03-18T18:51:15Z
author_association: COLLABORATOR
diff_hunk: @@ -24,35 +23,57 @@ const ENCRYPTION_TAG_LENGTH: usize = 16;
 /// An incoming connection with a verified peer handshake.
 pub struct IncomingConnection<C: Scheme, Si: Sink, St: Stream> {
     config: Config<C>,
-    handshake: IncomingHandshake<Si, St, C>,
+    sink: Si,
+    stream: St,
+    deadline: SystemTime,
+    ephemeral_public_key: x25519::PublicKey,
+    peer_public_key: C::PublicKey,
 }
 
 impl<C: Scheme, Si: Sink, St: Stream> IncomingConnection<C, Si, St> {
-    /// Verify the handshake of an incoming connection.
-    pub async fn verify<R: Rng + CryptoRng + Spawner + Clock>(
-        context: &R,
+    pub async fn verify<E: Clock + Spawner>(
+        context: &E,
         config: Config<C>,
         sink: Si,
-        stream: St,
+        mut stream: St,
     ) -> Result<Self, Error> {
-        let handshake = IncomingHandshake::verify(
+        // Set handshake deadline
+        let deadline = context.current() + config.handshake_timeout;
+
+        // Wait for up to handshake timeout for response
+        let msg = select! {
+            _ = context.sleep_until(deadline) => { return Err(Error::HandshakeTimeout) },
+            result = recv_frame(&mut stream, config.max_message_size) => { result? },
+        };
+
+        // Verify handshake message from peer
+        let signed_handshake =
path: stream/src/public_key/connection.rs
user: login: patrick-ogrady
body: I think it might be worth commenting this under the // Verify... because auditor's might not appreciate/realize this.

SO nice
created_at: 2025-03-18T18:52:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -101,25 +122,22 @@ impl<Si: Sink, St: Stream> Connection<Si, St> {
 
         // Generate shared secret
         let secret = x25519::new(&mut context);
-        let ephemeral = x25519_dalek::PublicKey::from(&secret);
 
         // Send handshake
-        let timestamp = context.current().epoch_millis();
-        let msg = create_handshake(
+        let msg = handshake::Signed::sign(
             &mut config.crypto,
             &config.namespace,
-            timestamp,
-            peer.clone(),
-            ephemeral,
-        )?;
+            handshake::Info::<C>::new(peer.clone(), &secret, context.current().epoch_millis()),
path: stream/src/public_key/connection.rs
user: login: patrick-ogrady
body: imo it is nice to give timestamp its own variable (given this function signature is already complex).
created_at: 2025-03-18T18:54:35Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-03-18T19:13:36Z
body: ## [Codecov](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) Report
Attention: Patch coverage is `99.79633%` with `1 line` in your changes missing coverage. Please review.
> Project coverage is 89.34%. Comparing base [(`5b09315`)](https://app.codecov.io/gh/commonwarexyz/monorepo/commit/5b09315df818cd0c084cc2b8086b708940631eb1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) to head [(`4520676`)](https://app.codecov.io/gh/commonwarexyz/monorepo/commit/4520676ba7eec2c4a69ea2ae9bb36ce647ea0420?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz).
> Report is 1 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) | Patch % | Lines |
|---|---|---|
| [p2p/src/simulated/network.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=p2p%2Fsrc%2Fsimulated%2Fnetwork.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cDJwL3NyYy9zaW11bGF0ZWQvbmV0d29yay5ycw==) | 0.00% | [1 Missing :warning: ](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) |

```diff
@@            Coverage Diff             @@
##             main     #624      +/-   ##
==========================================
+ Coverage   89.25%   89.34%   +0.09%     
==========================================
  Files         138      138              
  Lines       36210    36266      +56     
==========================================
+ Hits        32318    32403      +85     
+ Misses       3892     3863      -29     
```

| [Files with missing lines](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) | Coverage Δ | |
|---|---|---|
| [runtime/src/deterministic.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=runtime%2Fsrc%2Fdeterministic.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvZGV0ZXJtaW5pc3RpYy5ycw==) | `99.40% <100.00%> (ø)` | |
| [stream/src/public\_key/connection.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=stream%2Fsrc%2Fpublic_key%2Fconnection.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-c3RyZWFtL3NyYy9wdWJsaWNfa2V5L2Nvbm5lY3Rpb24ucnM=) | `99.54% <100.00%> (+6.38%)` | :arrow_up: |
| [stream/src/public\_key/handshake.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=stream%2Fsrc%2Fpublic_key%2Fhandshake.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-c3RyZWFtL3NyYy9wdWJsaWNfa2V5L2hhbmRzaGFrZS5ycw==) | `100.00% <100.00%> (+2.46%)` | :arrow_up: |
| [stream/src/public\_key/x25519.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=stream%2Fsrc%2Fpublic_key%2Fx25519.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-c3RyZWFtL3NyYy9wdWJsaWNfa2V5L3gyNTUxOS5ycw==) | `100.00% <100.00%> (ø)` | |
| [stream/src/utils/codec.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=stream%2Fsrc%2Futils%2Fcodec.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-c3RyZWFtL3NyYy91dGlscy9jb2RlYy5ycw==) | `100.00% <100.00%> (ø)` | |
| [p2p/src/simulated/network.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?src=pr&el=tree&filepath=p2p%2Fsrc%2Fsimulated%2Fnetwork.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cDJwL3NyYy9zaW11bGF0ZWQvbmV0d29yay5ycw==) | `94.87% <0.00%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz)

------

[Continue to review full report in Codecov by Sentry](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?dropdown=coverage&src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz). Last update [5b09315...4520676](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/624?dropdown=coverage&src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz).

<details><summary>🚀 New features to boost your workflow: </summary>

- ❄ [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE