macltok

mod feed;

#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]

mod index;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

pub(crate) mod store;

mod text;

use std::{

    collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque},

    sync::Arc,

    time::Duration,

};

use futures_util::{

    StreamExt,

    future::Abortable,

    stream::{AbortHandle, FuturesUnordered, SelectAll, abortable},

};

use matrix_sdk::{

    Client, Room, RoomState, SendOutsideWasm,

    deserialized_responses::TimelineEvent as SdkTimelineEvent,

    event_cache::{RoomEventCacheSubscriber, RoomEventCacheUpdate},

    room::{Messages, MessagesOptions},

};

use matrix_sdk_base::RoomInfoNotableUpdateReasons;

use matrix_sdk_ui::eyeball_im::VectorDiff;

use ruma::{

    OwnedRoomId, OwnedUserId, UInt,

    api::client::filter::RoomEventFilter,

    events::{

        ignored_user_list::IgnoredUserListEventContent, relation::RelationType,

        room::encrypted::OriginalSyncRoomEncryptedEvent,

    },

    serde::Raw,

};

use serde_json::Value;

use sha2::{Digest, Sha256};

use tokio::sync::{Mutex, broadcast, mpsc, oneshot};

use self::{

    feed::{Effect, Observed, Redactions},

    index::{Changes, Doc, Hit, KeywordIndex},

};

use crate::{

    CommandOutput, EngineConfig, Error, ErrorKind, Event, RoomSearchCoverage, SearchBackfill,

    SearchCloseReason, SearchCommand, SearchCoverage, SearchHit, SearchInMemory, SearchMatch,

    SearchPage, SearchScope, SearchSession,

    engine::{MediaService, SearchStore, context, executor::EngineBoxFuture},

    normalise::{MediaRegistration, MediaScope},

    transport::EventSender,

};

const MAX_LIMIT: u16 = 50;

/// a consumer that never closes its searches still holds at most this many

const MAX_SEARCHES: usize = 8;

/// hits read back at once while a page is put together

const PREVIEWS_AT_ONCE: usize = 8;

const DECRYPTS_AT_ONCE: usize = 8;

const COMMIT_AFTER_MS: u64 = 5_000;

const COMMIT_BUFFERED: usize = 2_000;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

const COMMIT_RETRY_MS: u64 = 60_000;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

const COMMIT_ATTEMPTS: u32 = 5;

const IDLE_SUBSCRIPTION_MS: u64 = 5 * 60_000;

const BACKFILL_RETRY_MS: u64 = 60_000;

const BACKFILL_RETRY_MAX_MS: u64 = 15 * 60_000;

#[cfg(any(feature = "sqlite", feature = "indexeddb"))]

const SCHEMA: u64 = 2;

/// waiting undecryptable events are saved in this many blobs, so one arriving rewrites a sixteenth

const UTD_SHARDS: u8 = 16;

const BACKFILL_TYPES: [&str; 6] = [

    "m.room.message",

    "m.room.encrypted",

    "m.room.redaction",

    "m.sticker",

    "m.poll.start",

    "org.matrix.msc3381.poll.start",

];

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

type BoxStream<T> = futures_util::stream::BoxStream<'static, T>;

#[cfg(target_family = "wasm")]

type BoxStream<T> = futures_util::stream::LocalBoxStream<'static, T>;

/// a local index of message text. keyword search over tantivy is the only one today, an

/// embedding index can stand beside it later without the public api noticing

#[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]

pub(crate) trait MessageIndex: SendOutsideWasm {

    fn get(&self, key: &str) -> Option<Doc>;

    fn by_content(&self, content: &str) -> Option<Doc>;

    /// replaces whatever `doc.key` held

    fn put(&mut self, doc: Doc);

    fn remove(&mut self, key: &str) -> bool;

    /// every provisional doc waiting on `target`, whoever sent its edit

    fn remove_provisional(&mut self, target: &str) -> bool;

    fn buffered(&self) -> usize;

    /// newest first, strictly older than `before`. also hands back the analysed query terms

    fn search(

        &mut self,

        query: &str,

        rooms: Option<&[String]>,

        before: Option<(u64, u64)>,

        limit: usize,

    ) -> Result<(Vec<Hit>, Vec<String>), Error>;

    fn maintain(&mut self) -> Result<(), Error>;

    fn changes(&mut self) -> Result<Changes, Error>;

    fn committed(&mut self, changes: &Changes);

    /// stops tracking what a commit would need, once nothing will be committed again

    fn detach(&mut self);

}

#[derive(Clone, Debug, Default)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

struct RoomProgress {

    oldest: Option<u64>,

    newest: Option<u64>,

    /// where the next backfill page starts, `None` before the first one

    tail: Option<String>,

    started: bool,

    start_reached: bool,

}

#[derive(Clone, Debug)]

#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]

struct Pending {

    room: String,

    session_id: String,

    ts: u64,

    raw: String,

}

#[derive(Clone, Copy)]

struct Budget {

    pages_per_room: u16,

    page_size: u16,

    interval_ms: u64,

    max_pending: usize,

}

impl Budget {

    fn from_config(config: &EngineConfig) -> Self {

        Self {

            pages_per_room: config.search_backfill_pages_per_room,

            page_size: config.search_backfill_page_size.max(1),

            interval_ms: config.search_backfill_interval_ms,

            max_pending: config.max_search_undecrypted,

        }

    }

}

struct RawPage {

    hits: Vec<Hit>,

    terms: Vec<String>,

    more: bool,

    coverage: SearchCoverage,

}

enum Request {

    Search {

        query: String,

        rooms: Option<Vec<String>>,

        before: Option<(u64, u64)>,

        limit: usize,

        coverage: Vec<OwnedRoomId>,

        reply: oneshot::Sender<Result<RawPage, Error>>,

    },

    Coverage(Vec<OwnedRoomId>, oneshot::Sender<SearchCoverage>),

    /// a hit that turned out redacted when its preview was read

    Forget(String),

    Stop(oneshot::Sender<()>),

}

#[derive(Clone)]

pub(crate) struct Handle {

    requests: mpsc::Sender<Request>,

    client: Client,

    media: Arc<MediaService>,

    sessions: Arc<Sessions>,

}

enum Work {

    Backfill {

        room: OwnedRoomId,

        head: bool,

        from: Option<String>,

        result: Result<Messages, Error>,

    },

    /// the room was left while its page was in flight

    BackfillCancelled,

    Decrypted {

        key: String,

        room: OwnedRoomId,

        value: Option<Value>,

    },

    /// the original first, then whatever edits of it the cache or server still has

    Refetched {

        room: OwnedRoomId,

        values: Vec<Value>,

    },

}

enum RoomUpdate {

    Events(Vec<SdkTimelineEvent>),

    Lagged,

    Closed,

}

struct Failure {

    message: String,

    retry_at_ms: u64,

    delay_ms: u64,

}

enum Persistence {

    Store(SearchStore),

    Memory(SearchInMemory),

}

pub(crate) struct Service {

    client: Client,

    index: Box<dyn MessageIndex>,

    persistence: Persistence,

    #[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]

    commit_failures: u32,

    /// storage keys a load found nothing pointing at, dropped with the next commit

    #[cfg_attr(not(any(feature = "sqlite", feature = "indexeddb")), allow(dead_code))]

    orphans: Vec<String>,

    rooms: BTreeMap<String, RoomProgress>,

    /// each room's newest indexed message when this engine started, what a head pass looks for

    baselines: HashMap<String, u64>,

    pending: BTreeMap<String, Pending>,

    dirty_shards: BTreeSet<u8>,

    decrypt_queue: VecDeque<String>,

    decrypting: usize,

    redactions: Redactions,

    dirty_since: Option<u64>,

    commit_after: u64,

    budget: Budget,

    requests: mpsc::Receiver<Request>,

    subscriptions: HashMap<OwnedRoomId, (AbortHandle, u64)>,

    updates: SelectAll<BoxStream<(OwnedRoomId, RoomUpdate)>>,

    work: FuturesUnordered<EngineBoxFuture<'static, Work>>,

    pages: HashMap<OwnedRoomId, u16>,

    heads: VecDeque<OwnedRoomId>,

    headed: HashSet<OwnedRoomId>,

    head_tokens: HashMap<OwnedRoomId, String>,

    backfilling: Option<(OwnedRoomId, AbortHandle)>,

    next_backfill: u64,

    failures: HashMap<OwnedRoomId, Failure>,

    turn: usize,

}

fn now() -> u64 {

    ruma::MilliSecondsSinceUnixEpoch::now().get().into()

}

fn value_of(event: &SdkTimelineEvent) -> Option<Value> {

    serde_json::from_str(event.raw().json().get()).ok()

}

fn encode_cursor((ts, tie): (u64, u64)) -> String {

    format!("{ts:x}.{tie:x}")

}

fn decode_cursor(cursor: &str) -> Option<(u64, u64)> {

    let (ts, tie) = cursor.split_once('.')?;

    Some((

        u64::from_str_radix(ts, 16).ok()?,

        u64::from_str_radix(tie, 16).ok()?,

    ))

}

fn shard(key: &str) -> u8 {

    Sha256::digest(key.as_bytes())[0] % UTD_SHARDS

}

/// a 400 to a page with a saved token means the server no longer takes the token

fn token_rejected(error: &Error) -> bool {

    error

        .server

        .as_ref()

        .is_some_and(|server| server.status == 400 || server.errcode == "M_INVALID_PARAM")

}

impl Handle {

    pub(crate) async fn execute(&self, command: SearchCommand) -> Result<CommandOutput, Error> {

        match command {

            SearchCommand::Search {

                query,

                scope,

                limit,

            } => {

                let media_scope = self.media.context_scope();

                let prefix = media_prefix(&media_scope);

                let (page, media) = self.search(&query, &scope, None, limit, &prefix).await?;

                let search_id = self

                    .sessions

                    .open(&self.media, query, scope, media_scope, media)

                    .await;

                Ok(CommandOutput::SearchSession(SearchSession {

                    search_id,

                    page,

                }))

            }

            SearchCommand::More {

                search_id,

                cursor,

                limit,

            } => {

                let (query, scope, prefix) = self.sessions.get(&search_id).await?;

                let (page, media) = self

                    .search(&query, &scope, Some(cursor.as_str()), limit, &prefix)

                    .await?;

                self.sessions

                    .register(&self.media, &search_id, media)

                    .await?;

                Ok(CommandOutput::SearchPage(page))

            }

            SearchCommand::Close { search_id } => {

                self.sessions.close(&self.media, &search_id).await?;

                Ok(CommandOutput::Done)

            }

            SearchCommand::Coverage(scope) => {

                let (reply, response) = oneshot::channel();

                self.requests

                    .send(Request::Coverage(self.scope_rooms(&scope), reply))

                    .await

                    .map_err(|_| Error::interrupted())?;

                Ok(CommandOutput::SearchCoverage(

                    response.await.map_err(|_| Error::interrupted())?,

                ))

            }

        }

    }

    /// the media service is cleared beside this, so nothing is released one by one

    pub(crate) async fn stop(&self) {

        let (reply, response) = oneshot::channel();

        if self.requests.send(Request::Stop(reply)).await.is_ok() {

            let _ = response.await;

        }

        self.sessions.open.lock().await.clear();

    }

    fn scope_rooms(&self, scope: &SearchScope) -> Vec<OwnedRoomId> {

        match scope {

            SearchScope::Room(room_id) => vec![room_id.clone()],

            SearchScope::Rooms(room_ids) => room_ids.clone(),

            SearchScope::All => self

                .client

                .joined_rooms()

                .into_iter()

                .map(|room| room.room_id().to_owned())

                .collect(),

        }

    }

    async fn search(

        &self,

        query: &str,

        scope: &SearchScope,

        cursor: Option<&str>,

        limit: u16,

        media_prefix: &str,

    ) -> Result<(SearchPage, Vec<MediaRegistration>), Error> {

        if limit == 0 || limit > MAX_LIMIT {

            return Err(Error::new(

                ErrorKind::Invalid,

                "search limit must be between 1 and 50",

            ));

        }

        let before = match cursor {

            Some(cursor) => Some(

                decode_cursor(cursor)

                    .ok_or_else(|| Error::new(ErrorKind::Invalid, "search cursor is invalid"))?,

            ),

            None => None,

        };

        let rooms = self.scope_rooms(scope);

        let filter = Some(rooms.iter().map(ToString::to_string).collect());

        let (reply, response) = oneshot::channel();

        self.requests

            .send(Request::Search {

                query: query.to_owned(),

                rooms: filter,

                before,

                limit: usize::from(limit),

                coverage: rooms,

                reply,

            })

            .await

            .map_err(|_| Error::interrupted())?;

        let raw = response.await.map_err(|_| Error::interrupted())??;

        let next_cursor = raw

            .more

            .then(|| raw.hits.last().map(|hit| encode_cursor((hit.ts, hit.tie))))

            .flatten();

        let ignored = self.ignored_users().await;

        let current_user_id = self.client.user_id().map(ToOwned::to_owned);

        let terms = &raw.terms;

        let previews: Vec<_> = futures_util::stream::iter(

            raw.hits

                .into_iter()

                .filter(|hit| !ignored.iter().any(|user| user.as_str() == hit.sender)),

        )

        .map(|hit| self.preview(hit, current_user_id.clone(), terms, media_prefix))

        .buffered(PREVIEWS_AT_ONCE)

        .collect()

        .await;

        let mut hits = Vec::new();

        let mut media = Vec::new();

        for preview in previews {

            match preview {

                Preview::Shown(hit, registrations) => {

                    hits.push(hit);

                    media.extend(registrations);

                }

                Preview::Redacted(key) => {

                    let _ = self.requests.send(Request::Forget(key)).await;

                }

                Preview::Unavailable => {}

            }

        }

        let page = SearchPage {

            hits,

            next_cursor,

            coverage: raw.coverage,

        };

        Ok((page, media))

    }

    async fn ignored_users(&self) -> Vec<OwnedUserId> {

        self.client

            .account()

            .account_data::<IgnoredUserListEventContent>()

            .await

            .ok()

            .flatten()

            .and_then(|raw| raw.deserialize().ok())

            .map(|content| content.ignored_users.into_keys().collect())

            .unwrap_or_default()

    }

    async fn preview(

        &self,

        hit: Hit,

        current_user_id: Option<OwnedUserId>,

        terms: &[String],

        media_prefix: &str,

    ) -> Preview {

        let (Ok(room_id), Ok(event_id)) = (

            OwnedRoomId::try_from(hit.room.as_str()),

            ruma::OwnedEventId::try_from(hit.key.as_str()),

        ) else {

            return Preview::Unavailable;

        };

        let Some(room) = self.client.get_room(&room_id) else {

            return Preview::Unavailable;

        };

        let Some((mut original, highlight)) = load(&room, &hit.key).await else {

            return Preview::Unavailable;

        };

        if original.pointer("/unsigned/redacted_because").is_some() {

            return Preview::Redacted(hit.key);

        }

        // the index already settled which edit is newest, and a bundled one can be older

        if let Some(relations) = original

            .pointer_mut("/unsigned/m.relations")

            .and_then(Value::as_object_mut)

        {

            relations.remove("m.replace");

        }

        let mut values = vec![original];

        if hit.content != hit.key