add login message 2 sec timeout

lib/transport/src/auth.rs

@@ -6,6 +6,8 @@
 //! This is trivial for Periphery -> Core connection, but presents a challenge
 //! for Core -> Periphery, where untrusted TLS certs are being used.
 
+use std::time::Duration;
+
 use anyhow::Context;
 use axum::http::{HeaderMap, HeaderValue};
 use base64::{Engine, prelude::BASE64_STANDARD};
@@ -31,6 +33,8 @@ pub trait LoginFlow {
   ) -> impl Future<Output = anyhow::Result<()>>;
 }
 
+const AUTH_TIMEOUT: Duration = Duration::from_secs(2);
+
 pub struct ServerLoginFlow;
 
 impl LoginFlow for ServerLoginFlow {
@@ -58,7 +62,7 @@ impl LoginFlow for ServerLoginFlow {
 
       // Receive and read handshake_m1
       let handshake_m1 = socket
-        .recv_bytes()
+        .recv_bytes_with_timeout(AUTH_TIMEOUT)
         .await
         .context("Failed to get handshake_m1")?;
       match MessageState::from_byte(
@@ -86,7 +90,7 @@ impl LoginFlow for ServerLoginFlow {
 
       // Receive and read handshake_m3
       let handshake_m3 = socket
-        .recv_bytes()
+        .recv_bytes_with_timeout(AUTH_TIMEOUT)
         .await
         .context("Failed to get handshake_m3")?;
       match MessageState::from_byte(
@@ -152,7 +156,7 @@ impl LoginFlow for ClientLoginFlow {
     let res = async {
       // Receive nonce from server
       let nonce = socket
-        .recv_bytes()
+        .recv_bytes_with_timeout(AUTH_TIMEOUT)
         .await
         .context("Failed to receive connection nonce")?;
 
@@ -176,7 +180,7 @@ impl LoginFlow for ClientLoginFlow {
 
       // Receive and read handshake_m2
       let handshake_m2 = socket
-        .recv_bytes()
+        .recv_bytes_with_timeout(AUTH_TIMEOUT)
         .await
         .context("Failed to get handshake_m2")?;
       match MessageState::from_byte(
@@ -212,7 +216,7 @@ impl LoginFlow for ClientLoginFlow {
 
       // Receive login state message and return based on value
       let state_msg = socket
-        .recv_bytes()
+        .recv_bytes_with_timeout(AUTH_TIMEOUT)
         .await
         .context("Failed to receive authentication state message")?;
       let state = state_msg.last().context(

lib/transport/src/websocket/mod.rs

@@ -1,8 +1,11 @@
 //! Wrappers to normalize behavior of websockets between Tungstenite and Axum,
 //! as well as streamline process of handling socket messages.
 
-use anyhow::anyhow;
+use std::time::Duration;
+
+use anyhow::{Context, anyhow};
 use bytes::Bytes;
+use futures_util::FutureExt;
 
 pub mod axum;
 pub mod tungstenite;
@@ -34,8 +37,7 @@ pub trait Websocket {
     Output = Result<WebsocketMessage<Self::CloseFrame>, Self::Error>,
   >;
 
-  /// Looping receiver for websocket messages which only returns
-  /// on significant messages.
+  /// Looping receiver for websocket messages which only returns on bytes.
   fn recv_bytes(
     &mut self,
   ) -> impl Future<Output = Result<Bytes, anyhow::Error>> {
@@ -52,6 +54,16 @@ pub trait Websocket {
     }
   }
 
+  /// Looping receiver for websocket messages which only returns on bytes.
+  /// Includes timeout.
+  fn recv_bytes_with_timeout(
+    &mut self,
+    timeout: Duration,
+  ) -> impl Future<Output = Result<Bytes, anyhow::Error>> {
+    tokio::time::timeout(timeout, self.recv_bytes())
+      .map(|res| res.context("Failed to receive bytes").flatten())
+  }
+
   /// Streamlined sending on bytes
   fn send(
     &mut self,