matrix-rust-sdk

use std::{

    collections::{BTreeMap, HashMap},

    io::Cursor,

    mem::ManuallyDrop,

    ops::Deref,

    sync::Arc,

    time::Duration,

};

use js_int::UInt;

use matrix_sdk_common::deserialized_responses::AlgorithmInfo;

use matrix_sdk_crypto::{

    CollectStrategy, DecryptionSettings, LocalTrust, OlmMachine as InnerMachine,

    UserIdentity as SdkUserIdentity,

    backups::{

        MegolmV1BackupKey as RustBackupKey, SignatureState,

        SignatureVerification as RustSignatureCheckResult,

    },

    decrypt_room_key_export, encrypt_room_key_export,

    olm::ExportedRoomKey,

    store::types::{BackupDecryptionKey, Changes},

    types::requests::ToDeviceRequest,

};

use ruma::{

    DeviceKeyAlgorithm, EventId, OneTimeKeyAlgorithm, OwnedTransactionId, OwnedUserId, RoomId,

    UserId,

    api::{

        IncomingResponse,

        client::{

            backup::add_backup_keys::v3::Response as KeysBackupResponse,

            keys::{

                claim_keys::v3::Response as KeysClaimResponse,

                get_keys::v3::Response as KeysQueryResponse,

                upload_keys::v3::Response as KeysUploadResponse,

                upload_signatures::v3::Response as SignatureUploadResponse,

            },

            message::send_message_event::v3::Response as RoomMessageResponse,

            sync::sync_events::{DeviceLists as RumaDeviceLists, v3::ToDevice},

            to_device::send_event_to_device::v3::Response as ToDeviceResponse,

        },

    },

    events::{

        AnyMessageLikeEvent, AnySyncMessageLikeEvent, AnyTimelineEvent, MessageLikeEvent,

        key::verification::VerificationMethod, room::message::MessageType,

    },

    serde::Raw,

    to_device::DeviceIdOrAllDevices,

};

use serde::{Deserialize, Serialize};

use serde_json::{Value, value::RawValue};

use tokio::runtime::Runtime;

use zeroize::Zeroize;

use crate::{

    BackupKeys, BackupRecoveryKey, BootstrapCrossSigningResult, CrossSigningKeyExport,

    CrossSigningStatus, DecodeError, DecryptedEvent, Device, DeviceLists, EncryptionSettings,

    EventEncryptionAlgorithm, KeyImportError, KeysImportResult, MegolmV1BackupKey,

    ProgressListener, Request, RequestType, RequestVerificationResult, RoomKeyCounts, RoomSettings,

    Sas, SignatureUploadRequest, StartSasResult, UserIdentity, Verification, VerificationRequest,

    dehydrated_devices::DehydratedDevices,

    error::{

        CryptoStoreError, DecryptionError, SecretImportError, SecretsBundleExportError,

        SignatureError,

    },

    parse_user_id,

    responses::{OwnedResponse, response_from_string},

};

/// The return value for the [`OlmMachine::receive_sync_changes()`] method.

///

/// Will contain various information about the `/sync` changes the

/// [`OlmMachine`] processed.

#[derive(uniffi::Record)]

pub struct SyncChangesResult {

    /// The, now possibly decrypted, to-device events the [`OlmMachine`]

    /// received, decrypted, and processed.

    to_device_events: Vec<String>,

    /// Information about the room keys that were extracted out of the to-device

    /// events.

    room_key_infos: Vec<RoomKeyInfo>,

}

/// Information on a room key that has been received or imported.

#[derive(uniffi::Record)]

pub struct RoomKeyInfo {

    /// The [messaging algorithm] that this key is used for. Will be one of the

    /// `m.megolm.*` algorithms.

    ///

    /// [messaging algorithm]: https://spec.matrix.org/v1.6/client-server-api/#messaging-algorithms

    pub algorithm: String,

    /// The room where the key is used.

    pub room_id: String,

    /// The Curve25519 key of the device which initiated the session originally.

    pub sender_key: String,

    /// The ID of the session that the key is for.

    pub session_id: String,

}

impl From<matrix_sdk_crypto::store::types::RoomKeyInfo> for RoomKeyInfo {

    fn from(value: matrix_sdk_crypto::store::types::RoomKeyInfo) -> Self {

        Self {

            algorithm: value.algorithm.to_string(),

            room_id: value.room_id.to_string(),

            sender_key: value.sender_key.to_base64(),

            session_id: value.session_id,

        }

    }

}

/// A high level state machine that handles E2EE for Matrix.

#[derive(uniffi::Object)]

pub struct OlmMachine {

    pub(crate) inner: ManuallyDrop<InnerMachine>,

    pub(crate) runtime: Runtime,

}

impl Drop for OlmMachine {

    fn drop(&mut self) {

        // Dropping the inner OlmMachine must happen within a tokio context

        // because deadpool drops sqlite connections in the DB pool on tokio's

        // blocking threadpool to avoid blocking async worker threads.

        let _guard = self.runtime.enter();

        // SAFETY: self.inner is never used again, which is the only requirement

        //         for ManuallyDrop::drop to be used safely.

        unsafe {

            ManuallyDrop::drop(&mut self.inner);

        }

    }

}

/// A pair of outgoing room key requests, both of those are sendToDevice

/// requests.

#[derive(uniffi::Record)]

pub struct KeyRequestPair {

    /// The optional cancellation, this is None if no previous key request was

    /// sent out for this key, thus it doesn't need to be cancelled.

    pub cancellation: Option<Request>,

    /// The actual key request.

    pub key_request: Request,

}

/// The result of a signature verification of a signed JSON object.

#[derive(Clone, Debug, PartialEq, Eq, uniffi::Record)]

pub struct SignatureVerification {

    /// The result of the signature verification using the public key of our own

    /// device.

    pub device_signature: SignatureState,

    /// The result of the signature verification using the public key of our own

    /// user identity.

    pub user_identity_signature: SignatureState,

    /// The result of the signature verification using public keys of other

    /// devices we own.

    pub other_devices_signatures: HashMap<String, SignatureState>,

    /// Is the signed JSON object trusted.

    ///

    /// This flag tells us if the result has a valid signature from any of the

    /// following:

    ///

    /// * Our own device

    /// * Our own user identity, provided the identity is trusted as well

    /// * Any of our own devices, provided the device is trusted as well

    pub trusted: bool,

}

impl From<RustSignatureCheckResult> for SignatureVerification {

    fn from(r: RustSignatureCheckResult) -> Self {

        let trusted = r.trusted();

        Self {

            device_signature: r.device_signature,

            user_identity_signature: r.user_identity_signature,

            other_devices_signatures: r

                .other_signatures

                .into_iter()

                .map(|(k, v)| (k.to_string(), v))

                .collect(),

            trusted,

        }

    }

}

#[matrix_sdk_ffi_macros::export]

impl OlmMachine {

    /// Create a new `OlmMachine`

    ///

    /// # Arguments

    ///

    /// * `user_id` - The unique ID of the user that owns this machine.

    ///

    /// * `device_id` - The unique ID of the device that owns this machine.

    ///

    /// * `path` - The path where the state of the machine should be persisted.

    ///

    /// * `passphrase` - The passphrase that should be used to encrypt the data

    ///   at rest in the crypto store. **Warning**, if no passphrase is given,

    ///   the store and all its data will remain unencrypted.

    #[uniffi::constructor]

    pub fn new(

        user_id: String,

        device_id: String,

        path: String,

        mut passphrase: Option<String>,

    ) -> Result<Arc<Self>, CryptoStoreError> {

        let user_id = parse_user_id(&user_id)?;

        let device_id = device_id.as_str().into();

        let runtime = Runtime::new().expect("Couldn't create a tokio runtime");

        let store = runtime

            .block_on(matrix_sdk_sqlite::SqliteCryptoStore::open(path, passphrase.as_deref()))?;

        passphrase.zeroize();

        let inner = runtime.block_on(InnerMachine::with_store(

            &user_id,

            device_id,

            Arc::new(store),

            None,

        ))?;

        Ok(Arc::new(OlmMachine { inner: ManuallyDrop::new(inner), runtime }))

    }

    /// Get the user ID of the owner of this `OlmMachine`.

    pub fn user_id(&self) -> String {

        self.inner.user_id().to_string()

    }

    /// Get the device ID of the device of this `OlmMachine`.

    pub fn device_id(&self) -> String {

        self.inner.device_id().to_string()

    }

    /// Get our own identity keys.

    pub fn identity_keys(&self) -> HashMap<String, String> {

        let identity_keys = self.inner.identity_keys();

        let curve_key = identity_keys.curve25519.to_base64();

        let ed25519_key = identity_keys.ed25519.to_base64();

        HashMap::from([("ed25519".to_owned(), ed25519_key), ("curve25519".to_owned(), curve_key)])

    }

    /// Get the status of the private cross signing keys.

    ///

    /// This can be used to check which private cross signing keys we have

    /// stored locally.

    pub fn cross_signing_status(&self) -> CrossSigningStatus {

        self.runtime.block_on(self.inner.cross_signing_status()).into()

    }

    /// Get a cross signing user identity for the given user ID.

    ///

    /// # Arguments

    ///

    /// * `user_id` - The unique id of the user that the identity belongs to

    ///

    /// * `timeout` - The time in seconds we should wait before returning if the

    ///   user's device list has been marked as stale. Passing a 0 as the

    ///   timeout means that we won't wait at all. **Note**, this assumes that

    ///   the requests from [`OlmMachine::outgoing_requests`] are being

    ///   processed and sent out. Namely, this waits for a `/keys/query`

    ///   response to be received.

    pub fn get_identity(

        &self,

        user_id: String,

        timeout: u32,

    ) -> Result<Option<UserIdentity>, CryptoStoreError> {

        let user_id = parse_user_id(&user_id)?;

        let timeout = if timeout == 0 { None } else { Some(Duration::from_secs(timeout.into())) };

        Ok(

            if let Some(identity) =

                self.runtime.block_on(self.inner.get_identity(&user_id, timeout))?

            {

                Some(self.runtime.block_on(UserIdentity::from_rust(identity))?)

            } else {

                None

            },

        )

    }

    /// Check if a user identity is considered to be verified by us.

    pub fn is_identity_verified(&self, user_id: String) -> Result<bool, CryptoStoreError> {

        let user_id = parse_user_id(&user_id)?;

        Ok(

            if let Some(identity) =

                self.runtime.block_on(self.inner.get_identity(&user_id, None))?

            {

                identity.is_verified()

            } else {

                false

            },

        )

    }

    /// Manually the user with the given user ID.

    ///

    /// This method will attempt to sign the user identity using either our

    /// private cross signing key, for other user identities, or our device keys

    /// for our own user identity.

    ///

    /// This method can fail if we don't have the private part of our

    /// user-signing key.

    ///

    /// Returns a request that needs to be sent out for the user identity to be

    /// marked as verified.

    pub fn verify_identity(

        &self,

        user_id: String,

    ) -> Result<SignatureUploadRequest, SignatureError> {

        let user_id = UserId::parse(user_id)?;

        let user_identity = self.runtime.block_on(self.inner.get_identity(&user_id, None))?;

        if let Some(user_identity) = user_identity {

            Ok(match user_identity {

                SdkUserIdentity::Own(i) => self.runtime.block_on(i.verify())?,

                SdkUserIdentity::Other(i) => self.runtime.block_on(i.verify())?,

            }

            .into())

        } else {

            Err(SignatureError::UnknownUserIdentity(user_id.to_string()))

        }

    }

    /// Get a `Device` from the store.

    ///

    /// # Arguments

    ///

    /// * `user_id` - The id of the device owner.

    ///

    /// * `device_id` - The id of the device itself.

    ///

    /// * `timeout` - The time in seconds we should wait before returning if the

    ///   user's device list has been marked as stale. Passing a 0 as the

    ///   timeout means that we won't wait at all. **Note**, this assumes that

    ///   the requests from [`OlmMachine::outgoing_requests`] are being

    ///   processed and sent out. Namely, this waits for a `/keys/query`

    ///   response to be received.

    pub fn get_device(

        &self,

        user_id: String,

        device_id: String,

        timeout: u32,

    ) -> Result<Option<Device>, CryptoStoreError> {

        let user_id = parse_user_id(&user_id)?;

        let timeout = if timeout == 0 { None } else { Some(Duration::from_secs(timeout.into())) };

        Ok(self

            .runtime

            .block_on(self.inner.get_device(&user_id, device_id.as_str().into(), timeout))?

            .map(|d| d.into()))

    }

    /// Manually verify the device of the given user with the given device ID.

    ///

    /// This method will attempt to sign the device using our private cross

    /// signing key.

    ///

    /// This method will always fail if the device belongs to someone else, we

    /// can only sign our own devices.

    ///

    /// It can also fail if we don't have the private part of our self-signing

    /// key.

    ///

    /// Returns a request that needs to be sent out for the device to be marked

    /// as verified.

    pub fn verify_device(

        &self,

        user_id: String,

        device_id: String,

    ) -> Result<SignatureUploadRequest, SignatureError> {

        let user_id = UserId::parse(user_id)?;

        let device = self.runtime.block_on(self.inner.get_device(

            &user_id,

            device_id.as_str().into(),

            None,

        ))?;

        if let Some(device) = device {

            Ok(self.runtime.block_on(device.verify())?.into())

        } else {

            Err(SignatureError::UnknownDevice(user_id, device_id))

        }

    }

    /// Set local trust state for the device of the given user without creating

    /// or uploading any signatures if verified

    pub fn set_local_trust(

        &self,

        user_id: String,

        device_id: String,

        trust_state: LocalTrust,

    ) -> Result<(), CryptoStoreError> {

        let user_id = parse_user_id(&user_id)?;

        let device = self.runtime.block_on(self.inner.get_device(

            &user_id,

            device_id.as_str().into(),

            None,

        ))?;

        if let Some(device) = device {

            self.runtime.block_on(device.set_local_trust(trust_state))?;

        }

        Ok(())

    }

    /// Get all devices of an user.

    ///

    /// # Arguments

    ///

    /// * `user_id` - The id of the device owner.

    ///

    /// * `timeout` - The time in seconds we should wait before returning if the

    ///   user's device list has been marked as stale. Passing a 0 as the

    ///   timeout means that we won't wait at all. **Note**, this assumes that

    ///   the requests from [`OlmMachine::outgoing_requests`] are being

    ///   processed and sent out. Namely, this waits for a `/keys/query`

    ///   response to be received.

    pub fn get_user_devices(

        &self,

        user_id: String,

        timeout: u32,

    ) -> Result<Vec<Device>, CryptoStoreError> {

        let user_id = parse_user_id(&user_id)?;

        let timeout = if timeout == 0 { None } else { Some(Duration::from_secs(timeout.into())) };

        Ok(self

            .runtime

            .block_on(self.inner.get_user_devices(&user_id, timeout))?

            .devices()

            .map(|d| d.into())

            .collect())

    }

    /// Get the list of outgoing requests that need to be sent to the

    /// homeserver.

    ///

    /// After the request was sent out and a successful response was received

    /// the response body should be passed back to the state machine using the

    /// [mark_request_as_sent()](Self::mark_request_as_sent) method.

    ///

    /// **Note**: This method call should be locked per call.

    pub fn outgoing_requests(&self) -> Result<Vec<Request>, CryptoStoreError> {

        Ok(self

            .runtime

            .block_on(self.inner.outgoing_requests())?

            .into_iter()

            .map(|r| r.into())

            .collect())

    }

    /// Mark a request that was sent to the server as sent.

    ///

    /// # Arguments

    ///

    /// * `request_id` - The unique ID of the request that was sent out. This

    ///   needs to be an UUID.

    ///

    /// * `request_type` - The type of the request that was sent out.

    ///

    /// * `response_body` - The body of the response that was received.

    pub fn mark_request_as_sent(

        &self,

        request_id: String,

        request_type: RequestType,

        response_body: String,

    ) -> Result<(), CryptoStoreError> {

        let id: OwnedTransactionId = request_id.into();

        let response = response_from_string(&response_body);

        let response: OwnedResponse = match request_type {

            RequestType::KeysUpload => {

                KeysUploadResponse::try_from_http_response(response).map(Into::into)

            }

            RequestType::KeysQuery => {

                KeysQueryResponse::try_from_http_response(response).map(Into::into)

            }

            RequestType::ToDevice => {

                ToDeviceResponse::try_from_http_response(response).map(Into::into)

            }

            RequestType::KeysClaim => {

                KeysClaimResponse::try_from_http_response(response).map(Into::into)

            }

            RequestType::SignatureUpload => {

                SignatureUploadResponse::try_from_http_response(response).map(Into::into)

            }

            RequestType::KeysBackup => {

                KeysBackupResponse::try_from_http_response(response).map(Into::into)

            }

            RequestType::RoomMessage => {

                RoomMessageResponse::try_from_http_response(response).map(Into::into)

            }

        }

        .expect("Can't convert json string to response");

        self.runtime.block_on(self.inner.mark_request_as_sent(&id, &response))?;

        Ok(())

    }

    /// Let the state machine know about E2EE related sync changes that we

    /// received from the server.

    ///

    /// This needs to be called after every sync, ideally before processing

    /// any other sync changes.

    ///