Reauth on invalid session (#50)

* Reauth on invalid session

* fix to actually use new token and retry auth

* fix unused vars

2 files changed, 90 insertions, 13 deletions

diff --git a/azalea-client/src/account.rs b/azalea-client/src/account.rs
index 42bfe6fc..0d758507 100755
--- a/azalea-client/src/account.rs
+++ b/azalea-client/src/account.rs
@@ -1,6 +1,9 @@
 //! Connect to Minecraft servers.
 
+use std::sync::Arc;
+
 use crate::get_mc_dir;
+use parking_lot::Mutex;
 use uuid::Uuid;
 
 /// Something that can join Minecraft servers.
@@ -24,9 +27,25 @@ pub struct Account {
     pub username: String,
     /// The access token for authentication. You can obtain one of these
     /// manually from azalea-auth.
-    pub access_token: Option<String>,
+    ///
+    /// This is an Arc<Mutex> so it can be modified by [`Self::refresh`].
+    pub access_token: Option<Arc<Mutex<String>>>,
     /// Only required for online-mode accounts.
     pub uuid: Option<uuid::Uuid>,
+
+    /// The parameters (i.e. email) that were passed for creating this
+    /// [`Account`]. This is used to for automatic reauthentication when we get
+    /// "Invalid Session" errors. If you don't need that feature (like in
+    /// offline mode), then you can set this to `AuthOpts::default()`.
+    pub auth_opts: AuthOpts,
+}
+
+/// The parameters that were passed for creating the associated [`Account`].
+#[derive(Clone, Debug)]
+pub enum AuthOpts {
+    Offline { username: String },
+    // this is an enum so legacy Mojang auth can be added in the future
+    Microsoft { email: String },
 }
 
 impl Account {
@@ -38,6 +57,9 @@ impl Account {
             username: username.to_string(),
             access_token: None,
             uuid: None,
+            auth_opts: AuthOpts::Offline {
+                username: username.to_string(),
+            },
         }
     }
 
@@ -62,8 +84,36 @@ impl Account {
         .await?;
         Ok(Self {
             username: auth_result.profile.name,
-            access_token: Some(auth_result.access_token),
+            access_token: Some(Arc::new(Mutex::new(auth_result.access_token))),
             uuid: Some(Uuid::parse_str(&auth_result.profile.id).expect("Invalid UUID")),
+            auth_opts: AuthOpts::Microsoft {
+                email: email.to_string(),
+            },
         })
     }
+
+    /// Refresh the access_token for this account to be valid again.
+    ///
+    /// This requires the `auth_opts` field to be set correctly (which is done
+    /// by default if you used the constructor functions). Note that if the
+    /// Account is offline-mode, this function won't do anything.
+    pub async fn refresh(&self) -> Result<(), azalea_auth::AuthError> {
+        match &self.auth_opts {
+            // offline mode doesn't need to refresh so just don't do anything lol
+            AuthOpts::Offline { .. } => Ok(()),
+            AuthOpts::Microsoft { email } => {
+                let new_account = Account::microsoft(email).await?;
+                let access_token = self
+                    .access_token.as_ref()
+                    .expect("Access token should always be set for Microsoft accounts");
+                let new_access_token = new_account
+                    .access_token
+                    .expect("Access token should always be set for Microsoft accounts")
+                    .lock()
+                    .clone();
+                *access_token.lock() = new_access_token;
+                Ok(())
+            }
+        }
+    }
 }
diff --git a/azalea-client/src/client.rs b/azalea-client/src/client.rs
index 17fb4840..4f0a2ed6 100644
--- a/azalea-client/src/client.rs
+++ b/azalea-client/src/client.rs
@@ -1,6 +1,6 @@
 pub use crate::chat::ChatPacket;
 use crate::{movement::WalkDirection, plugins::PluginStates, Account, PlayerInfo};
-use azalea_auth::game_profile::GameProfile;
+use azalea_auth::{game_profile::GameProfile, sessionserver::SessionServerError};
 use azalea_core::{ChunkPos, ResourceLocation, Vec3};
 use azalea_protocol::{
     connect::{Connection, ConnectionError, ReadConnection, WriteConnection},
@@ -142,6 +142,8 @@ pub enum JoinError {
     SessionServer(#[from] azalea_auth::sessionserver::SessionServerError),
     #[error("The given address could not be parsed into a ServerAddress")]
     InvalidAddress,
+    #[error("Couldn't refresh access token: {0}")]
+    Auth(#[from] azalea_auth::AuthError),
 }
 
 #[derive(Error, Debug)]
@@ -239,7 +241,11 @@ impl Client {
         Ok((client, rx))
     }
 
-    /// Do a handshake with the server and get to the game state from the initial handshake state.
+    /// Do a handshake with the server and get to the game state from the
+    /// initial handshake state.
+    ///
+    /// This will also automatically refresh the account's access token if
+    /// it's expired.
     pub async fn handshake(
         mut conn: Connection<ClientboundHandshakePacket, ServerboundHandshakePacket>,
         account: &Account,
@@ -282,15 +288,36 @@ impl Client {
                     let e = azalea_crypto::encrypt(&p.public_key, &p.nonce).unwrap();
 
                     if let Some(access_token) = &account.access_token {
-                        conn.authenticate(
-                            access_token,
-                            &account
-                                .uuid
-                                .expect("Uuid must be present if access token is present."),
-                            e.secret_key,
-                            p,
-                        )
-                        .await?;
+                        // keep track of the number of times we tried
+                        // authenticating so we can give up after too many
+                        let mut attempts: usize = 1;
+
+                        while let Err(e) = {
+                            let access_token = access_token.lock().clone();
+                            conn.authenticate(
+                                &access_token,
+                                &account
+                                    .uuid
+                                    .expect("Uuid must be present if access token is present."),
+                                e.secret_key,
+                                &p,
+                            )
+                            .await
+                        } {
+                            if attempts >= 2 {
+                                // if this is the second attempt and we failed
+                                // both times, give up
+                                return Err(e.into());
+                            }
+                            if let SessionServerError::InvalidSession = e {
+                                // uh oh, we got an invalid session and have
+                                // to reauthenticate now
+                                account.refresh().await?;
+                            } else {
+                                return Err(e.into());
+                            }
+                            attempts += 1;
+                        }
                     }
 
                     conn.write(