macltok

#[cfg(feature = "sqlite")]

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

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

use matrix_sdk_store_encryption::StoreCipher;

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

#[cfg(feature = "sqlite")]

type RecordLocation = PathBuf;

#[cfg(feature = "indexeddb")]

type RecordLocation = String;

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

type Lock<T> = std::sync::Mutex<T>;

#[cfg(target_family = "wasm")]

type Lock<T> = tokio::sync::Mutex<T>;

/// 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: Lock<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: Lock::new(None),

        }

    }

    fn open(

        &self,

        sealer: &mut Option<Sealer>,

        sealed: Option<Vec<u8>>,

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

        let Some(sealed) = sealed else {

            return Ok(None);

        };

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

        *sealer = Some(opened);

        Ok(Some(record))

    }

    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)

    }

}

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

impl EncryptedSessionStore {

    fn lock(&self) -> std::sync::MutexGuard<'_, Option<Sealer>> {

        self.sealer

            .lock()

            .unwrap_or_else(std::sync::PoisonError::into_inner)

    }

    fn read(&self, sealer: &mut Option<Sealer>) -> Result<Option<SessionRecord>, AccountError> {

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

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

        self.open(sealer, sealed)

    }

    fn write(

        &self,

        sealer: &mut Option<Sealer>,

        record: &SessionRecord,

    ) -> Result<(), AccountError> {

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

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

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

    }

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

        self.read(&mut self.lock())

    }

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

        self.write(&mut self.lock(), record)

    }

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

        let mut sealer = self.lock();

        match self.read(&mut sealer)? {

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

            None => Ok(()),

        }

    }

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

        let mut sealer = self.lock();

        backend::remove(&self.record)

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

        *sealer = None;

        Ok(())

    }

}

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

impl SessionStore for EncryptedSessionStore {

    fn store(&self) -> &StoreConfig {

        &self.store

    }

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

        Box::pin(async move { self.load_record() })

    }

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

        Box::pin(async move { self.save_record(record) })

    }

    /// blocks the calling worker for the write and its fsyncs, about a millisecond

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

        Some(self.save_record(record))

    }

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

        Box::pin(async move { self.soft_log_out_record() })

    }

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

        Box::pin(async move {

            self.forget_record()?;

            self.store.delete().await

        })

    }

}

#[cfg(target_family = "wasm")]

impl EncryptedSessionStore {

    async fn read(

        &self,

        sealer: &mut Option<Sealer>,

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

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

            .await

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

        self.open(sealer, sealed)

    }

    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))

    }

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

        self.read(&mut *self.sealer.lock().await).await

    }

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

        self.write(&mut *self.sealer.lock().await, 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| error.into_account_error("removing the session record"))?;

        *sealer = None;

        self.store.delete().await

    }

}

#[cfg(target_family = "wasm")]

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))

    }

    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::{

        ffi::OsString,

        fs::{self, File, OpenOptions},

        io::{self, ErrorKind, Write},

        path::{Path, PathBuf},

        sync::atomic::{AtomicU64, Ordering},

    };

    static WRITES: AtomicU64 = AtomicU64::new(0);

    pub(super) 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) fn write(path: &Path, sealed: &[u8]) -> io::Result<()> {

        let parent = match path.parent() {

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

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

        };

        fs::create_dir_all(parent)?;

        let (temporary, mut file) = temporary_file(path)?;

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

        if written.is_err() {

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

        }

        written

    }

    fn temporary_file(path: &Path) -> io::Result<(PathBuf, File)> {

        loop {

            let mut name = OsString::from(path.as_os_str());

            let write = WRITES.fetch_add(1, Ordering::Relaxed);

            name.push(format!(".{}.{write}.tmp", std::process::id()));

            let temporary = PathBuf::from(name);

            match private_file(&temporary) {

                Ok(file) => return Ok((temporary, file)),

                Err(error) if error.kind() == ErrorKind::AlreadyExists => {}

                Err(error) => return Err(error),

            }

        }

    }

    fn replace(

        path: &Path,

        temporary: &Path,

        file: &mut File,

        parent: &Path,

        sealed: &[u8],

    ) -> io::Result<()> {

        file.write_all(sealed)?;

        file.sync_all()?;

        fs::rename(temporary, path)?;

        sync_directory(parent)

    }

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

        match fs::remove_file(path) {

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

            _ => Ok(()),

        }

    }

    /// `create_new` is `O_EXCL`, which posix says never follows a symlink, even a dangling one

    #[cfg(unix)]

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

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

        OpenOptions::new()

            .write(true)

            .create_new(true)

            .mode(0o600)

            .open(path)

    }

    #[cfg(not(unix))]

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

        OpenOptions::new().write(true).create_new(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

    }

    /// the crate's `Database` never closes itself on drop, and an open one blocks deleting it

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

        let database = open(name).await?;

        let sealed = get(&database).await;

        database.close();

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

    }

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

        let database = open(name).await?;

        let written = put(&database, sealed).await;

        database.close();

        Ok(written?)

    }

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

        super::delete_database(name).await

    }

    async fn get(database: &Database) -> Result<Option<Uint8Array>, Error> {

        database

            .transaction(RECORDS)

            .with_mode(TransactionMode::Readonly)

            .build()?

            .object_store(RECORDS)?

            .get(KEY)

            .await

    }

    async fn put(database: &Database, sealed: &[u8]) -> Result<(), Error> {