matrix-rust-sdk/matrix_sdk_base/src/state/json_store.rs

use std::{
    collections::HashMap,
    fmt, fs,
    path::{Path, PathBuf},
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
};

use matrix_sdk_common::{async_trait, identifiers::RoomId, locks::RwLock};
use tokio::{fs as async_fs, io::AsyncWriteExt};

use super::{AllRooms, ClientState, StateStore};
use crate::{Error, Result, Room, RoomState, Session};

/// A default `StateStore` implementation that serializes state as json
/// and saves it to disk.
///
/// When logged in the `JsonStore` appends the user_id to its folder path,
/// so all files are saved in `my_client/user_id_localpart/*`.
pub struct JsonStore {
    path: Arc<RwLock<PathBuf>>,
    user_path_set: AtomicBool,
}

impl JsonStore {
    /// Create a `JsonStore` to store the client and room state.
    ///
    /// Checks if the provided path exists and creates the directories if not.
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let p = path.as_ref();
        if !p.exists() {
            fs::create_dir_all(p)?;
        }
        Ok(Self {
            path: Arc::new(RwLock::new(p.to_path_buf())),
            user_path_set: AtomicBool::new(false),
        })
    }

    /// Build a path for a file where the Room state to be stored in.
    async fn build_room_path(&self, room_state: &str, room_id: &RoomId) -> PathBuf {
        let mut path = self.path.read().await.clone();

        path.push("rooms");
        path.push(room_state);
        path.push(JsonStore::sanitize_room_id(room_id));
        path.set_extension("json");

        path
    }

    /// Build a path for the file where the Client state to be stored in.
    async fn build_client_path(&self) -> PathBuf {
        let mut path = self.path.read().await.clone();
        path.push("client");
        path.set_extension("json");

        path
    }

    /// Replace common characters that can't be used in a file name with an
    /// underscore.
    fn sanitize_room_id(room_id: &RoomId) -> String {
        room_id.as_str().replace(
            &['.', ':', '<', '>', '"', '/', '\\', '|', '?', '*'][..],
            "_",
        )
    }
}

impl fmt::Debug for JsonStore {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("JsonStore")
            .field("path", &self.path)
            .finish()
    }
}

#[async_trait]
impl StateStore for JsonStore {
    async fn load_client_state(&self, sess: &Session) -> Result<Option<ClientState>> {
        if !self.user_path_set.load(Ordering::SeqCst) {
            self.user_path_set.swap(true, Ordering::SeqCst);
            self.path.write().await.push(sess.user_id.localpart())
        }

        let path = self.build_client_path().await;

        let json = async_fs::read_to_string(path)
            .await
            .map_or(String::default(), |s| s);
        if json.is_empty() {
            Ok(None)
        } else {
            serde_json::from_str(&json).map(Some).map_err(Error::from)
        }
    }

    async fn load_all_rooms(&self) -> Result<AllRooms> {
        let mut path = self.path.read().await.clone();
        path.push("rooms");

        let mut joined = HashMap::new();
        let mut left = HashMap::new();
        let mut invited = HashMap::new();
        for room_state_type in &["joined", "invited", "left"] {
            path.push(room_state_type);
            // don't load rooms that aren't saved yet
            if !path.exists() {
                path.pop();
                continue;
            }

            for file in fs::read_dir(&path)? {
                let file = file?.path();

                if file.is_dir() {
                    continue;
                }

                let json = async_fs::read_to_string(&file).await?;
                let room = serde_json::from_str::<Room>(&json).map_err(Error::from)?;
                let room_id = room.room_id.clone();

                match *room_state_type {
                    "joined" => joined.insert(room_id, room),
                    "invited" => invited.insert(room_id, room),
                    "left" => left.insert(room_id, room),
                    _ => unreachable!("an array with 3 const elements was altered in JsonStore"),
                };
            }
            path.pop();
        }

        Ok(AllRooms {
            joined,
            left,
            invited,
        })
    }

    async fn store_client_state(&self, state: ClientState) -> Result<()> {
        let path = self.build_client_path().await;

        if !path.exists() {
            let mut dir = path.clone();
            dir.pop();
            async_fs::create_dir_all(dir).await?;
        }

        let json = serde_json::to_string(&state).map_err(Error::from)?;

        let mut file = async_fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)
            .await?;
        file.write_all(json.as_bytes()).await.map_err(Error::from)
    }

    async fn store_room_state(&self, room: RoomState<&Room>) -> Result<()> {
        let (room, room_state) = match room {
            RoomState::Joined(room) => (room, "joined"),
            RoomState::Invited(room) => (room, "invited"),
            RoomState::Left(room) => (room, "left"),
        };

        if !self.user_path_set.load(Ordering::SeqCst) {
            self.user_path_set.swap(true, Ordering::SeqCst);
            self.path.write().await.push(room.own_user_id.localpart())
        }

        let path = self.build_room_path(room_state, &room.room_id).await;

        if !path.exists() {
            let mut dir = path.clone();
            dir.pop();
            async_fs::create_dir_all(dir).await?;
        }

        let json = serde_json::to_string(&room).map_err(Error::from)?;

        let mut file = async_fs::OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(path)
            .await?;
        file.write_all(json.as_bytes()).await.map_err(Error::from)
    }

    async fn delete_room_state(&self, room: RoomState<&RoomId>) -> Result<()> {
        let (room_id, room_state) = match &room {
            RoomState::Joined(id) => (id, "joined"),
            RoomState::Invited(id) => (id, "invited"),
            RoomState::Left(id) => (id, "left"),
        };

        if !self.user_path_set.load(Ordering::SeqCst) {
            return Err(Error::StateStore("path for JsonStore not set".into()));
        }

        let path = self.build_room_path(room_state, room_id).await;

        if !path.exists() {
            return Err(Error::StateStore(format!("file {:?} not found", path)));
        }

        tokio::fs::remove_file(path).await.map_err(Error::from)
    }
}

#[cfg(test)]
mod test {}