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, 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<Option<String>> {
241 let Ok(buf) = base64::engine::general_purpose::STANDARD.decode(data) else {
242 return Ok(None);
243 };
244 let name = if let Ok(format) = image::guess_format(&buf) {
245 if let Some(ext) = format.extensions_str().first() {
246 format!("file.{ext}")
247 } else {
248 String::new()
249 }
250 } else {
251 String::new()
252 };
253 let blob = BlobObject::create_and_deduplicate_from_bytes(context, &buf, &name)?;
254 Ok(Some(blob.as_name().to_string()))
255 }
256
257 pub async fn recode_to_avatar_size(&mut self, context: &Context) -> Result<()> {
259 let (img_wh, max_bytes) =
260 match MediaQuality::from_i32(context.get_config_int(Config::MediaQuality).await?)
261 .unwrap_or_default()
262 {
263 MediaQuality::Balanced => (
264 constants::BALANCED_AVATAR_SIZE,
265 constants::BALANCED_AVATAR_BYTES,
266 ),
267 MediaQuality::Worse => {
268 (constants::WORSE_AVATAR_SIZE, constants::WORSE_AVATAR_BYTES)
269 }
270 };
271
272 let viewtype = &mut Viewtype::Image;
273 let is_avatar = true;
274 self.check_or_recode_to_size(
275 context, None, viewtype, img_wh, max_bytes, is_avatar,
277 )?;
278
279 Ok(())
280 }
281
282 pub async fn check_or_recode_image(
292 &mut self,
293 context: &Context,
294 name: Option<String>,
295 viewtype: &mut Viewtype,
296 ) -> Result<String> {
297 let (img_wh, max_bytes) =
298 match MediaQuality::from_i32(context.get_config_int(Config::MediaQuality).await?)
299 .unwrap_or_default()
300 {
301 MediaQuality::Balanced => (
302 constants::BALANCED_IMAGE_SIZE,
303 constants::BALANCED_IMAGE_BYTES,
304 ),
305 MediaQuality::Worse => (constants::WORSE_IMAGE_SIZE, constants::WORSE_IMAGE_BYTES),
306 };
307 let is_avatar = false;
308 self.check_or_recode_to_size(context, name, viewtype, img_wh, max_bytes, is_avatar)
309 }
310
311 fn check_or_recode_to_size(
323 &mut self,
324 context: &Context,
325 name: Option<String>,
326 viewtype: &mut Viewtype,
327 mut img_wh: u32,
328 max_bytes: usize,
329 is_avatar: bool,
330 ) -> Result<String> {
331 let mut add_white_bg = is_avatar;
333 let mut no_exif = false;
334 let no_exif_ref = &mut no_exif;
335 let mut name = name.unwrap_or_else(|| self.name.clone());
336 let original_name = name.clone();
337 let vt = &mut *viewtype;
338 let res: Result<String> = tokio::task::block_in_place(move || {
339 let mut file = std::fs::File::open(self.to_abs_path())?;
340 let (nr_bytes, exif) = image_metadata(&file)?;
341 *no_exif_ref = exif.is_none();
342 file.rewind()?;
345 let imgreader = ImageReader::new(std::io::BufReader::new(&file)).with_guessed_format();
346 let imgreader = match imgreader {
347 Ok(ir) => ir,
348 _ => {
349 file.rewind()?;
350 ImageReader::with_format(
351 std::io::BufReader::new(&file),
352 ImageFormat::from_path(self.to_abs_path())?,
353 )
354 }
355 };
356 let fmt = imgreader.format().context("Unknown format")?;
357 if *vt == Viewtype::File {
358 *vt = Viewtype::Image;
359 return Ok(name);
360 }
361 let mut img = imgreader.decode().context("image decode failure")?;
362 let orientation = exif.as_ref().map(|exif| exif_orientation(exif, context));
363 let mut encoded = Vec::new();
364
365 if *vt == Viewtype::Sticker {
366 let x_max = img.width().saturating_sub(1);
367 let y_max = img.height().saturating_sub(1);
368 if !img.in_bounds(x_max, y_max)
369 || !(img.get_pixel(0, 0).0[3] == 0
370 || img.get_pixel(x_max, 0).0[3] == 0
371 || img.get_pixel(0, y_max).0[3] == 0
372 || img.get_pixel(x_max, y_max).0[3] == 0)
373 {
374 *vt = Viewtype::Image;
375 } else {
376 return Ok(name);
379 }
380 }
381
382 img = match orientation {
383 Some(90) => img.rotate90(),
384 Some(180) => img.rotate180(),
385 Some(270) => img.rotate270(),
386 _ => img,
387 };
388
389 let exceeds_wh = img.width() > img_wh || img.height() > img_wh;
390 let exceeds_max_bytes = nr_bytes > max_bytes as u64;
391
392 let jpeg_quality = 75;
393 let ofmt = match fmt {
394 ImageFormat::Png if !exceeds_max_bytes => ImageOutputFormat::Png,
395 ImageFormat::Jpeg => {
396 add_white_bg = false;
397 ImageOutputFormat::Jpeg {
398 quality: jpeg_quality,
399 }
400 }
401 _ => ImageOutputFormat::Jpeg {
402 quality: jpeg_quality,
403 },
404 };
405 let do_scale = exceeds_max_bytes
412 || is_avatar
413 && (exceeds_wh
414 || exif.is_some() && {
415 if mem::take(&mut add_white_bg) {
416 self::add_white_bg(&mut img);
417 }
418 encoded_img_exceeds_bytes(
419 context,
420 &img,
421 ofmt.clone(),
422 max_bytes,
423 &mut encoded,
424 )?
425 });
426
427 if do_scale {
428 if !exceeds_wh {
429 img_wh = max(img.width(), img.height());
430 if matches!(fmt, ImageFormat::Jpeg) || !encoded.is_empty() {
433 img_wh = img_wh * 2 / 3;
434 }
435 }
436
437 loop {
438 if mem::take(&mut add_white_bg) {
439 self::add_white_bg(&mut img);
440 }
441
442 let new_img = if is_avatar {
451 img.resize(img_wh, img_wh, image::imageops::FilterType::Triangle)
452 } else {
453 img.thumbnail(img_wh, img_wh)
454 };
455
456 if encoded_img_exceeds_bytes(
457 context,
458 &new_img,
459 ofmt.clone(),
460 max_bytes,
461 &mut encoded,
462 )? && is_avatar
463 {
464 if img_wh < 20 {
465 return Err(format_err!(
466 "Failed to scale image to below {max_bytes}B.",
467 ));
468 }
469
470 img_wh = img_wh * 2 / 3;
471 } else {
472 info!(
473 context,
474 "Final scaled-down image size: {}B ({}px).",
475 encoded.len(),
476 img_wh
477 );
478 break;
479 }
480 }
481 }
482
483 if do_scale || exif.is_some() {
484 if !matches!(fmt, ImageFormat::Jpeg)
486 && matches!(ofmt, ImageOutputFormat::Jpeg { .. })
487 {
488 name = Path::new(&name)
489 .with_extension("jpg")
490 .to_string_lossy()
491 .into_owned();
492 }
493
494 if encoded.is_empty() {
495 if mem::take(&mut add_white_bg) {
496 self::add_white_bg(&mut img);
497 }
498 encode_img(&img, ofmt, &mut encoded)?;
499 }
500
501 self.name = BlobObject::create_and_deduplicate_from_bytes(context, &encoded, &name)
502 .context("failed to write recoded blob to file")?
503 .name;
504 }
505
506 Ok(name)
507 });
508 match res {
509 Ok(_) => res,
510 Err(err) => {
511 if !is_avatar && no_exif {
512 error!(
513 context,
514 "Cannot check/recode image, using original data: {err:#}.",
515 );
516 *viewtype = Viewtype::File;
517 Ok(original_name)
518 } else {
519 Err(err)
520 }
521 }
522 }
523 }
524}
525
526fn file_hash(src: &Path) -> Result<blake3::Hash> {
527 ensure!(
528 !src.starts_with("$BLOBDIR/"),
529 "Use `get_abs_path()` to get the absolute path of the blobfile"
530 );
531 let mut hasher = blake3::Hasher::new();
532 let mut src_file = std::fs::File::open(src)
533 .with_context(|| format!("Failed to open file {}", src.display()))?;
534 hasher
535 .update_reader(&mut src_file)
536 .context("update_reader")?;
537 let hash = hasher.finalize();
538 Ok(hash)
539}
540
541fn image_metadata(file: &std::fs::File) -> Result<(u64, Option<exif::Exif>)> {
543 let len = file.metadata()?.len();
544 let mut bufreader = std::io::BufReader::new(file);
545 let exif = exif::Reader::new()
546 .continue_on_error(true)
547 .read_from_container(&mut bufreader)
548 .or_else(|e| e.distill_partial_result(|_errors| {}))
549 .ok();
550 Ok((len, exif))
551}
552
553fn exif_orientation(exif: &exif::Exif, context: &Context) -> i32 {
554 if let Some(orientation) = exif.get_field(exif::Tag::Orientation, exif::In::PRIMARY) {
555 match orientation.value.get_uint(0) {
558 Some(3) => return 180,
559 Some(6) => return 90,
560 Some(8) => return 270,
561 other => warn!(context, "Exif orientation value ignored: {other:?}."),
562 }
563 }
564 0
565}
566
567pub(crate) struct BlobDirContents<'a> {
574 inner: Vec<PathBuf>,
575 context: &'a Context,
576}
577
578impl<'a> BlobDirContents<'a> {
579 pub(crate) async fn new(context: &'a Context) -> Result<BlobDirContents<'a>> {
580 let readdir = fs::read_dir(context.get_blobdir()).await?;
581 let inner = ReadDirStream::new(readdir)
582 .filter_map(|entry| async move {
583 match entry {
584 Ok(entry) => Some(entry),
585 Err(err) => {
586 error!(context, "Failed to read blob file: {err}.");
587 None
588 }
589 }
590 })
591 .filter_map(|entry| async move {
592 match entry.file_type().await.ok()?.is_file() {
593 true => Some(entry.path()),
594 false => {
595 warn!(
596 context,
597 "Export: Found blob dir entry {} that is not a file, ignoring.",
598 entry.path().display()
599 );
600 None
601 }
602 }
603 })
604 .collect()
605 .await;
606 Ok(Self { inner, context })
607 }
608
609 pub(crate) fn iter(&self) -> BlobDirIter<'_> {
610 BlobDirIter::new(self.context, self.inner.iter())
611 }
612}
613
614pub(crate) struct BlobDirIter<'a> {
616 iter: std::slice::Iter<'a, PathBuf>,
617 context: &'a Context,
618}
619
620impl<'a> BlobDirIter<'a> {
621 fn new(context: &'a Context, iter: std::slice::Iter<'a, PathBuf>) -> BlobDirIter<'a> {
622 Self { iter, context }
623 }
624}
625
626impl<'a> Iterator for BlobDirIter<'a> {
627 type Item = BlobObject<'a>;
628
629 fn next(&mut self) -> Option<Self::Item> {
630 for path in self.iter.by_ref() {
631 match BlobObject::from_path(self.context, path) {
634 Ok(blob) => return Some(blob),
635 Err(err) => warn!(self.context, "{err}"),
636 }
637 }
638 None
639 }
640}
641
642impl FusedIterator for BlobDirIter<'_> {}
643
644fn encode_img(
645 img: &DynamicImage,
646 fmt: ImageOutputFormat,
647 encoded: &mut Vec<u8>,
648) -> anyhow::Result<()> {
649 encoded.clear();
650 let mut buf = Cursor::new(encoded);
651 match fmt {
652 ImageOutputFormat::Png => img.write_to(&mut buf, ImageFormat::Png)?,
653 ImageOutputFormat::Jpeg { quality } => {
654 let encoder = JpegEncoder::new_with_quality(&mut buf, quality);
655 img.clone().into_rgb8().write_with_encoder(encoder)?;
659 }
660 }
661 Ok(())
662}
663
664fn encoded_img_exceeds_bytes(
665 context: &Context,
666 img: &DynamicImage,
667 fmt: ImageOutputFormat,
668 max_bytes: usize,
669 encoded: &mut Vec<u8>,
670) -> anyhow::Result<bool> {
671 encode_img(img, fmt, encoded)?;
672 if encoded.len() > max_bytes {
673 info!(
674 context,
675 "Image size {}B ({}x{}px) exceeds {}B, need to scale down.",
676 encoded.len(),
677 img.width(),
678 img.height(),
679 max_bytes,
680 );
681 return Ok(true);
682 }
683 Ok(false)
684}
685
686fn add_white_bg(img: &mut DynamicImage) {
688 for y in 0..img.height() {
689 for x in 0..img.width() {
690 let mut p = Rgba([255u8, 255, 255, 255]);
691 p.blend(&img.get_pixel(x, y));
692 img.put_pixel(x, y, p);
693 }
694 }
695}
696
697#[cfg(test)]
698mod blob_tests;