grapevine/src/database/key_value/uiaa.rs

use ruma::{
    api::client::{error::ErrorKind, uiaa::UiaaInfo},
    CanonicalJsonValue, DeviceId, UserId,
};

use crate::{database::KeyValueDatabase, service, Error, Result};

impl service::uiaa::Data for KeyValueDatabase {
    fn set_uiaa_request(
        &self,
        user_id: &UserId,
        device_id: &DeviceId,
        session: &str,
        request: &CanonicalJsonValue,
    ) -> Result<()> {
        self.userdevicesessionid_uiaarequest.write().unwrap().insert(
            (user_id.to_owned(), device_id.to_owned(), session.to_owned()),
            request.to_owned(),
        );

        Ok(())
    }

    fn get_uiaa_request(
        &self,
        user_id: &UserId,
        device_id: &DeviceId,
        session: &str,
    ) -> Option<CanonicalJsonValue> {
        self.userdevicesessionid_uiaarequest
            .read()
            .unwrap()
            .get(&(
                user_id.to_owned(),
                device_id.to_owned(),
                session.to_owned(),
            ))
            .map(ToOwned::to_owned)
    }

    fn update_uiaa_session(
        &self,
        user_id: &UserId,
        device_id: &DeviceId,
        session: &str,
        uiaainfo: Option<&UiaaInfo>,
    ) -> Result<()> {
        let mut userdevicesessionid = user_id.as_bytes().to_vec();
        userdevicesessionid.push(0xFF);
        userdevicesessionid.extend_from_slice(device_id.as_bytes());
        userdevicesessionid.push(0xFF);
        userdevicesessionid.extend_from_slice(session.as_bytes());

        if let Some(uiaainfo) = uiaainfo {
            self.userdevicesessionid_uiaainfo.insert(
                &userdevicesessionid,
                &serde_json::to_vec(&uiaainfo)
                    .expect("UiaaInfo::to_vec always works"),
            )?;
        } else {
            self.userdevicesessionid_uiaainfo.remove(&userdevicesessionid)?;
        }

        Ok(())
    }

    fn get_uiaa_session(
        &self,
        user_id: &UserId,
        device_id: &DeviceId,
        session: &str,
    ) -> Result<UiaaInfo> {
        let mut userdevicesessionid = user_id.as_bytes().to_vec();
        userdevicesessionid.push(0xFF);
        userdevicesessionid.extend_from_slice(device_id.as_bytes());
        userdevicesessionid.push(0xFF);
        userdevicesessionid.extend_from_slice(session.as_bytes());

        serde_json::from_slice(
            &self
                .userdevicesessionid_uiaainfo
                .get(&userdevicesessionid)?
                .ok_or(Error::BadRequest(
                    ErrorKind::forbidden(),
                    "UIAA session does not exist.",
                ))?,
        )
        .map_err(|_| {
            Error::bad_database("UiaaInfo in userdeviceid_uiaainfo is invalid.")
        })
    }
}