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 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(|message| {

            message.own

                && message.transaction_id.is_some()

                && message

                    .media

                    .as_ref()

                    .is_some_and(|media| media.filename == "pending.png")

        })

    })

    .await;

    let media = support::event_messages(&event)

        .into_iter()

        .find(|message| message.own && message.media.is_some())

        .and_then(|message| message.media.as_ref())

        .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()

        .await;

    harness

        .server

        .mock_upload()

        .expect_mime_type("image/png")

        .ok(mxc_uri!("mxc://localhost/video-thumb"))

        .expect(1)

        .mount()

        .await;

    harness

        .server

        .mock_upload()

        .expect_mime_type("video/mp4")

        .ok(mxc_uri!("mxc://localhost/video-upload"))

        .expect(1)

        .mount()

        .await;

    harness

        .server

        .mock_room_send()

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

        .expect(1)

        .mount()

        .await;

    let mut upload = AttachmentUpload::new(

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

        "clip.mp4",

        "video/mp4",

        b"video bytes".to_vec(),

    );

    upload.details = Some(UploadDetails::Video {

        width: Some(1920),

        height: Some(1080),

        duration_ms: Some(3500),

        blurhash: Some("blurhash".to_owned()),

    });

    upload.thumbnail = Some(UploadThumbnail::new(

        b"poster bytes".to_vec(),

        "image/png",

        320,

        180,

    ));

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

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

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

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

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

    assert_eq!(content["info"]["duration"], 3500);

    assert_eq!(content["info"]["size"], 11);

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

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

    assert_eq!(content["info"]["thumbnail_info"]["size"], 12);

    assert_eq!(

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

        "mxc://localhost/video-thumb"

    );

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

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

    assert!(requests.iter().any(|request| {

        request.url.path().ends_with("/_matrix/media/v3/upload") && request.body == b"poster bytes"

    }));

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

}

#[tokio::test]

async fn voice_upload_emits_voice_details_and_correlated_progress() {

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

    harness

        .server

        .mock_authenticated_media_config()

        .ok_default()

        .mount()

        .await;

    harness

        .server

        .mock_upload()

        .expect_mime_type("audio/ogg")

        .respond_with(

            ResponseTemplate::new(200)

                .set_body_json(json!({ "content_uri": "mxc://localhost/voice-upload" }))

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

        )

        .expect(1)

        .mount()

        .await;

    harness

        .server

        .mock_room_send()

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

        .expect(1)

        .mount()

        .await;

    let data = b"voice bytes".to_vec();

    let mut upload = AttachmentUpload::new(

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

        "voice.ogg",

        "audio/ogg",

        data.clone(),

    );

    upload.details = Some(UploadDetails::Voice {

        duration_ms: 4200,

        waveform: vec![0, 512, 1024],

    });

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

    let pending = harness

        .handle

        .pending_send_ids(room_id!("!media:localhost"))

        .await

        .unwrap();

    assert_eq!(pending.len(), 1);

    let transaction_id = pending[0].clone();

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

        matches!(event, Event::MediaUploadProgress { transaction_id: seen, complete: true, .. } if seen == &transaction_id)

    })

    .await;

    assert!(matches!(progress, Event::MediaUploadProgress {

        room_id,

        transaction_id: seen,

        uploaded_bytes,

        total_bytes,

        complete: true,

    } if room_id == room_id!("!media:localhost") && seen == transaction_id && uploaded_bytes == data.len() as u64 && total_bytes == data.len() as u64));

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

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

    assert_eq!(content["info"]["duration"], 4200);

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

    assert_eq!(content["org.matrix.msc1767.audio"]["duration"], 4200);

    assert_eq!(

        content["org.matrix.msc1767.audio"]["waveform"],

        json!([0, 512, 1024])

    );

    assert_eq!(content["org.matrix.msc3245.voice"], json!({}));

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

}

#[tokio::test]

async fn mismatched_upload_metadata_and_invalid_waveform_fail_without_uploading() {

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

    let mut mismatched = AttachmentUpload::new(

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

        "clip.mp4",

        "video/mp4",

        b"video bytes".to_vec(),

    );

    mismatched.details = Some(UploadDetails::Image {

        width: Some(32),

        height: Some(32),

        blurhash: None,

    });

    let error = harness.handle.upload(mismatched).await.unwrap_err();

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

    let mut voice = AttachmentUpload::new(

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

        "voice.ogg",

        "audio/ogg",

        b"voice bytes".to_vec(),

    );

    voice.details = Some(UploadDetails::Voice {

        duration_ms: 1000,

        waveform: vec![1025],

    });

    let error = harness.handle.upload(voice).await.unwrap_err();

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

    assert!(

        !harness

            .server

            .received_requests()

            .await

            .unwrap()

            .iter()

            .any(|request| { request.url.path().ends_with("/_matrix/media/v3/upload") })

    );

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

}

#[tokio::test]

async fn upload_over_media_limit_fails_without_uploading() {

    let mut harness = start_with_config(

        plain_image(),

        EngineConfig::default()

            .with_initial_history_limit(0)

            .with_max_media_bytes(4),

    )

    .await;

    let upload = AttachmentUpload::new(

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

        "large.bin",

        "application/octet-stream",

        b"too large".to_vec(),

    );

    let error = harness.handle.upload(upload).await.unwrap_err();

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

    assert!(

        !harness

            .server

            .received_requests()

            .await

            .unwrap()

            .iter()

            .any(|request| { request.url.path().ends_with("/_matrix/media/v3/upload") })

    );

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

}

#[tokio::test]

async fn audio_thumbnail_fails_without_uploading() {

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

    let mut upload = AttachmentUpload::new(

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

        "voice.ogg",

        "audio/ogg",

        b"voice bytes".to_vec(),

    );

    upload.thumbnail = Some(UploadThumbnail::new(

        b"thumbnail".to_vec(),

        "image/png",

        80,

        80,

    ));

    let error = harness.handle.upload(upload).await.unwrap_err();

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

    assert!(

        !harness

            .server

            .received_requests()

            .await

            .unwrap()

            .iter()

            .any(|request| { request.url.path().ends_with("/_matrix/media/v3/upload") })

    );

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

}

#[tokio::test]

async fn invalid_upload_mime_fails_without_uploading() {

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

    let upload = AttachmentUpload::new(

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

        "photo.jpg",

        "not a MIME type",

        b"photo data".to_vec(),

    );

    let error = harness.handle.upload(upload).await.unwrap_err();

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

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

}

async fn start(content: RoomMessageEventContent) -> MediaHarness {

    start_with_config(

        content,

        EngineConfig::default().with_initial_history_limit(0),

    )

    .await

}

async fn start_with_config(content: RoomMessageEventContent, config: EngineConfig) -> MediaHarness {

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

    server

        .mock_sliding_sync()

        .ok(v5::Response::new("1".to_owned()))

        .mount()

        .await;

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

    let room_id = room_id!("!media:localhost");

    server.mock_room_state_encryption().plain().mount().await;

    server.sync_joined_room(&client, room_id).await;

    let sync_client = client.clone();

    let (engine, handle, mut events) = Engine::from_client(client, config).await.unwrap();

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

    handle.open_room(room_id.to_owned()).await.unwrap();

    support::receive_until(

        &mut events,

        |event| matches!(event, Event::TimelineReset { timeline_id: reset, .. } if reset.room_id() == room_id),

    )

    .await;

    server

        .sync_room(

            &sync_client,

            JoinedRoomBuilder::new(room_id)

                .add_timeline_event(EventFactory::new().event(content).sender(&ALICE)),

        )

        .await;

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

        support::event_messages(event)

            .into_iter()

            .any(|message| message.media.is_some())

    })

    .await;

    let media_id = support::event_messages(&event)

        .into_iter()

        .find_map(|message| message.media.as_ref().map(|media| media.id.clone()))

        .unwrap();

    MediaHarness {

        server,

        client: sync_client,

        handle,

        events,

        running,

        media_id,

    }

}

async fn mount_downloads(server: &MatrixMockServer) {

    server.mock_media_download().ok_image().mount().await;

    server.mock_authed_media_download().ok_image().mount().await;

}

fn thumbnail_size() -> MediaSize {

    MediaSize::Thumbnail {

        width: 800,

        height: 800,

    }

}

async fn fetch(harness: &MediaHarness, size: MediaSize) -> Arc<[u8]> {

    harness

        .handle

        .media(&harness.media_id, size)

        .await

        .unwrap()

        .bytes

}

async fn media_paths(server: &MatrixMockServer) -> Vec<String> {

    server

        .received_requests()

        .await

        .unwrap()

        .into_iter()

        .map(|request| request.url.path().to_owned())

        .filter(|path| path.contains("/media/"))

        .collect()

}

async fn uploaded_content(server: &MatrixMockServer) -> Value {

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

        loop {

            if let Some(request) = server

                .received_requests()

                .await

                .unwrap()

                .iter()

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

            {

                return serde_json::from_slice(&request.body).unwrap();

            }

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

        }

    })

    .await

    .expect("queued media was not sent")

}

fn file_with_thumbnail() -> RoomMessageEventContent {

    let mut info = FileInfo::new();

    info.thumbnail_source = Some(MediaSource::Plain(

        mxc_uri!("mxc://localhost/explicit-thumbnail").to_owned(),

    ));

    RoomMessageEventContent::new(MessageType::File(

        FileMessageEventContent::plain(

            "file.bin".to_owned(),

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

        )

        .info(Box::new(info)),

    ))

}

fn plain_image() -> RoomMessageEventContent {

    RoomMessageEventContent::new(MessageType::Image(ImageMessageEventContent::plain(

        "image.jpg".to_owned(),

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

    )))

}

fn plain_audio() -> RoomMessageEventContent {

    RoomMessageEventContent::new(MessageType::Audio(AudioMessageEventContent::plain(

        "audio.ogg".to_owned(),

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

    )))

}