macltok

mod support;

use std::{

    collections::{BTreeMap, HashSet},

    sync::{

        Arc, Mutex,

        atomic::{AtomicUsize, Ordering},

    },

    time::Duration,

};

#[cfg(feature = "sqlite")]

use macltok::{EncryptedSessionStore, PasswordLogin, Restore, Session, SessionStore, StoreConfig};

use macltok::{Engine, EngineConfig, Event, MentionReason, MentionsPage};

use matrix_sdk::{

    Client,

    ruma::{OwnedRoomId, OwnedUserId, event_id, events::Mentions, room_id, user_id},

    test_utils::mocks::{MatrixMockServer, RoomContextResponseTemplate},

};

use matrix_sdk_test::{ALICE, JoinedRoomBuilder, event_factory::EventFactory};

use ruma::push::{

    Action, EventMatchConditionData, HighlightTweakValue, NewConditionalPushRule,

    NewPatternedPushRule, PushCondition, Ruleset, SoundTweakValue, Tweak,

};

use serde_json::{Value, json};

#[cfg(feature = "sqlite")]

use support::receive_until;

use support::shut_down;

use wiremock::{

    Mock, Request, Respond, ResponseTemplate,

    matchers::{method, path, path_regex, query_param},

};

#[cfg(feature = "serde")]

use macltok::{MentionCoverage, MentionsBoundary, MentionsCommand};

const SLIDING_SYNC_PATH: &str = "/_matrix/client/unstable/org.matrix.simplified_msc3575/sync";

struct SlidingFeed {

    rooms: Arc<Mutex<BTreeMap<OwnedRoomId, Vec<Value>>>>,

    required_state: Arc<Mutex<BTreeMap<OwnedRoomId, Vec<Value>>>>,

    requests: Arc<AtomicUsize>,

}

impl SlidingFeed {

    fn new(room_id: OwnedRoomId, events: Vec<Value>) -> Self {

        Self {

            rooms: Arc::new(Mutex::new(BTreeMap::from([(room_id, events)]))),

            required_state: Arc::new(Mutex::new(BTreeMap::new())),

            requests: Arc::new(AtomicUsize::new(0)),

        }

    }

    fn replace(&self, room_id: &ruma::RoomId, events: Vec<Value>) {

        self.rooms

            .lock()

            .unwrap()

            .insert(room_id.to_owned(), events);

    }

    fn add_room(&self, room_id: &ruma::RoomId, events: Vec<Value>) {

        self.rooms

            .lock()

            .unwrap()

            .insert(room_id.to_owned(), events);

    }

    fn set_required_state(&self, room_id: &ruma::RoomId, events: Vec<Value>) {

        self.required_state

            .lock()

            .unwrap()

            .insert(room_id.to_owned(), events);

    }

    async fn mount(&self, server: &MatrixMockServer) {

        Mock::given(method("POST"))

            .and(path(SLIDING_SYNC_PATH))

            .respond_with(self.clone())

            .mount(server.server())

            .await;

    }

}

impl Clone for SlidingFeed {

    fn clone(&self) -> Self {

        Self {

            rooms: self.rooms.clone(),

            required_state: self.required_state.clone(),

            requests: self.requests.clone(),

        }

    }

}

impl Respond for SlidingFeed {

    fn respond(&self, _request: &Request) -> ResponseTemplate {

        let request = self.requests.fetch_add(1, Ordering::AcqRel);

        let required_state = self.required_state.lock().unwrap().clone();

        let rooms: serde_json::Map<String, Value> = self

            .rooms

            .lock()

            .unwrap()

            .iter()

            .map(|(room_id, events)| {

                (

                    room_id.to_string(),

                    json!({

                        "initial": request == 0,

                        "required_state": required_state

                            .get(room_id)

                            .cloned()

                            .unwrap_or_default(),

                        "timeline": events,

                    }),

                )

            })

            .collect();

        ResponseTemplate::new(200)

            .set_body_json(json!({

                "txn_id": null,

                "pos": format!("mentions-{request}"),

                "lists": {},

                "rooms": rooms,

                "extensions": {},

            }))

            .set_delay(Duration::from_millis(20))

    }

}

async fn joined_client(server: &MatrixMockServer, room_id: &ruma::RoomId) -> (Client, OwnedUserId) {

    let client = server.client_builder().build().await;

    let own_user = client.user_id().unwrap().to_owned();

    let room_id = room_id.to_owned();

    server

        .mock_sync()

        .ok_and_run(&client, {

            let own_user = own_user.clone();

            move |builder| {

                let factory = EventFactory::new().room(&room_id);

                builder.add_global_account_data(

                    factory.push_rules(Ruleset::server_default(&own_user)),

                );

                builder.add_joined_room(

                    JoinedRoomBuilder::new(&room_id).add_state_event(factory.member(&own_user)),

                );

            }

        })

        .await;

    (client, own_user)

}

fn message(

    event_id: &str,

    sender: &str,

    timestamp: u64,

    body: &str,

    mentioned: Option<&str>,

) -> Value {

    let mut content = json!({ "msgtype": "m.text", "body": body });

    if let Some(mentioned) = mentioned {

        content["m.mentions"] = json!({ "user_ids": [mentioned] });

    }

    json!({

        "type": "m.room.message",

        "event_id": event_id,

        "sender": sender,

        "origin_server_ts": timestamp,

        "content": content,

    })

}

fn room_mention(event_id: &str, sender: &str, timestamp: u64) -> Value {

    let mut event = message(event_id, sender, timestamp, "room mention", None);

    event["content"]["m.mentions"] = json!({ "room": true });

    event

}

async fn mentions_page(

    events: &mut macltok::EventStream,

    room_id: &ruma::RoomId,

    predicate: impl Fn(&MentionsPage) -> bool,

) -> MentionsPage {

    mentions_page_with_diagnostics(events, room_id, predicate, None).await

}

async fn mentions_page_with_diagnostics(

    events: &mut macltok::EventStream,

    room_id: &ruma::RoomId,

    predicate: impl Fn(&MentionsPage) -> bool,

    diagnostic_server: Option<&MatrixMockServer>,

) -> MentionsPage {

    let mut last_page = None;

    let result = tokio::time::timeout(Duration::from_secs(5), async {

        while let Some(event) = events.recv().await {

            match event {

                Event::TimelineReset { timeline_id, .. } if timeline_id.room_id() == room_id => {

                    panic!("the unopened room produced a timeline reset");

                }

                Event::MentionsChanged(update) => {

                    let snapshot = update.snapshot;

                    let matched = predicate(&snapshot);

                    last_page = Some(snapshot.clone());

                    if matched {

                        return snapshot;

                    }

                }

                _ => {}

            }

        }

        panic!("event stream closed before the expected mentions page");

    })

    .await;

    match result {

        Ok(page) => page,

        Err(_) => {

            let last_state = last_page.map(|page: MentionsPage| {

                (

                    page.revision,

                    page.unread_count,

                    page.entries

                        .iter()

                        .map(|entry| {

                            (

                                entry.key.room_id.to_string(),

                                entry.key.event_id.to_string(),

                                entry.capture.clone(),

                                entry.revision,

                            )

                        })

                        .collect::<Vec<_>>(),

                )

            });

            let requests = if let Some(server) = diagnostic_server {

                server.server().received_requests().await.map(|requests| {

                    requests

                        .iter()

                        .filter(|request| {

                            let path = request.url.path();

                            path.contains(room_id.as_str())

                                && (path.contains("/context/") || path.ends_with("/messages"))

                        })

                        .map(|request| format!("{} {}", request.method, request.url))

                        .collect::<Vec<_>>()

                })

            } else {

                None

            };

            panic!(

                "expected mentions update did not arrive, last state {last_state:?}, requests {requests:?}"

            );

        }

    }

}

fn entry_ids(page: &MentionsPage) -> HashSet<String> {

    page.entries

        .iter()

        .map(|entry| entry.key.event_id.to_string())

        .collect()

}

async fn wait_for_requests(feed: &SlidingFeed, minimum: usize) {

    tokio::time::timeout(Duration::from_secs(5), async {

        while feed.requests.load(Ordering::Acquire) < minimum {

            tokio::task::yield_now().await;

        }

    })

    .await

    .expect("the sliding sync feed did not receive another request");

}

async fn mount_messages_page(

    server: &MatrixMockServer,

    room_id: &ruma::RoomId,

    direction: &str,

    token: &str,

    limit: u16,

) {

    Mock::given(method("GET"))

        .and(path_regex(format!(

            r"^/_matrix/client/v3/rooms/{room_id}/messages$"

        )))

        .and(query_param("dir", direction))

        .and(query_param("from", token))

        .and(move |request: &Request| {

            request

                .url

                .query_pairs()

                .find(|(key, _)| key == "limit")

                .map_or(limit == 10, |(_, value)| value == limit.to_string())

        })

        .respond_with(ResponseTemplate::new(200).set_body_json(json!({

            "start": token,

            "end": format!("{token}-next"),

            "chunk": [],

            "state": [],

        })))

        .expect(1)

        .mount(server.server())

        .await;

}

#[tokio::test]

async fn unopened_room_mentions_are_deduplicated_and_read_through_uses_a_revision_boundary() {

    let server = MatrixMockServer::new().await;

    let room_id = room_id!("!mentions-inbox:localhost");

    let (client, own_user) = joined_client(&server, room_id).await;

    let sender = user_id!("@alice:localhost");

    let initial = vec![

        message(

            "$mention",

            sender.as_str(),

            2_000,

            "direct mention",

            Some(own_user.as_str()),

        ),

        message(

            "$own-mention",

            own_user.as_str(),

            2_100,

            "own event",

            Some(own_user.as_str()),

        ),

        message(

            "$unrelated",

            sender.as_str(),

            2_200,

            "ordinary message",

            None,

        ),

    ];

    let feed = SlidingFeed::new(room_id.to_owned(), initial.clone());

    feed.mount(&server).await;

    let (engine, handle, mut events) = Engine::from_client(

        client,

        EngineConfig::default()

            .with_initial_history_limit(0)

            .with_search_backfill_pages_per_room(0),

    )

    .await

    .unwrap();

    let running = tokio::spawn(engine.run());

    let first = mentions_page(&mut events, room_id, |page| !page.entries.is_empty()).await;

    assert_eq!(entry_ids(&first), HashSet::from(["$mention".to_owned()]));

    assert_eq!(first.unread_count, 1);

    let first_boundary = first.boundary.clone();

    wait_for_requests(&feed, 3).await;

    let replayed = handle.mentions_page(None, 50).await.unwrap();

    assert_eq!(entry_ids(&replayed), HashSet::from(["$mention".to_owned()]));

    let key = replayed.entries[0].key.clone();

    handle

        .set_mentions_read(vec![key.clone()], true)

        .await

        .unwrap();

    let read = handle.mentions_page(None, 50).await.unwrap();

    assert_eq!(read.unread_count, 0);

    feed.replace(

        room_id,

        [

            initial,

            vec![message(

                "$older-mention",

                sender.as_str(),

                1_000,

                "older direct mention",

                Some(own_user.as_str()),

            )],

        ]

        .concat(),

    );

    let updated = mentions_page(&mut events, room_id, |page| page.entries.len() == 2).await;

    assert_eq!(

        entry_ids(&updated),

        HashSet::from(["$mention".to_owned(), "$older-mention".to_owned(),])

    );

    handle

        .mark_mentions_read_through(first_boundary)

        .await

        .unwrap();

    let after_boundary = handle.mentions_page(None, 50).await.unwrap();

    let read_states: Vec<_> = after_boundary

        .entries

        .iter()

        .map(|entry| (entry.key.event_id.to_string(), entry.read))

        .collect();

    assert!(read_states.contains(&("$mention".to_owned(), true)));

    assert!(read_states.contains(&("$older-mention".to_owned(), false)));

    assert_eq!(after_boundary.unread_count, 1);

    shut_down(&handle, &mut events, running).await;

}

#[tokio::test]

async fn edits_can_remove_and_restore_a_mention_until_redaction_removes_its_target() {

    let server = MatrixMockServer::new().await;

    let room_id = room_id!("!mentions-edits:localhost");

    let (client, own_user) = joined_client(&server, room_id).await;

    let sender = user_id!("@alice:localhost");

    let original = message(

        "$edited-target",

        sender.as_str(),

        1_000,

        "original mention",

        Some(own_user.as_str()),

    );

    let removal = json!({

        "type": "m.room.message",

        "event_id": "$remove-edit",

        "sender": sender,

        "origin_server_ts": 2_000,

        "content": {

            "msgtype": "m.text",

            "body": "* edited without a mention",

            "m.relates_to": { "rel_type": "m.replace", "event_id": "$edited-target" },

            "m.new_content": { "msgtype": "m.text", "body": "edited without a mention" },

        },

    });

    let addition = json!({

        "type": "m.room.message",

        "event_id": "$add-edit",

        "sender": sender,

        "origin_server_ts": 3_000,

        "content": {

            "msgtype": "m.text",

            "body": "* edited to mention you",

            "m.relates_to": { "rel_type": "m.replace", "event_id": "$edited-target" },

            "m.new_content": {

                "msgtype": "m.text",

                "body": "edited to mention you",

                "m.mentions": { "user_ids": [own_user.as_str()] },

            },

        },

    });

    let redaction = json!({

        "type": "m.room.redaction",

        "event_id": "$redaction",

        "sender": sender,

        "origin_server_ts": 4_000,

        "redacts": "$edited-target",

        "content": {},

    });

    let feed = SlidingFeed::new(room_id.to_owned(), vec![original.clone()]);

    feed.mount(&server).await;

    let (engine, handle, mut events) = Engine::from_client(

        client,

        EngineConfig::default()

            .with_initial_history_limit(0)

            .with_search_backfill_pages_per_room(0),

    )

    .await

    .unwrap();

    let running = tokio::spawn(engine.run());

    let discovered = mentions_page(&mut events, room_id, |page| !page.entries.is_empty()).await;

    assert_eq!(

        entry_ids(&discovered),

        HashSet::from(["$edited-target".to_owned()])

    );

    feed.replace(room_id, vec![original.clone(), removal.clone()]);

    let removed = mentions_page(&mut events, room_id, |page| {

        page.revision > discovered.revision && page.entries.is_empty()

    })

    .await;

    assert!(removed.entries.is_empty());

    assert_eq!(

        handle

            .mention_context(room_id, ruma::event_id!("$edited-target"))

            .await

            .unwrap_err()

            .kind,

        macltok::ErrorKind::EventNotFound

    );

    feed.replace(room_id, vec![original.clone(), removal, addition]);

    let restored = mentions_page(&mut events, room_id, |page| {

        page.revision > removed.revision && entry_ids(page).contains("$edited-target")