macltok

mod support;

use std::{sync::Arc, time::Duration};

use macltok::{

    AttachmentUpload, ClientHandle, Engine, EngineConfig, EngineError, ErrorKind, Event,

    EventStream, MediaId, MediaSize, UploadDetails, UploadThumbnail,

};

use matrix_sdk::{

    media::{MediaFormat, MediaRequestParameters},

    ruma::{

        api::client::sync::sync_events::v5,

        event_id,

        events::room::{

            EncryptedFile, MediaSource,

            message::{

                AudioMessageEventContent, FileInfo, FileMessageEventContent,

                ImageMessageEventContent, MessageType, RoomMessageEventContent,

            },

        },

        media::Method,

        mxc_uri, room_id,

    },

    test_utils::mocks::MatrixMockServer,

};

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

use serde_json::{Value, json};

use wiremock::{

    Mock, ResponseTemplate,

    matchers::{method, path_regex},

};

struct MediaHarness {

    server: MatrixMockServer,

    client: matrix_sdk::Client,

    handle: ClientHandle,

    events: EventStream,

    running: tokio::task::JoinHandle<Result<(), EngineError>>,

    media_id: MediaId,

}

#[tokio::test]

async fn full_size_uses_full_source() {

    let mut harness = start(file_with_thumbnail()).await;

    mount_downloads(&harness.server).await;

    let bytes = fetch(&harness, MediaSize::Full).await;

    assert_eq!(&*bytes, b"binaryjpegfullimagedata");

    assert!(

        media_paths(&harness.server)

            .await

            .iter()

            .any(|path| path.ends_with("/localhost/full"))

    );

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

}

#[tokio::test]

async fn encrypted_media_is_cached_as_plaintext_like_the_sdk_does() {

    use std::io::Read as _;

    let plaintext = b"very secret image".to_vec();

    let mut cursor = std::io::Cursor::new(plaintext.clone());

    let mut encryptor = matrix_sdk_crypto::AttachmentEncryptor::new(&mut cursor);

    let mut ciphertext = Vec::new();

    encryptor.read_to_end(&mut ciphertext).unwrap();

    let keys = encryptor.finish();

    let file = EncryptedFile::new(

        mxc_uri!("mxc://localhost/secret").to_owned(),

        keys.encryption_info,

        keys.hashes,

    );

    let source = MediaSource::Encrypted(Box::new(file.clone()));

    let mut harness = start(RoomMessageEventContent::new(MessageType::Image(

        ImageMessageEventContent::encrypted("secret.png".to_owned(), file),

    )))

    .await;

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

        .and(path_regex("/download/localhost/secret$"))

        .respond_with(ResponseTemplate::new(200).set_body_bytes(ciphertext))

        .mount(harness.server.server())

        .await;

    let from_sdk = harness

        .client

        .media()

        .get_media_content(

            &MediaRequestParameters {

                source,

                format: MediaFormat::File,

            },

            true,

        )

        .await

        .unwrap();

    assert_eq!(from_sdk, plaintext);

    assert_eq!(&*fetch(&harness, MediaSize::Full).await, &plaintext[..]);

    assert_eq!(media_paths(&harness.server).await.len(), 1);

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

}

#[tokio::test]

async fn a_plain_event_cannot_poison_an_encrypted_files_cache() {

    use std::io::Read as _;

    let plaintext = b"very secret image".to_vec();

    let mut cursor = std::io::Cursor::new(plaintext.clone());

    let mut encryptor = matrix_sdk_crypto::AttachmentEncryptor::new(&mut cursor);

    let mut ciphertext = Vec::new();

    encryptor.read_to_end(&mut ciphertext).unwrap();

    let keys = encryptor.finish();

    let uri = mxc_uri!("mxc://localhost/secret").to_owned();

    let file = EncryptedFile::new(uri.clone(), keys.encryption_info, keys.hashes);

    let mut harness = start(RoomMessageEventContent::new(MessageType::Image(

        ImageMessageEventContent::plain("bait.png".to_owned(), uri),

    )))

    .await;

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

        .and(path_regex("/download/localhost/secret$"))

        .respond_with(ResponseTemplate::new(200).set_body_bytes(ciphertext.clone()))

        .mount(harness.server.server())

        .await;

    assert_eq!(&*fetch(&harness, MediaSize::Full).await, &ciphertext[..]);

    harness

        .server

        .sync_room(

            &harness.client,

            JoinedRoomBuilder::new(room_id!("!media:localhost")).add_timeline_event(

                EventFactory::new()

                    .event(RoomMessageEventContent::new(MessageType::Image(

                        ImageMessageEventContent::encrypted("secret.png".to_owned(), file),

                    )))

                    .sender(&ALICE),

            ),

        )

        .await;

    let event = support::receive_until(&mut harness.events, |event| {

        support::event_messages(event).into_iter().any(|message| {

            support::media_of(message).is_some_and(|media| media.filename == "secret.png")

        })

    })

    .await;

    let media_id = support::event_messages(&event)

        .into_iter()

        .find_map(|message| {

            support::media_of(message).filter(|media| media.filename == "secret.png")

        })

        .map(|media| media.id.clone())

        .unwrap();

    let media = harness

        .handle

        .media(&media_id, MediaSize::Full)

        .await

        .unwrap();

    assert_eq!(&*media.bytes, &plaintext[..]);

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

}

#[tokio::test]

async fn simultaneous_fetches_share_one_download() {

    let mut harness = start(plain_image()).await;

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

        .and(path_regex("/download/localhost/image$"))

        .respond_with(

            ResponseTemplate::new(200)

                .set_body_bytes(b"shared image bytes".to_vec())

                .set_delay(Duration::from_millis(300)),

        )

        .mount(harness.server.server())

        .await;

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

        tokio::join!(

            harness.handle.media(&harness.media_id, MediaSize::Full),

            harness.handle.media(&harness.media_id, MediaSize::Full),

        )

    })

    .await

    .expect("both fetches did not complete");

    assert_eq!(&*first.unwrap().bytes, b"shared image bytes");

    assert_eq!(&*second.unwrap().bytes, b"shared image bytes");

    let downloads = harness

        .server

        .received_requests()

        .await

        .unwrap()

        .iter()

        .filter(|request| request.url.path().ends_with("/download/localhost/image"))

        .count();

    assert_eq!(downloads, 1, "duplicate in-flight media download");

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

}

#[tokio::test]

async fn downloaded_media_over_limit_is_rejected() {

    let mut harness = start_with_config(

        plain_image(),

        EngineConfig::default()

            .with_initial_history_limit(0)

            .with_max_media_bytes(4),

    )

    .await;

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

        .and(path_regex("/download/localhost/image$"))

        .respond_with(ResponseTemplate::new(200).set_body_bytes(b"too large".to_vec()))

        .mount(harness.server.server())

        .await;

    let error = harness

        .handle

        .media(&harness.media_id, MediaSize::Full)

        .await

        .unwrap_err();

    assert_eq!(error.kind, ErrorKind::MediaTooLarge);

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

}

#[tokio::test]

async fn thumbnail_size_prefers_explicit_thumbnail() {

    let mut harness = start(file_with_thumbnail()).await;

    mount_downloads(&harness.server).await;

    let bytes = fetch(&harness, thumbnail_size()).await;

    assert_eq!(&*bytes, b"binaryjpegfullimagedata");

    let paths = media_paths(&harness.server).await;

    assert!(

        paths

            .iter()

            .any(|path| path.ends_with("/localhost/explicit-thumbnail"))

    );

    assert!(!paths.iter().any(|path| path.ends_with("/localhost/full")));

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

}

#[tokio::test]

async fn plain_image_thumbnail_uses_thumbnail_endpoint() {

    let mut harness = start(plain_image()).await;

    harness

        .server

        .mock_media_thumbnail(Method::Scale, 800, 800, false)

        .ok()

        .mount()

        .await;

    harness

        .server

        .mock_authed_media_thumbnail(Method::Scale, 800, 800, false)

        .ok()

        .mount()

        .await;

    let bytes = fetch(&harness, thumbnail_size()).await;

    assert_eq!(&*bytes, b"binaryjpegthumbnaildata");

    assert!(

        media_paths(&harness.server)

            .await

            .iter()

            .any(|path| path.contains("/thumbnail/localhost/image"))

    );

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

}

#[tokio::test]

async fn audio_thumbnail_fetch_fails_with_media_unavailable() {

    let mut harness = start(plain_audio()).await;

    let error = harness

        .handle

        .media(&harness.media_id, thumbnail_size())

        .await

        .unwrap_err();

    assert_eq!(error.kind, ErrorKind::MediaUnavailable);

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

}

#[tokio::test]

async fn media_upload_is_queued_and_sent() {

    let mut harness = start(plain_image()).await;

    harness

        .server

        .mock_authenticated_media_config()

        .ok_default()

        .mount()

        .await;

    harness

        .server

        .mock_upload()

        .expect_mime_type("image/jpeg")

        .ok(mxc_uri!("mxc://localhost/uploaded"))

        .expect(1)

        .mount()

        .await;

    harness

        .server

        .mock_room_send()

        .ok(event_id!("$uploaded:localhost"))

        .expect(1)

        .mount()

        .await;

    let mut upload = AttachmentUpload::new(

        room_id!("!media:localhost").to_owned(),

        "photo.jpg",

        "image/jpeg",

        b"photo data".to_vec(),

    );

    upload.caption = Some("holiday".to_owned());

    harness.handle.upload(upload).await.unwrap();

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

        loop {

            let requests = harness.server.received_requests().await.unwrap();

            if requests.iter().any(|request| {

                request.url.path().contains("/send/m.room.message/")

                    && String::from_utf8_lossy(&request.body).contains("mxc://localhost/uploaded")

            }) {

                break;

            }

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

        }

    })

    .await

    .expect("queued media was not sent");

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

}

#[tokio::test]

async fn queued_offline_attachment_exposes_fetchable_local_media() {

    let mut harness = start(plain_image()).await;

    harness.client.send_queue().set_enabled(false).await;

    let original = b"offline image bytes".to_vec();

    let mut upload = AttachmentUpload::new(

        room_id!("!media:localhost").to_owned(),

        "pending.png",

        "image/png",

        original.clone(),

    );

    upload.details = Some(UploadDetails::Image {

        width: Some(32),

        height: Some(16),

        blurhash: None,

    });

    harness.handle.upload(upload).await.unwrap();

    let event = support::receive_until(&mut harness.events, |event| {

        support::event_messages(event).into_iter().any(|event| {

            event.own

                && event.transaction_id.is_some()

                && support::media_of(event).is_some_and(|media| media.filename == "pending.png")

        })

    })

    .await;

    let media = support::event_messages(&event)

        .into_iter()

        .find(|event| event.own && support::media_of(event).is_some())

        .and_then(support::media_of)

        .unwrap();

    assert_eq!((media.width, media.height), (Some(32), Some(16)));

    assert_eq!(media.size_bytes, Some(original.len() as u64));

    let media_id = media.id.clone();

    let fetched = harness

        .handle

        .media(&media_id, MediaSize::Full)

        .await

        .unwrap();

    assert_eq!(&*fetched.bytes, &original[..]);

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

}

#[tokio::test]

async fn image_upload_uses_caller_supplied_metadata_and_thumbnail() {

    let mut harness = start(plain_image()).await;

    harness

        .server

        .mock_authenticated_media_config()

        .ok_default()

        .mount()

        .await;

    harness

        .server

        .mock_upload()

        .expect_mime_type("image/png")

        .respond_with(|request: &wiremock::Request| {

            let uri = if request.body == b"thumbnail bytes" {

                "mxc://localhost/image-thumb"

            } else {

                "mxc://localhost/image-upload"

            };

            ResponseTemplate::new(200).set_body_json(json!({ "content_uri": uri }))

        })

        .expect(2)

        .mount()

        .await;

    harness

        .server

        .mock_room_send()

        .ok(event_id!("$image-upload:localhost"))

        .expect(1)

        .mount()

        .await;

    let original = b"image bytes".to_vec();

    let thumbnail_bytes = b"thumbnail bytes".to_vec();

    let mut upload = AttachmentUpload::new(

        room_id!("!media:localhost").to_owned(),

        "large.png",

        "image/png",

        original.clone(),

    );

    upload.details = Some(UploadDetails::Image {

        width: Some(900),

        height: Some(450),

        blurhash: None,

    });

    upload.thumbnail = Some(UploadThumbnail::new(

        thumbnail_bytes.clone(),

        "image/png",

        800,

        400,

    ));

    harness.handle.upload(upload).await.unwrap();

    let content = uploaded_content(&harness.server).await;

    assert_eq!(content["msgtype"], "m.image");

    assert_eq!(content["info"]["w"], 900);

    assert_eq!(content["info"]["h"], 450);

    assert_eq!(content["info"]["size"], original.len());

    assert_eq!(content["info"]["thumbnail_info"]["w"], 800);

    assert_eq!(content["info"]["thumbnail_info"]["h"], 400);

    assert_eq!(content["info"]["thumbnail_info"]["mimetype"], "image/png");

    assert_eq!(

        content["info"]["thumbnail_info"]["size"],

        thumbnail_bytes.len()

    );

    assert_eq!(

        content["info"]["thumbnail_url"],

        "mxc://localhost/image-thumb"

    );

    assert_eq!(content["url"], "mxc://localhost/image-upload");

    let requests = harness.server.received_requests().await.unwrap();

    let uploads = requests

        .iter()

        .filter(|request| request.url.path().ends_with("/_matrix/media/v3/upload"))

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

    assert_eq!(uploads.len(), 2);

    assert!(uploads.iter().any(|request| request.body == original));

    assert!(

        uploads

            .iter()

            .any(|request| request.body == thumbnail_bytes)

    );

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

}

#[tokio::test]

async fn video_upload_preserves_metadata_and_supplied_thumbnail() {

    let mut harness = start(plain_image()).await;

    harness

        .server

        .mock_authenticated_media_config()

        .ok_default()

        .mount()