Move key exchange into connect

Author: bjorn3

src/lib.rs

@@ -23,49 +23,37 @@ pub struct Connection {
 }
 
 impl Connection {
-    /// Create a new [`Connection`]
-    pub fn new(
+    /// Create a new [`Connection`] and do the initial key exchange
+    pub async fn connect(
         stream: TcpStream,
         addr: SocketAddr,
         host_key: Arc<Ed25519KeyPair>,
-    ) -> anyhow::Result<Self> {
-        stream.set_nodelay(true)?;
+    ) -> Result<Connection, ()> {
+        if let Err(error) = stream.set_nodelay(true) {
+            warn!(addr = %addr, %error, "failed to set nodelay");
+            return Err(());
+        }
 
         let (stream_read, stream_write) = stream.into_split();
 
-        Ok(Self {
+        let mut connection = Self {
             stream_read: DecryptingReader::new(stream_read),
             stream_write: EncryptingWriter::new(stream_write),
             addr,
             host_key,
-        })
-    }
+        };
 
-    /// Drive the connection forward
-    pub async fn run(mut self) {
         let mut exchange = digest::Context::new(&digest::SHA256);
         let state = VersionExchange::default();
-        let Ok(state) = state.advance(&mut exchange, &mut self).await else {
-            return;
-        };
+        let state = state.advance(&mut exchange, &mut connection).await?;
+        let state = state.advance(&mut exchange, &mut connection).await?;
+        state.advance(exchange, &mut connection).await?;
 
-        let Ok(state) = state.advance(&mut exchange, &mut self).await else {
-            return;
-        };
-
-        let Ok(()) = state.advance(exchange, &mut self).await else {
-            return;
-        };
-
-        todo!();
+        Ok(connection)
     }
 
-    pub(crate) async fn connect(
-        stream: TcpStream,
-        addr: SocketAddr,
-        host_key: Arc<Ed25519KeyPair>,
-    ) -> anyhow::Result<Self> {
-        // complete connection till kex finished (incl sending the newkeys message)
+    /// Drive the connection forward
+    pub async fn run(mut self) {
         todo!()
     }
 

src/main.rs

@@ -58,8 +58,11 @@ async fn main() -> anyhow::Result<()> {
         match listener.accept().await {
             Ok((stream, addr)) => {
                 debug!(%addr, "accepted connection");
-                let conn = Connection::new(stream, addr, host_key.clone())?;
-                conn.run().await; // FIXME(aws/aws-lc-rs#975) use tokio::spawn() once StreamingDecryptingKey is Send
+                // FIXME(aws/aws-lc-rs#975) use tokio::spawn() once StreamingDecryptingKey is Send
+                let Ok(conn) = Connection::connect(stream, addr, host_key.clone()).await else {
+                    continue; // Some kind of error happened. Has been logged already.
+                };
+                conn.run().await;
             }
             Err(error) => {
                 warn!(%error, "failed to accept connection");