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