Improve reconnection / backoff logic

src/lib.rs

@@ -428,14 +428,15 @@ impl ProxyClient {
         )>(1024);
 
         // Connect to the proxy server and provide / verify attestation
-        let (mut sender, mut measurements, mut remote_attestation_type) = Self::setup_connection(
-            connector.clone(),
-            target.clone(),
-            cert_chain.clone(),
-            attestation_generator.clone(),
-            attestation_verifier.clone(),
-        )
-        .await?;
+        let (mut sender, mut measurements, mut remote_attestation_type) =
+            Self::setup_connection_with_backoff(
+                connector.clone(),
+                target.clone(),
+                cert_chain.clone(),
+                attestation_generator.clone(),
+                attestation_verifier.clone(),
+            )
+            .await?;
 
         tokio::spawn(async move {
             // Read an incoming request from the channel (from the source client)
@@ -483,15 +484,24 @@ impl ProxyClient {
                 // If the connection to the proxy server failed, reconnect
                 if should_reconnect {
                     // Reconnect to the server - retrying indefinately with a backoff
+
                     (sender, measurements, remote_attestation_type) =
-                        Self::setup_connection_with_backoff(
+                        match Self::setup_connection_with_backoff(
                             connector.clone(),
                             target.clone(),
                             cert_chain.clone(),
                             attestation_generator.clone(),
                             attestation_verifier.clone(),
                         )
-                        .await;
+                        .await
+                        {
+                            Ok((sender, measurements, remote_attestation_type)) => {
+                                (sender, measurements, remote_attestation_type)
+                            }
+                            Err(_) => {
+                                break;
+                            }
+                        }
                 }
             }
         });
@@ -562,7 +572,7 @@ impl ProxyClient {
         cert_chain: Option<Vec<CertificateDer<'static>>>,
         attestation_generator: AttestationGenerator,
         attestation_verifier: AttestationVerifier,
-    ) -> (Http2Sender, Option<MultiMeasurements>, AttestationType) {
+    ) -> Result<(Http2Sender, Option<MultiMeasurements>, AttestationType), ProxyError> {
         let mut delay = Duration::from_secs(1);
         let max_delay = Duration::from_secs(SERVER_RECONNECT_MAX_BACKOFF_SECS);
 
@@ -577,14 +587,19 @@ impl ProxyClient {
             .await
             {
                 Ok(output) => {
-                    return output;
+                    return Ok(output);
                 }
                 Err(e) => {
-                    warn!("Reconnect failed: {e}. Retrying in {:#?}...", delay);
-                    tokio::time::sleep(delay).await;
-
-                    // increase delay for next time (exponential), but clamp to max_delay
-                    delay = std::cmp::min(delay * 2, max_delay);
+                    if matches!(e, ProxyError::Io(_)) {
+                        warn!("Reconnect failed: {e}. Retrying in {:#?}...", delay);
+                        tokio::time::sleep(delay).await;
+
+                        // increase delay for next time (exponential), but clamp to max_delay
+                        delay = std::cmp::min(delay * 2, max_delay);
+                    } else {
+                        // If we get a non-IO error, bail
+                        return Err(e);
+                    }
                 }
             }
         }