From e02e4ae224173fe63612e7542d656b0eadabc719 Mon Sep 17 00:00:00 2001
From: Nick Steel <nick@nsteel.co.uk>
Date: Mon, 23 Sep 2024 19:44:51 +0100
Subject: [PATCH] core: retry Session access-point connection (#1345)

Tries multiple access-points with a retry for each one, unless there was a login failure.
---
 CHANGELOG.md               |  3 ++-
 core/src/connection/mod.rs | 23 +++++++++++++++++++++++
 core/src/session.rs        | 29 +++++++++++++++++++++++------
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 681762d3b..7c70c72cb 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -56,7 +56,8 @@ https://github.com/librespot-org/librespot
 - [core] `FileId` is moved out of `SpotifyId`. For now it will be re-exported.
 - [core] Report actual platform data on login
 - [core] Support `Session` authentication with a Spotify access token
-- [core] `Credentials.username` is now an `Option` (breaking) 
+- [core] `Credentials.username` is now an `Option` (breaking)
+- [core] `Session::connect` tries multiple access points, retrying each one.
 - [main] `autoplay {on|off}` now acts as an override. If unspecified, `librespot`
   now follows the setting in the Connect client that controls it. (breaking)
 - [metadata] Most metadata is now retrieved with the `spclient` (breaking)
diff --git a/core/src/connection/mod.rs b/core/src/connection/mod.rs
index b2e0356b9..8fb596f95 100644
--- a/core/src/connection/mod.rs
+++ b/core/src/connection/mod.rs
@@ -68,6 +68,29 @@ pub async fn connect(host: &str, port: u16, proxy: Option<&Url>) -> io::Result<T
     handshake(socket).await
 }
 
+pub async fn connect_with_retry(
+    host: &str,
+    port: u16,
+    proxy: Option<&Url>,
+    max_retries: u8,
+) -> io::Result<Transport> {
+    let mut num_retries = 0;
+    loop {
+        match connect(host, port, proxy).await {
+            Ok(f) => return Ok(f),
+            Err(e) => {
+                debug!("Connection failed: {e}");
+                if num_retries < max_retries {
+                    num_retries += 1;
+                    debug!("Retry access point...");
+                    continue;
+                }
+                return Err(e);
+            }
+        }
+    }
+}
+
 pub async fn authenticate(
     transport: &mut Transport,
     credentials: Credentials,
diff --git a/core/src/session.rs b/core/src/session.rs
index 3944ce83e..f934ed7b3 100644
--- a/core/src/session.rs
+++ b/core/src/session.rs
@@ -144,13 +144,15 @@ impl Session {
 
     async fn connect_inner(
         &self,
-        access_point: SocketAddress,
+        access_point: &SocketAddress,
         credentials: Credentials,
     ) -> Result<(Credentials, Transport), Error> {
-        let mut transport = connection::connect(
+        const MAX_RETRIES: u8 = 1;
+        let mut transport = connection::connect_with_retry(
             &access_point.0,
             access_point.1,
             self.config().proxy.as_ref(),
+            MAX_RETRIES,
         )
         .await?;
         let mut reusable_credentials = connection::authenticate(
@@ -165,10 +167,11 @@ impl Session {
             trace!(
                 "Reconnect using stored credentials as token authed sessions cannot use keymaster."
             );
-            transport = connection::connect(
+            transport = connection::connect_with_retry(
                 &access_point.0,
                 access_point.1,
                 self.config().proxy.as_ref(),
+                MAX_RETRIES,
             )
             .await?;
             reusable_credentials = connection::authenticate(
@@ -187,19 +190,32 @@ impl Session {
         credentials: Credentials,
         store_credentials: bool,
     ) -> Result<(), Error> {
+        // There currently happen to be 6 APs but anything will do to avoid an infinite loop.
+        const MAX_AP_TRIES: u8 = 6;
+        let mut num_ap_tries = 0;
         let (reusable_credentials, transport) = loop {
             let ap = self.apresolver().resolve("accesspoint").await?;
             info!("Connecting to AP \"{}:{}\"", ap.0, ap.1);
-            match self.connect_inner(ap, credentials.clone()).await {
+            match self.connect_inner(&ap, credentials.clone()).await {
                 Ok(ct) => break ct,
                 Err(e) => {
+                    num_ap_tries += 1;
+                    if MAX_AP_TRIES == num_ap_tries {
+                        error!("Tried too many access points");
+                        return Err(e);
+                    }
                     if let Some(AuthenticationError::LoginFailed(ErrorCode::TryAnotherAP)) =
                         e.error.downcast_ref::<AuthenticationError>()
                     {
                         warn!("Instructed to try another access point...");
                         continue;
-                    } else {
+                    } else if let Some(AuthenticationError::LoginFailed(..)) =
+                        e.error.downcast_ref::<AuthenticationError>()
+                    {
                         return Err(e);
+                    } else {
+                        warn!("Try another access point...");
+                        continue;
                     }
                 }
             }
@@ -567,7 +583,7 @@ impl Session {
     }
 
     pub fn shutdown(&self) {
-        debug!("Invalidating session");
+        debug!("Shutdown: Invalidating session");
         self.0.data.write().invalid = true;
         self.mercury().shutdown();
         self.channel().shutdown();
@@ -624,6 +640,7 @@ where
                     return Poll::Ready(Ok(()));
                 }
                 Some(Err(e)) => {
+                    error!("Connection to server closed.");
                     session.shutdown();
                     return Poll::Ready(Err(e));
                 }