1use core::cmp::max;
4use std::io::{Cursor, Seek};
5use std::iter::FusedIterator;
6use std::mem;
7use std::path::{Path, PathBuf};
8
9use anyhow::{Context as _, Result, ensure, format_err};
10use base64::Engine as _;
11use futures::StreamExt;
12use image::ImageReader;
13use image::codecs::jpeg::JpegEncoder;
14use image::{DynamicImage, GenericImage, GenericImageView, ImageFormat, Pixel, Rgba};
15use num_traits::FromPrimitive;
16use tokio::{fs, task};
17use tokio_stream::wrappers::ReadDirStream;
18
19use crate::config::Config;
20use crate::constants::{self, MediaQuality};
21use crate::context::Context;
22use crate::events::EventType;
23use crate::log::{LogExt, error, info, warn};
24use crate::message::Viewtype;
25use crate::tools::sanitize_filename;
26
27#[derive(Debug, Clone, PartialEq, Eq)]
34pub struct BlobObject<'a> {
35    blobdir: &'a Path,
36
37    name: String,
41}
42
43#[derive(Debug, Clone)]
44enum ImageOutputFormat {
45    Png,
46    Jpeg { quality: u8 },
47}
48
49impl<'a> BlobObject<'a> {
50    pub fn create_and_deduplicate(
61        context: &'a Context,
62        src: &Path,
63        original_name: &Path,
64    ) -> Result<BlobObject<'a>> {
65        task::block_in_place(|| {
70            let temp_path;
71            let src_in_blobdir: &Path;
72            let blobdir = context.get_blobdir();
73
74            if src.starts_with(blobdir) {
75                src_in_blobdir = src;
76            } else {
77                info!(
78                    context,
79                    "Source file not in blobdir. Copying instead of moving in order to prevent moving a file that was still needed."
80                );
81                temp_path = blobdir.join(format!("tmp-{}", rand::random::<u64>()));
82                if std::fs::copy(src, &temp_path).is_err() {
83                    std::fs::create_dir_all(blobdir).log_err(context).ok();
85                    std::fs::copy(src, &temp_path).context("Copying new blobfile failed")?;
86                };
87                src_in_blobdir = &temp_path;
88            }
89
90            let hash = file_hash(src_in_blobdir)?.to_hex();
91            let hash = hash.as_str();
92            let hash = hash.get(0..31).unwrap_or(hash);
93            let new_file =
94                if let Some(extension) = original_name.extension().filter(|e| e.len() <= 32) {
95                    let extension = extension.to_string_lossy().to_lowercase();
96                    let extension = sanitize_filename(&extension);
97                    format!("$BLOBDIR/{hash}.{extension}")
98                } else {
99                    format!("$BLOBDIR/{hash}")
100                };
101
102            let blob = BlobObject {
103                blobdir,
104                name: new_file,
105            };
106            let new_path = blob.to_abs_path();
107
108            std::fs::rename(src_in_blobdir, &new_path)?;
111
112            context.emit_event(EventType::NewBlobFile(blob.as_name().to_string()));
113            Ok(blob)
114        })
115    }
116
117    pub fn create_and_deduplicate_from_bytes(
127        context: &'a Context,
128        data: &[u8],
129        original_name: &str,
130    ) -> Result<BlobObject<'a>> {
131        task::block_in_place(|| {
132            let blobdir = context.get_blobdir();
133            let temp_path = blobdir.join(format!("tmp-{}", rand::random::<u64>()));
134            if std::fs::write(&temp_path, data).is_err() {
135                std::fs::create_dir_all(blobdir).log_err(context).ok();
137                std::fs::write(&temp_path, data).context("writing new blobfile failed")?;
138            };
139
140            BlobObject::create_and_deduplicate(context, &temp_path, Path::new(original_name))
141        })
142    }
143
144    pub fn from_path(context: &'a Context, path: &Path) -> Result<BlobObject<'a>> {
151        let rel_path = path
152            .strip_prefix(context.get_blobdir())
153            .with_context(|| format!("wrong blobdir: {}", path.display()))?;
154        let name = rel_path.to_str().context("wrong name")?;
155        if !BlobObject::is_acceptible_blob_name(name) {
156            return Err(format_err!("bad blob name: {}", rel_path.display()));
157        }
158        BlobObject::from_name(context, name)
159    }
160
161    pub fn from_name(context: &'a Context, name: &str) -> Result<BlobObject<'a>> {
168        let name = match name.starts_with("$BLOBDIR/") {
169            true => name.splitn(2, '/').last().unwrap(),
170            false => name,
171        };
172        if !BlobObject::is_acceptible_blob_name(name) {
173            return Err(format_err!("not an acceptable blob name: {name}"));
174        }
175        Ok(BlobObject {
176            blobdir: context.get_blobdir(),
177            name: format!("$BLOBDIR/{name}"),
178        })
179    }
180
181    pub fn to_abs_path(&self) -> PathBuf {
183        let fname = Path::new(&self.name).strip_prefix("$BLOBDIR/").unwrap();
184        self.blobdir.join(fname)
185    }
186
187    #[allow(rustdoc::private_intra_doc_links)]
198    pub fn as_name(&self) -> &str {
200        &self.name
201    }
202
203    pub fn suffix(&self) -> Option<&str> {
208        let ext = self.name.rsplit('.').next();
209        if ext == Some(&self.name) { None } else { ext }
210    }
211
212    fn is_acceptible_blob_name(name: &str) -> bool {
220        if name.find('/').is_some() {
221            return false;
222        }
223        if name.find('\\').is_some() {
224            return false;
225        }
226        if name.find('\0').is_some() {
227            return false;
228        }
229        true
230    }
231
232    pub(crate) fn store_from_base64(context: &Context, data: &str) -> Result<String> {
238        let buf = base64::engine::general_purpose::STANDARD.decode(data)?;
239        let name = if let Ok(format) = image::guess_format(&buf) {
240            if let Some(ext) = format.extensions_str().first() {
241                format!("file.{ext}")
242            } else {
243                String::new()
244            }
245        } else {
246            String::new()
247        };
248        let blob = BlobObject::create_and_deduplicate_from_bytes(context, &buf, &name)?;
249        Ok(blob.as_name().to_string())
250    }
251
252    pub async fn recode_to_avatar_size(&mut self, context: &Context) -> Result<()> {
254        let (img_wh, max_bytes) =
255            match MediaQuality::from_i32(context.get_config_int(Config::MediaQuality).await?)
256                .unwrap_or_default()
257            {
258                MediaQuality::Balanced => (
259                    constants::BALANCED_AVATAR_SIZE,
260                    constants::BALANCED_AVATAR_BYTES,
261                ),
262                MediaQuality::Worse => {
263                    (constants::WORSE_AVATAR_SIZE, constants::WORSE_AVATAR_BYTES)
264                }
265            };
266
267        let viewtype = &mut Viewtype::Image;
268        let is_avatar = true;
269        self.check_or_recode_to_size(
270            context, None, viewtype, img_wh, max_bytes, is_avatar,
272        )?;
273
274        Ok(())
275    }
276
277    pub async fn check_or_recode_image(
287        &mut self,
288        context: &Context,
289        name: Option<String>,
290        viewtype: &mut Viewtype,
291    ) -> Result<String> {
292        let (img_wh, max_bytes) =
293            match MediaQuality::from_i32(context.get_config_int(Config::MediaQuality).await?)
294                .unwrap_or_default()
295            {
296                MediaQuality::Balanced => (
297                    constants::BALANCED_IMAGE_SIZE,
298                    constants::BALANCED_IMAGE_BYTES,
299                ),
300                MediaQuality::Worse => (constants::WORSE_IMAGE_SIZE, constants::WORSE_IMAGE_BYTES),
301            };
302        let is_avatar = false;
303        self.check_or_recode_to_size(context, name, viewtype, img_wh, max_bytes, is_avatar)
304    }
305
306    fn check_or_recode_to_size(
318        &mut self,
319        context: &Context,
320        name: Option<String>,
321        viewtype: &mut Viewtype,
322        mut img_wh: u32,
323        max_bytes: usize,
324        is_avatar: bool,
325    ) -> Result<String> {
326        let mut add_white_bg = is_avatar;
328        let mut no_exif = false;
329        let no_exif_ref = &mut no_exif;
330        let mut name = name.unwrap_or_else(|| self.name.clone());
331        let original_name = name.clone();
332        let vt = &mut *viewtype;
333        let res: Result<String> = tokio::task::block_in_place(move || {
334            let mut file = std::fs::File::open(self.to_abs_path())?;
335            let (nr_bytes, exif) = image_metadata(&file)?;
336            *no_exif_ref = exif.is_none();
337            file.rewind()?;
340            let imgreader = ImageReader::new(std::io::BufReader::new(&file)).with_guessed_format();
341            let imgreader = match imgreader {
342                Ok(ir) => ir,
343                _ => {
344                    file.rewind()?;
345                    ImageReader::with_format(
346                        std::io::BufReader::new(&file),
347                        ImageFormat::from_path(self.to_abs_path())?,
348                    )
349                }
350            };
351            let fmt = imgreader.format().context("Unknown format")?;
352            if *vt == Viewtype::File {
353                *vt = Viewtype::Image;
354                return Ok(name);
355            }
356            let mut img = imgreader.decode().context("image decode failure")?;
357            let orientation = exif.as_ref().map(|exif| exif_orientation(exif, context));
358            let mut encoded = Vec::new();
359
360            if *vt == Viewtype::Sticker {
361                let x_max = img.width().saturating_sub(1);
362                let y_max = img.height().saturating_sub(1);
363                if !img.in_bounds(x_max, y_max)
364                    || !(img.get_pixel(0, 0).0[3] == 0
365                        || img.get_pixel(x_max, 0).0[3] == 0
366                        || img.get_pixel(0, y_max).0[3] == 0
367                        || img.get_pixel(x_max, y_max).0[3] == 0)
368                {
369                    *vt = Viewtype::Image;
370                } else {
371                    return Ok(name);
374                }
375            }
376
377            img = match orientation {
378                Some(90) => img.rotate90(),
379                Some(180) => img.rotate180(),
380                Some(270) => img.rotate270(),
381                _ => img,
382            };
383
384            let exceeds_wh = img.width() > img_wh || img.height() > img_wh;
385            let exceeds_max_bytes = nr_bytes > max_bytes as u64;
386
387            let jpeg_quality = 75;
388            let ofmt = match fmt {
389                ImageFormat::Png if !exceeds_max_bytes => ImageOutputFormat::Png,
390                ImageFormat::Jpeg => {
391                    add_white_bg = false;
392                    ImageOutputFormat::Jpeg {
393                        quality: jpeg_quality,
394                    }
395                }
396                _ => ImageOutputFormat::Jpeg {
397                    quality: jpeg_quality,
398                },
399            };
400            let do_scale = exceeds_max_bytes
407                || is_avatar
408                    && (exceeds_wh
409                        || exif.is_some() && {
410                            if mem::take(&mut add_white_bg) {
411                                self::add_white_bg(&mut img);
412                            }
413                            encoded_img_exceeds_bytes(
414                                context,
415                                &img,
416                                ofmt.clone(),
417                                max_bytes,
418                                &mut encoded,
419                            )?
420                        });
421
422            if do_scale {
423                if !exceeds_wh {
424                    img_wh = max(img.width(), img.height());
425                    if matches!(fmt, ImageFormat::Jpeg) || !encoded.is_empty() {
428                        img_wh = img_wh * 2 / 3;
429                    }
430                }
431
432                loop {
433                    if mem::take(&mut add_white_bg) {
434                        self::add_white_bg(&mut img);
435                    }
436
437                    let new_img = if is_avatar {
446                        img.resize(img_wh, img_wh, image::imageops::FilterType::Triangle)
447                    } else {
448                        img.thumbnail(img_wh, img_wh)
449                    };
450
451                    if encoded_img_exceeds_bytes(
452                        context,
453                        &new_img,
454                        ofmt.clone(),
455                        max_bytes,
456                        &mut encoded,
457                    )? && is_avatar
458                    {
459                        if img_wh < 20 {
460                            return Err(format_err!(
461                                "Failed to scale image to below {max_bytes}B.",
462                            ));
463                        }
464
465                        img_wh = img_wh * 2 / 3;
466                    } else {
467                        info!(
468                            context,
469                            "Final scaled-down image size: {}B ({}px).",
470                            encoded.len(),
471                            img_wh
472                        );
473                        break;
474                    }
475                }
476            }
477
478            if do_scale || exif.is_some() {
479                if !matches!(fmt, ImageFormat::Jpeg)
481                    && matches!(ofmt, ImageOutputFormat::Jpeg { .. })
482                {
483                    name = Path::new(&name)
484                        .with_extension("jpg")
485                        .to_string_lossy()
486                        .into_owned();
487                }
488
489                if encoded.is_empty() {
490                    if mem::take(&mut add_white_bg) {
491                        self::add_white_bg(&mut img);
492                    }
493                    encode_img(&img, ofmt, &mut encoded)?;
494                }
495
496                self.name = BlobObject::create_and_deduplicate_from_bytes(context, &encoded, &name)
497                    .context("failed to write recoded blob to file")?
498                    .name;
499            }
500
501            Ok(name)
502        });
503        match res {
504            Ok(_) => res,
505            Err(err) => {
506                if !is_avatar && no_exif {
507                    error!(
508                        context,
509                        "Cannot check/recode image, using original data: {err:#}.",
510                    );
511                    *viewtype = Viewtype::File;
512                    Ok(original_name)
513                } else {
514                    Err(err)
515                }
516            }
517        }
518    }
519}
520
521fn file_hash(src: &Path) -> Result<blake3::Hash> {
522    ensure!(
523        !src.starts_with("$BLOBDIR/"),
524        "Use `get_abs_path()` to get the absolute path of the blobfile"
525    );
526    let mut hasher = blake3::Hasher::new();
527    let mut src_file = std::fs::File::open(src)
528        .with_context(|| format!("Failed to open file {}", src.display()))?;
529    hasher
530        .update_reader(&mut src_file)
531        .context("update_reader")?;
532    let hash = hasher.finalize();
533    Ok(hash)
534}
535
536fn image_metadata(file: &std::fs::File) -> Result<(u64, Option<exif::Exif>)> {
538    let len = file.metadata()?.len();
539    let mut bufreader = std::io::BufReader::new(file);
540    let exif = exif::Reader::new()
541        .continue_on_error(true)
542        .read_from_container(&mut bufreader)
543        .or_else(|e| e.distill_partial_result(|_errors| {}))
544        .ok();
545    Ok((len, exif))
546}
547
548fn exif_orientation(exif: &exif::Exif, context: &Context) -> i32 {
549    if let Some(orientation) = exif.get_field(exif::Tag::Orientation, exif::In::PRIMARY) {
550        match orientation.value.get_uint(0) {
553            Some(3) => return 180,
554            Some(6) => return 90,
555            Some(8) => return 270,
556            other => warn!(context, "Exif orientation value ignored: {other:?}."),
557        }
558    }
559    0
560}
561
562pub(crate) struct BlobDirContents<'a> {
569    inner: Vec<PathBuf>,
570    context: &'a Context,
571}
572
573impl<'a> BlobDirContents<'a> {
574    pub(crate) async fn new(context: &'a Context) -> Result<BlobDirContents<'a>> {
575        let readdir = fs::read_dir(context.get_blobdir()).await?;
576        let inner = ReadDirStream::new(readdir)
577            .filter_map(|entry| async move {
578                match entry {
579                    Ok(entry) => Some(entry),
580                    Err(err) => {
581                        error!(context, "Failed to read blob file: {err}.");
582                        None
583                    }
584                }
585            })
586            .filter_map(|entry| async move {
587                match entry.file_type().await.ok()?.is_file() {
588                    true => Some(entry.path()),
589                    false => {
590                        warn!(
591                            context,
592                            "Export: Found blob dir entry {} that is not a file, ignoring.",
593                            entry.path().display()
594                        );
595                        None
596                    }
597                }
598            })
599            .collect()
600            .await;
601        Ok(Self { inner, context })
602    }
603
604    pub(crate) fn iter(&self) -> BlobDirIter<'_> {
605        BlobDirIter::new(self.context, self.inner.iter())
606    }
607}
608
609pub(crate) struct BlobDirIter<'a> {
611    iter: std::slice::Iter<'a, PathBuf>,
612    context: &'a Context,
613}
614
615impl<'a> BlobDirIter<'a> {
616    fn new(context: &'a Context, iter: std::slice::Iter<'a, PathBuf>) -> BlobDirIter<'a> {
617        Self { iter, context }
618    }
619}
620
621impl<'a> Iterator for BlobDirIter<'a> {
622    type Item = BlobObject<'a>;
623
624    fn next(&mut self) -> Option<Self::Item> {
625        for path in self.iter.by_ref() {
626            match BlobObject::from_path(self.context, path) {
629                Ok(blob) => return Some(blob),
630                Err(err) => warn!(self.context, "{err}"),
631            }
632        }
633        None
634    }
635}
636
637impl FusedIterator for BlobDirIter<'_> {}
638
639fn encode_img(
640    img: &DynamicImage,
641    fmt: ImageOutputFormat,
642    encoded: &mut Vec<u8>,
643) -> anyhow::Result<()> {
644    encoded.clear();
645    let mut buf = Cursor::new(encoded);
646    match fmt {
647        ImageOutputFormat::Png => img.write_to(&mut buf, ImageFormat::Png)?,
648        ImageOutputFormat::Jpeg { quality } => {
649            let encoder = JpegEncoder::new_with_quality(&mut buf, quality);
650            img.clone().into_rgb8().write_with_encoder(encoder)?;
654        }
655    }
656    Ok(())
657}
658
659fn encoded_img_exceeds_bytes(
660    context: &Context,
661    img: &DynamicImage,
662    fmt: ImageOutputFormat,
663    max_bytes: usize,
664    encoded: &mut Vec<u8>,
665) -> anyhow::Result<bool> {
666    encode_img(img, fmt, encoded)?;
667    if encoded.len() > max_bytes {
668        info!(
669            context,
670            "Image size {}B ({}x{}px) exceeds {}B, need to scale down.",
671            encoded.len(),
672            img.width(),
673            img.height(),
674            max_bytes,
675        );
676        return Ok(true);
677    }
678    Ok(false)
679}
680
681fn add_white_bg(img: &mut DynamicImage) {
683    for y in 0..img.height() {
684        for x in 0..img.width() {
685            let mut p = Rgba([255u8, 255, 255, 255]);
686            p.blend(&img.get_pixel(x, y));
687            img.put_pixel(x, y, p);
688        }
689    }
690}
691
692#[cfg(test)]
693mod blob_tests;