grapevine/src/service/media.rs

use std::io::Cursor;

use image::imageops::FilterType;
use ruma::{
    api::client::error::ErrorKind, http_headers::ContentDisposition,
    OwnedMxcUri,
};
use tokio::{
    fs::{self, File},
    io::{AsyncReadExt, AsyncWriteExt},
};
use tracing::{debug, warn};

use crate::{services, utils, Error, Result};

mod data;

pub(crate) use data::Data;

#[derive(Debug, Eq, PartialEq)]
pub(crate) struct FileMeta {
    // This gets written to the database but we no longer read it
    //
    // TODO: Write a database migration to get rid of this and instead store
    // only the filename instead of the entire `Content-Disposition` header.
    #[allow(dead_code)]
    pub(crate) content_disposition: Option<String>,
    pub(crate) content_type: Option<String>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub(crate) struct MediaFileKey(Vec<u8>);

impl MediaFileKey {
    pub(crate) fn new(key: Vec<u8>) -> Self {
        Self(key)
    }

    pub(crate) fn as_bytes(&self) -> &[u8] {
        &self.0
    }
}

pub(crate) struct Service {
    pub(crate) db: &'static dyn Data,
}

impl Service {
    /// Uploads a file.
    #[tracing::instrument(skip(self, file))]
    pub(crate) async fn create(
        &self,
        mxc: OwnedMxcUri,
        content_disposition: Option<&ContentDisposition>,
        content_type: Option<String>,
        file: &[u8],
    ) -> Result<FileMeta> {
        let meta = FileMeta {
            content_disposition: content_disposition
                .map(ContentDisposition::to_string),
            content_type,
        };
        // Width, Height = 0 if it's not a thumbnail
        let key = self.db.create_file_metadata(mxc, 0, 0, &meta)?;

        self.write_content(&key, file).await?;
        Ok(meta)
    }

    /// Uploads or replaces a file thumbnail.
    #[tracing::instrument(skip(self, file))]
    pub(crate) async fn upload_thumbnail(
        &self,
        mxc: OwnedMxcUri,
        content_disposition: Option<String>,
        content_type: Option<String>,
        width: u32,
        height: u32,
        file: &[u8],
    ) -> Result<FileMeta> {
        let meta = FileMeta {
            content_disposition,
            content_type,
        };
        let key = self.db.create_file_metadata(mxc, width, height, &meta)?;

        self.write_content(&key, file).await?;
        Ok(meta)
    }

    /// Downloads a file.
    #[tracing::instrument(skip(self))]
    pub(crate) async fn get(
        &self,
        mxc: OwnedMxcUri,
    ) -> Result<Option<(FileMeta, Vec<u8>)>> {
        if let Some((meta, key)) = self.db.search_file_metadata(mxc, 0, 0)? {
            Ok(self.read_content(&key).await?.map(|data| (meta, data)))
        } else {
            Ok(None)
        }
    }

    /// Deletes a media object and all associated thumbnails.
    #[tracing::instrument(skip(self))]
    pub(crate) async fn delete(&self, mxc: OwnedMxcUri) -> Result<()> {
        let mut any_files = false;

        let thumbnails = self.db.search_thumbnails_metadata(mxc.clone())?;
        for (_, thumbnail_key) in thumbnails {
            any_files = true;
            self.delete_by_key(thumbnail_key.clone()).await.inspect_err(
                |error| {
                    warn!(
                        thumbnail_key = utils::u8_slice_to_hex(
                            thumbnail_key.as_bytes()
                        ),
                        %error,
                        "Failed to delete thumbnail media"
                    );
                },
            )?;
        }

        if let Some((_, key)) =
            self.db.search_file_metadata(mxc, 0, 0).inspect_err(
                |error| warn!(%error, "Failed to find original media key"),
            )?
        {
            any_files = true;
            self.delete_by_key(key).await.inspect_err(
                |error| warn!(%error, "Failed to delete original media"),
            )?;
        }

        if any_files {
            Ok(())
        } else {
            let error =
                Error::BadRequest(ErrorKind::NotFound, "Media not found");
            warn!(%error, "Failed to delete media");
            Err(error)
        }
    }

    /// Deletes a specific media key, which may or may not be a thumbnail.
    ///
    /// When deleting a non-thumbnail key with this method, the associated
    /// thumbnails are not deleted.
    async fn delete_by_key(&self, key: MediaFileKey) -> Result<()> {
        let path = services().globals.get_media_file(&key);
        match fs::remove_file(path).await {
            Ok(()) => (),
            // The file in the fs may already have been deleted by hand
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => (),
            // The file may have never existed in the fs because the name was
            // too long
            #[cfg(unix)]
            Err(e) if e.raw_os_error() == Some(nix::libc::ENAMETOOLONG) => (),
            other_error => other_error?,
        }
        self.db.delete_file_metadata(key)?;
        Ok(())
    }

    /// List all media stored in the database.
    ///
    /// Each MXC is listed once. Thumbnails are not included separately from the
    /// original media.
    #[tracing::instrument(skip(self))]
    pub(crate) fn iter_all(&self) -> impl Iterator<Item = Result<OwnedMxcUri>> {
        let mut prev_mxc = None;
        self.db
            .all_file_metadata()
            .map(|media| media.map(|(mxc, ..)| mxc))
            .filter(move |mxc| {
                if let Ok(mxc) = mxc {
                    // Skip mxcs that we have already seen. All files associated
                    // with a given mxc should appear consecutively in the db
                    // iterator, so we only need to check against the previous
                    // value.
                    if prev_mxc.as_ref() == Some(mxc) {
                        return false;
                    }
                    prev_mxc = Some(mxc.clone());
                }
                true
            })
    }

    /// Returns width, height of the thumbnail and whether it should be cropped.
    /// Returns None when the server should send the original file.
    fn thumbnail_properties(
        width: u32,
        height: u32,
    ) -> Option<(u32, u32, bool)> {
        match (width, height) {
            (0..=32, 0..=32) => Some((32, 32, true)),
            (0..=96, 0..=96) => Some((96, 96, true)),
            (0..=320, 0..=240) => Some((320, 240, false)),
            (0..=640, 0..=480) => Some((640, 480, false)),
            (0..=800, 0..=600) => Some((800, 600, false)),
            _ => None,
        }
    }

    /// Generates a thumbnail from the given image file contents. Returns
    /// `Ok(None)` if the input image should be used as-is.
    #[tracing::instrument(
        skip(file),
        fields(input_size = file.len(), original_width, original_height),
    )]
    fn generate_thumbnail(
        file: &[u8],
        width: u32,
        height: u32,
        crop: bool,
    ) -> Result<Option<Vec<u8>>> {
        let image = match image::load_from_memory(file) {
            Ok(image) => image,
            Err(error) => {
                warn!(%error, "Failed to parse source image");
                return Ok(None);
            }
        };

        let original_width = image.width();
        let original_height = image.height();
        tracing::Span::current().record("original_width", original_width);
        tracing::Span::current().record("original_height", original_height);

        if width > original_width || height > original_height {
            debug!("Requested thumbnail is larger than source image");
            return Ok(None);
        }

        let thumbnail = if crop {
            image.resize_to_fill(width, height, FilterType::CatmullRom)
        } else {
            let (exact_width, exact_height) = {
                // Copied from image::dynimage::resize_dimensions
                let use_width = (u64::from(width) * u64::from(original_height))
                    <= (u64::from(original_width) * u64::from(height));
                let intermediate = if use_width {
                    u64::from(original_height) * u64::from(width)
                        / u64::from(original_width)
                } else {
                    u64::from(original_width) * u64::from(height)
                        / u64::from(original_height)
                };
                if use_width {
                    if let Ok(intermediate) = u32::try_from(intermediate) {
                        (width, intermediate)
                    } else {
                        (
                            (u64::from(width) * u64::from(u32::MAX)
                                / intermediate)
                                .try_into()
                                .unwrap_or(u32::MAX),
                            u32::MAX,
                        )
                    }
                } else if let Ok(intermediate) = u32::try_from(intermediate) {
                    (intermediate, height)
                } else {
                    (
                        u32::MAX,
                        (u64::from(height) * u64::from(u32::MAX)
                            / intermediate)
                            .try_into()
                            .unwrap_or(u32::MAX),
                    )
                }
            };

            image.thumbnail_exact(exact_width, exact_height)
        };

        debug!("Serializing thumbnail as PNG");
        let mut thumbnail_bytes = Vec::new();
        thumbnail.write_to(
            &mut Cursor::new(&mut thumbnail_bytes),
            image::ImageFormat::Png,
        )?;

        Ok(Some(thumbnail_bytes))
    }

    /// Downloads a file's thumbnail.
    ///
    /// Here's an example on how it works:
    ///
    /// - Client requests an image with width=567, height=567
    /// - Server rounds that up to (800, 600), so it doesn't have to save too
    ///   many thumbnails
    /// - Server rounds that up again to (958, 600) to fix the aspect ratio
    ///   (only for width,height>96)
    /// - Server creates the thumbnail and sends it to the user
    ///
    /// For width,height <= 96 the server uses another thumbnailing algorithm
    /// which crops the image afterwards.
    #[tracing::instrument(skip(self))]
    pub(crate) async fn get_thumbnail(
        &self,
        mxc: OwnedMxcUri,
        width: u32,
        height: u32,
    ) -> Result<Option<(FileMeta, Vec<u8>)>> {
        // 0, 0 because that's the original file
        let (width, height, crop) =
            Self::thumbnail_properties(width, height).unwrap_or((0, 0, false));

        if let Some((meta, key)) =
            self.db.search_file_metadata(mxc.clone(), width, height)?
        {
            debug!("Using saved thumbnail");
            return Ok(self.read_content(&key).await?.map(|file| (meta, file)));
        }

        let Some((meta, key)) =
            self.db.search_file_metadata(mxc.clone(), 0, 0)?
        else {
            debug!("Original image not found, can't generate thumbnail");
            return Ok(None);
        };

        let Some(file) = self.read_content(&key).await? else {
            debug!("Original image not found, can't generate thumbnail");
            return Ok(None);
        };

        debug!("Generating thumbnail");
        let thumbnail_result = {
            let file = file.clone();
            let outer_span = tracing::span::Span::current();

            tokio::task::spawn_blocking(move || {
                outer_span.in_scope(|| {
                    Self::generate_thumbnail(&file, width, height, crop)
                })
            })
            .await
            .expect("failed to join thumbnailer task")
        };

        let Some(thumbnail_bytes) = thumbnail_result? else {
            debug!("Returning source image as-is");
            return Ok(Some((meta, file)));
        };

        debug!("Saving created thumbnail");

        // Save thumbnail in database so we don't have to generate it
        // again next time
        let thumbnail_key =
            self.db.create_file_metadata(mxc, width, height, &meta)?;
        self.write_content(&thumbnail_key, &thumbnail_bytes).await?;

        Ok(Some((meta, thumbnail_bytes.clone())))
    }

    /// Writes contents for a media file to the fs.
    ///
    /// If a file already existed with the specified key, it is replaced.
    async fn write_content(
        &self,
        key: &MediaFileKey,
        data: &[u8],
    ) -> Result<()> {
        let path = services().globals.get_media_file(key);
        let mut file = File::create(path).await?;

        file.write_all(data).await?;

        Ok(())
    }

    /// Returns the contents of a media file, read from the fs.
    ///
    /// If the file cannot be opened, returns `Ok(None)`. This is needed because
    /// before media deletion admin commands were implemented, the only way to
    /// delete abuse media was to remove the associated files from the fs. This
    /// leaves the db in an inconsistent state, where media keys exist in the db
    /// but their content files do not. We want to return `M_NOT_YET_UPLOADED`
    /// in this case rather than the `M_UNKNOWN` we would normally use for db
    /// consistency problems.
    async fn read_content(
        &self,
        key: &MediaFileKey,
    ) -> Result<Option<Vec<u8>>> {
        let path = services().globals.get_media_file(key);
        let mut file = match File::open(path).await {
            Ok(file) => file,
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
                return Ok(None)
            }
            Err(e) => return Err(e.into()),
        };

        let mut data = Vec::new();
        file.read_to_end(&mut data).await?;

        Ok(Some(data))
    }
}