Can save Matrix session after authentication

matrixgw_backend/src/matrix_connection/matrix_client.rs

@@ -1,15 +1,29 @@
 use crate::app_config::AppConfig;
 use crate::users::UserEmail;
 use crate::utils::rand_utils::rand_string;
+use anyhow::Context;
 use matrix_sdk::authentication::oauth::error::OAuthDiscoveryError;
-use matrix_sdk::authentication::oauth::{OAuthError, UrlOrQuery};
+use matrix_sdk::authentication::oauth::{ClientId, OAuthError, UrlOrQuery, UserSession};
 use matrix_sdk::ruma::serde::Raw;
 use matrix_sdk::{Client, ClientBuildError};
+use serde::{Deserialize, Serialize};
 use url::Url;
 
+/// The full session to persist.
+#[derive(Debug, Serialize, Deserialize, Clone)]
+struct StoredSession {
+    /// The OAuth 2.0 user session.
+    user_session: UserSession,
+
+    /// The OAuth 2.0 client ID.
+    client_id: ClientId,
+}
+
 /// Matrix Gateway session errors
 #[derive(thiserror::Error, Debug)]
 enum MatrixClientError {
+    #[error("Failed to destroy previous client data! {0}")]
+    DestroyPreviousData(Box<MatrixClientError>),
     #[error("Failed to create Matrix database storage directory! {0}")]
     CreateMatrixDbDir(std::io::Error),
     #[error("Failed to create database passphrase! {0}")]
@@ -18,6 +32,8 @@ enum MatrixClientError {
     ReadDbPassphrase(std::io::Error),
     #[error("Failed to build Matrix client! {0}")]
     BuildMatrixClient(ClientBuildError),
+    #[error("Failed to clear Matrix session file! {0}")]
+    ClearMatrixSessionFile(std::io::Error),
     #[error("Failed to clear Matrix database storage directory! {0}")]
     ClearMatrixDbDir(std::io::Error),
     #[error("Failed to remove database passphrase! {0}")]
@@ -30,6 +46,8 @@ enum MatrixClientError {
     BuildAuthRequest(OAuthError),
     #[error("Failed to finalize authentication! {0}")]
     FinishLogin(matrix_sdk::Error),
+    #[error("Failed to write session file! {0}")]
+    WriteSessionFile(std::io::Error),
 }
 
 #[derive(serde::Deserialize)]
@@ -47,6 +65,15 @@ pub struct MatrixClient {
 impl MatrixClient {
     /// Start to build Matrix client to initiate user authentication
     pub async fn build_client(email: &UserEmail) -> anyhow::Result<Self> {
+        // Check if we are restoring a previous state
+        let is_restoring = AppConfig::get()
+            .user_matrix_session_file_path(email)
+            .is_file();
+        if !is_restoring {
+            Self::destroy_data(email).map_err(MatrixClientError::DestroyPreviousData)?;
+        }
+
+        // Determine Matrix database path
         let db_path = AppConfig::get().user_matrix_db_path(email);
         std::fs::create_dir_all(&db_path).map_err(MatrixClientError::CreateMatrixDbDir)?;
 
@@ -69,34 +96,46 @@ impl MatrixClient {
             .map_err(MatrixClientError::BuildMatrixClient)?;
 
         // Check metadata
-        let server_metadata = client
-            .oauth()
+        let oauth = client.oauth();
+        let server_metadata = oauth
             .server_metadata()
             .await
             .map_err(MatrixClientError::FetchServerMetadata)?;
         log::info!("OAuth2 server issuer: {:?}", server_metadata.issuer);
 
+        // TODO : restore client ID to oauth if needed
         // TODO : restore client if client already existed
 
-        Ok(Self {
+        let client = Self {
             email: email.clone(),
             client,
-        })
+        };
+
+        // Automatically save session when token gets refreshed
+        client.setup_background_session_save().await;
+
+        Ok(client)
     }
 
-    /// Destroy this Matrix client instance
-    pub fn destroy(&self) -> anyhow::Result<()> {
-        let db_path = AppConfig::get().user_matrix_db_path(&self.email);
-        if db_path.is_file() {
+    /// Destroy Matrix client related data
+    fn destroy_data(email: &UserEmail) -> anyhow::Result<(), Box<MatrixClientError>> {
+        let session_path = AppConfig::get().user_matrix_session_file_path(email);
+        if session_path.is_file() {
+            std::fs::remove_file(&session_path)
+                .map_err(MatrixClientError::ClearMatrixSessionFile)?;
+        }
+
+        let db_path = AppConfig::get().user_matrix_db_path(email);
+        if db_path.is_dir() {
             std::fs::remove_dir_all(&db_path).map_err(MatrixClientError::ClearMatrixDbDir)?;
         }
 
-        let passphrase_path = AppConfig::get().user_matrix_passphrase_path(&self.email);
+        let passphrase_path = AppConfig::get().user_matrix_passphrase_path(email);
         if passphrase_path.is_file() {
             std::fs::remove_file(passphrase_path).map_err(MatrixClientError::ClearDbPassphrase)?;
         }
 
-        todo!()
+        Ok(())
     }
 
     /// Initiate OAuth authentication
@@ -137,6 +176,60 @@ impl MatrixClient {
             self.client.user_id().unwrap()
         );
 
+        // Persist session tokens
+        self.save_stored_session().await?;
+
+        Ok(())
+    }
+
+    /// Automatically persist session onto disk
+    pub async fn setup_background_session_save(&self) {
+        let this = self.clone();
+        tokio::spawn(async move {
+            while let Ok(update) = this.client.subscribe_to_session_changes().recv().await {
+                match update {
+                    matrix_sdk::SessionChange::UnknownToken { soft_logout } => {
+                        log::warn!("Received an unknown token error; soft logout? {soft_logout:?}");
+                    }
+                    matrix_sdk::SessionChange::TokensRefreshed => {
+                        // The tokens have been refreshed, persist them to disk.
+                        if let Err(err) = this.save_stored_session().await {
+                            log::error!("Unable to store a session in the background: {err}");
+                        }
+                    }
+                }
+            }
+        });
+    }
+
+    /// Update the session stored on the filesystem.
+    async fn save_stored_session(&self) -> anyhow::Result<()> {
+        log::debug!("Save the stored session for {:?}...", self.email);
+
+        let user_session: UserSession = self
+            .client
+            .oauth()
+            .user_session()
+            .context("A logged in client must have a session")?;
+
+        let stored_session = StoredSession {
+            user_session,
+            client_id: self
+                .client
+                .oauth()
+                .client_id()
+                .context("Client ID should be set at this point!")?
+                .clone(),
+        };
+
+        let serialized_session = serde_json::to_string(&stored_session)?;
+        std::fs::write(
+            AppConfig::get().user_matrix_session_file_path(&self.email),
+            serialized_session,
+        )
+        .map_err(MatrixClientError::WriteSessionFile)?;
+
+        log::debug!("Updating the stored session: done!");
         Ok(())
     }
 }