Merge pull request #664 from DefGuard/dev

merge dev -> main
DefGuard · Jul 8, 2024 · add6d2e · add6d2e
2 parents 39015c7 + f28c602
commit add6d2e
Show file tree

Hide file tree

Showing 6 changed files with 88 additions and 78 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/db/models/user.rs b/src/db/models/user.rs
@@ -1,4 +1,4 @@
-use std::{string::ToString, time::SystemTime};
+use std::{fmt, time::SystemTime};
 
 use argon2::{
     password_hash::{
@@ -39,15 +39,19 @@ pub enum MFAMethod {
     Email,
 }
 
-impl ToString for MFAMethod {
-    fn to_string(&self) -> String {
-        match self {
-            MFAMethod::None => "None".into(),
-            MFAMethod::OneTimePassword => "TOTP".into(),
-            MFAMethod::Web3 => "Web3".into(),
-            MFAMethod::Webauthn => "WebAuthn".into(),
-            MFAMethod::Email => "Email".into(),
-        }
+impl fmt::Display for MFAMethod {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "{}",
+            match self {
+                MFAMethod::None => "None",
+                MFAMethod::OneTimePassword => "TOTP",
+                MFAMethod::Web3 => "Web3",
+                MFAMethod::Webauthn => "WebAuthn",
+                MFAMethod::Email => "Email",
+            }
+        )
     }
 }
 

diff --git a/src/grpc/mod.rs b/src/grpc/mod.rs
@@ -19,7 +19,7 @@ use tokio::{
     },
     time::sleep,
 };
-use tokio_stream::{wrappers::UnboundedReceiverStream, StreamExt};
+use tokio_stream::wrappers::UnboundedReceiverStream;
 use tonic::{
     transport::{Certificate, ClientTlsConfig, Endpoint, Identity, Server, ServerTlsConfig},
     Status,
@@ -381,10 +381,14 @@ pub async fn run_grpc_bidi_stream(
         };
         info!("Connected to proxy at {}", endpoint.uri());
         let mut resp_stream = response.into_inner();
-        while let Some(received) = resp_stream.next().await {
-            info!("Received message from proxy");
-            match received {
-                Ok(received) => {
+        'message: loop {
+            match resp_stream.message().await {
+                Ok(None) => {
+                    info!("stream was closed by the sender");
+                    break 'message;
+                }
+                Ok(Some(received)) => {
+                    info!("Received message from proxy");
                     let payload = match received.payload {
                         // rpc StartEnrollment (EnrollmentStartRequest) returns (EnrollmentStartResponse)
                         Some(core_request::Payload::EnrollmentStart(request)) => {
@@ -509,7 +513,12 @@ pub async fn run_grpc_bidi_stream(
                     };
                     tx.send(req).unwrap();
                 }
-                Err(err) => error!("stream error {err}"),
+                Err(err) => {
+                    error!("stream error: {err}");
+                    debug!("waiting 10s to re-establish the connection");
+                    sleep(TEN_SECS).await;
+                    break 'message;
+                }
             }
         }
     }