macltok

#[cfg(feature = "sqlite")]

use std::path::{Path, PathBuf};

use matrix_sdk::{BoxFuture, SendOutsideWasm, SyncOutsideWasm};

use matrix_sdk_store_encryption::StoreCipher;

use tokio::sync::Mutex;

use crate::{AccountError, SessionRecord, StoreConfig, account::account_error};

#[cfg(feature = "sqlite")]

type RecordLocation = PathBuf;

#[cfg(feature = "indexeddb")]

type RecordLocation = String;

/// one login's record and the matrix store it belongs to. the engine rotates refresh tokens

/// and each works once, so a `save` has to land before the next refresh or the login is gone

pub trait SessionStore: SendOutsideWasm + SyncOutsideWasm {

    fn store(&self) -> &StoreConfig;

    fn load(&self) -> BoxFuture<'_, Result<Option<SessionRecord>, AccountError>>;

    /// replace the saved record, all or nothing

    fn save<'a>(&'a self, record: &'a SessionRecord) -> BoxFuture<'a, Result<(), AccountError>>;

    /// `save`, but without awaiting. the sdk calls it from inside a refresh, before the new

    /// refresh token is ever used, so a store that can write synchronously closes the gap

    /// a crash would otherwise lose the login through. `None` leaves the refresh to `save`

    #[cfg(not(target_family = "wasm"))]

    fn save_blocking(&self, _record: &SessionRecord) -> Option<Result<(), AccountError>> {

        None

    }

    /// swap the record for one that only says who to log in again as, keeping the store.

    /// `Session::restore` then answers `Restore::SoftLoggedOut`, even after a restart

    fn soft_log_out(&self) -> BoxFuture<'_, Result<(), AccountError>>;

    /// drop the record and delete the store, once nothing has the store open

    fn discard(&self) -> BoxFuture<'_, Result<(), AccountError>>;

}

/// keeps the record sealed with the store's passphrase, in a file natively and in an

/// IndexedDB database of its own in a browser

pub struct EncryptedSessionStore {

    record: RecordLocation,

    store: StoreConfig,

    sealer: Mutex<Option<Sealer>>,

}

struct Sealer {

    cipher: StoreCipher,

    exported: Vec<u8>,

}

impl EncryptedSessionStore {

    #[cfg(feature = "sqlite")]

    #[must_use]

    pub fn file(record: impl Into<PathBuf>, store: StoreConfig) -> Self {

        Self::new(record.into(), store)

    }

    /// the record lives beside the store, in `<name>::macltok-session`

    #[cfg(feature = "indexeddb")]

    #[must_use]

    pub fn indexeddb(store: StoreConfig) -> Self {

        let record = format!("{}::macltok-session", store.indexeddb_name());

        Self::new(record, store)

    }

    fn new(record: RecordLocation, store: StoreConfig) -> Self {

        Self {

            record,

            store,

            sealer: Mutex::new(None),

        }

    }

    async fn load_record(&self) -> Result<Option<SessionRecord>, AccountError> {

        let mut sealer = self.sealer.lock().await;

        self.read(&mut sealer).await

    }

    async fn save_record(&self, record: &SessionRecord) -> Result<(), AccountError> {

        let mut sealer = self.sealer.lock().await;

        self.write(&mut sealer, record).await

    }

    async fn soft_log_out_record(&self) -> Result<(), AccountError> {

        let mut sealer = self.sealer.lock().await;

        match self.read(&mut sealer).await? {

            Some(record) => self.write(&mut sealer, &record.soft_logged_out()).await,

            None => Ok(()),

        }

    }

    async fn discard_all(&self) -> Result<(), AccountError> {

        let mut sealer = self.sealer.lock().await;

        backend::remove(&self.record)

            .await

            .map_err(|error| account_error("removing the session record", error))?;

        *sealer = None;

        self.store.delete().await

    }

    async fn read(

        &self,

        sealer: &mut Option<Sealer>,

    ) -> Result<Option<SessionRecord>, AccountError> {

        let Some(sealed) = backend::read(&self.record)

            .await

            .map_err(|error| account_error("reading the session record", error))?

        else {

            return Ok(None);

        };

        let (opened, record) = unseal(self.store.passphrase(), &sealed)?;

        *sealer = Some(opened);

        Ok(Some(record))

    }

    async fn write(

        &self,

        sealer: &mut Option<Sealer>,

        record: &SessionRecord,

    ) -> Result<(), AccountError> {

        let sealed = self.seal(sealer, record)?;

        backend::write(&self.record, &sealed)

            .await

            .map_err(|error| account_error("writing the session record", error))

    }

    fn seal(

        &self,

        sealer: &mut Option<Sealer>,

        record: &SessionRecord,

    ) -> Result<Vec<u8>, AccountError> {

        let sealer = match sealer {

            Some(sealer) => sealer,

            empty => empty.insert(Sealer::new(self.store.passphrase())?),

        };

        sealer.seal(record)

    }

    /// the lock is only ever held around one small write, so losing the race to it just

    /// hands this refresh to the async save

    #[cfg(feature = "sqlite")]

    fn save_now(&self, record: &SessionRecord) -> Option<Result<(), AccountError>> {

        let mut sealer = self.sealer.try_lock().ok()?;

        Some(self.seal(&mut sealer, record).and_then(|sealed| {

            backend::write_blocking(&self.record, &sealed)

                .map_err(|error| account_error("writing the session record", error))

        }))

    }

}

impl SessionStore for EncryptedSessionStore {

    fn store(&self) -> &StoreConfig {

        &self.store

    }

    fn load(&self) -> BoxFuture<'_, Result<Option<SessionRecord>, AccountError>> {

        Box::pin(self.load_record())

    }

    fn save<'a>(&'a self, record: &'a SessionRecord) -> BoxFuture<'a, Result<(), AccountError>> {

        Box::pin(self.save_record(record))

    }

    #[cfg(feature = "sqlite")]

    fn save_blocking(&self, record: &SessionRecord) -> Option<Result<(), AccountError>> {

        self.save_now(record)

    }

    fn soft_log_out(&self) -> BoxFuture<'_, Result<(), AccountError>> {

        Box::pin(self.soft_log_out_record())

    }

    fn discard(&self) -> BoxFuture<'_, Result<(), AccountError>> {

        Box::pin(self.discard_all())

    }

}

impl Sealer {

    fn new(passphrase: &str) -> Result<Self, AccountError> {

        let cipher = StoreCipher::new()

            .map_err(|error| account_error("creating the session record cipher", error))?;

        let exported = cipher

            .export(passphrase)

            .map_err(|error| account_error("sealing the session record cipher", error))?;

        Ok(Self { cipher, exported })

    }

    fn seal(&self, record: &SessionRecord) -> Result<Vec<u8>, AccountError> {

        let value = self

            .cipher

            .encrypt_value(record)

            .map_err(|error| account_error("sealing the session record", error))?;

        let length = u32::try_from(self.exported.len())

            .expect("an exported store cipher is a few hundred bytes");

        let mut sealed = Vec::with_capacity(4 + self.exported.len() + value.len());

        sealed.extend_from_slice(&length.to_be_bytes());

        sealed.extend_from_slice(&self.exported);

        sealed.extend_from_slice(&value);

        Ok(sealed)

    }

}

fn unseal(passphrase: &str, sealed: &[u8]) -> Result<(Sealer, SessionRecord), AccountError> {

    let truncated = || {

        account_error(

            "opening the session record",

            std::io::Error::from(std::io::ErrorKind::UnexpectedEof),

        )

    };

    let (length, rest) = sealed.split_first_chunk::<4>().ok_or_else(truncated)?;

    let (exported, value) = rest

        .split_at_checked(u32::from_be_bytes(*length) as usize)

        .ok_or_else(truncated)?;

    let cipher = StoreCipher::import(passphrase, exported)

        .map_err(|error| account_error("opening the session record cipher", error))?;

    let record = cipher

        .decrypt_value(value)

        .map_err(|error| account_error("opening the session record", error))?;

    let sealer = Sealer {

        cipher,

        exported: exported.to_vec(),

    };

    Ok((sealer, record))

}

#[cfg(feature = "sqlite")]

mod backend {

    use std::{

        fs::{self, File, OpenOptions},

        io::{self, ErrorKind, Write},

        path::{Path, PathBuf},

    };

    pub(super) async fn read(path: &Path) -> io::Result<Option<Vec<u8>>> {

        match fs::read(path) {

            Ok(sealed) => Ok(Some(sealed)),

            Err(error) if error.kind() == ErrorKind::NotFound => Ok(None),

            Err(error) => Err(error),

        }

    }

    pub(super) async fn write(path: &Path, sealed: &[u8]) -> io::Result<()> {

        write_blocking(path, sealed)

    }

    pub(super) fn write_blocking(path: &Path, sealed: &[u8]) -> io::Result<()> {

        let mut temporary = path.as_os_str().to_owned();

        temporary.push(".tmp");

        let temporary = PathBuf::from(temporary);

        let parent = match path.parent() {

            Some(parent) if !parent.as_os_str().is_empty() => parent,

            _ => Path::new("."),

        };

        fs::create_dir_all(parent)?;

        let written = replace(path, &temporary, parent, sealed);

        if written.is_err() {

            let _ = fs::remove_file(&temporary);

        }

        written

    }

    fn replace(path: &Path, temporary: &Path, parent: &Path, sealed: &[u8]) -> io::Result<()> {

        let mut file = private_file(temporary)?;

        file.write_all(sealed)?;

        file.sync_all()?;

        fs::rename(temporary, path)?;

        sync_directory(parent)

    }

    pub(super) async fn remove(path: &Path) -> io::Result<()> {

        match fs::remove_file(path) {

            Err(error) if error.kind() != ErrorKind::NotFound => Err(error),

            _ => Ok(()),

        }

    }

    #[cfg(unix)]

    fn private_file(path: &Path) -> io::Result<File> {

        use std::os::unix::fs::OpenOptionsExt;

        OpenOptions::new()

            .write(true)

            .create(true)

            .truncate(true)

            .mode(0o600)

            .open(path)

    }

    #[cfg(not(unix))]

    fn private_file(path: &Path) -> io::Result<File> {

        OpenOptions::new()

            .write(true)

            .create(true)

            .truncate(true)

            .open(path)

    }

    /// a rename isn't durable until its directory is synced

    #[cfg(unix)]

    fn sync_directory(path: &Path) -> io::Result<()> {

        File::open(path)?.sync_all()

    }

    #[cfg(not(unix))]

    fn sync_directory(_: &Path) -> io::Result<()> {

        Ok(())

    }

}

#[cfg(feature = "indexeddb")]

mod backend {

    use indexed_db_futures::{

        Build,

        database::Database,

        error::{Error, OpenDbError},

        prelude::*,

        transaction::{Transaction, TransactionMode},

        typed_array::{Uint8Array, Uint8ArraySlice},

    };

    const RECORDS: &str = "records";

    const KEY: &str = "record";

    async fn open(name: &str) -> Result<Database, OpenDbError> {

        Database::open(name)

            .with_version(1u32)

            .with_on_upgrade_needed(|_, transaction: &Transaction<'_>| -> Result<(), Error> {

                transaction.db().create_object_store(RECORDS).build()?;

                Ok(())

            })

            .build()?

            .await

    }

    pub(super) async fn read(name: &str) -> Result<Option<Vec<u8>>, OpenDbError> {

        let database = open(name).await?;

        let sealed: Option<Uint8Array> = database

            .transaction(RECORDS)

            .with_mode(TransactionMode::Readonly)

            .build()?

            .object_store(RECORDS)?

            .get(KEY)

            .await?;

        Ok(sealed.map(Into::into))

    }

    pub(super) async fn write(name: &str, sealed: &[u8]) -> Result<(), OpenDbError> {

        let database = open(name).await?;

        let transaction = database

            .transaction(RECORDS)

            .with_mode(TransactionMode::Readwrite)

            .build()?;

        transaction

            .object_store(RECORDS)?

            .put(Uint8ArraySlice::new(sealed))

            .with_key(KEY)

            .build()?;

        Ok(transaction.commit().await?)

    }

    pub(super) async fn remove(name: &str) -> Result<(), OpenDbError> {

        Ok(Database::delete_by_name(name)?.await?)

    }

}

#[cfg(feature = "sqlite")]

pub(crate) fn delete_sqlite_store(path: &Path) -> std::io::Result<()> {

    use std::{fs, io::ErrorKind};

    // matrix-sdk-sqlite 0.18's database names, which it keeps private

    const DATABASES: [&str; 4] = [

        "matrix-sdk-state.sqlite3",

        "matrix-sdk-crypto.sqlite3",

        "matrix-sdk-event-cache.sqlite3",

        "matrix-sdk-media.sqlite3",

    ];

    for database in DATABASES {

        for suffix in ["", "-wal", "-shm", "-journal"] {

            match fs::remove_file(path.join(format!("{database}{suffix}"))) {

                Err(error) if error.kind() != ErrorKind::NotFound => return Err(error),

                _ => {}

            }

        }

    }

    match fs::remove_dir(path) {

        Err(error) if error.kind() != ErrorKind::NotFound => Err(error),

        _ => Ok(()),

    }

}

#[cfg(feature = "indexeddb")]

pub(crate) async fn delete_indexeddb_store(

    name: &str,

) -> Result<(), indexed_db_futures::error::OpenDbError> {

    use indexed_db_futures::database::Database;

    // the databases matrix-sdk-indexeddb 0.18 opens under one store name

    for suffix in [

        "",

        "::matrix-sdk-state",

        "::matrix-sdk-crypto",

        "::matrix-sdk-crypto-meta",

        "::event_cache",

        "::media",

    ] {

        Database::delete_by_name(&format!("{name}{suffix}"))?.await?;

    }

    Ok(())

}