notedeck

One damus client to rule them all
git clone git://jb55.com/notedeck
Log | Files | Refs | README | LICENSE

imgcache.rs (12439B)


      1 use crate::jobs::MediaJobSender;
      2 use crate::media::gif::AnimatedImgTexCache;
      3 use crate::media::images::ImageType;
      4 use crate::media::static_imgs::StaticImgTexCache;
      5 use crate::media::{
      6     AnimationMode, BlurCache, NoLoadingLatestTex, TrustedMediaLatestTex, UntrustedMediaLatestTex,
      7 };
      8 use crate::urls::{UrlCache, UrlMimes};
      9 use crate::ImageMetadata;
     10 use crate::ObfuscationType;
     11 use crate::RenderableMedia;
     12 use crate::Result;
     13 use egui::TextureHandle;
     14 use image::{Delay, Frame};
     15 
     16 use egui::ColorImage;
     17 
     18 use std::collections::HashMap;
     19 use std::fs::{self, create_dir_all, File};
     20 use std::sync::{Arc, Mutex};
     21 use std::time::{Duration, Instant, SystemTime};
     22 use std::{io, thread};
     23 
     24 use hex::ToHex;
     25 use sha2::Digest;
     26 use std::path::PathBuf;
     27 use std::path::{self, Path};
     28 use tracing::warn;
     29 
     30 pub struct TexturesCache {
     31     pub static_image: StaticImgTexCache,
     32     pub blurred: BlurCache,
     33     pub animated: AnimatedImgTexCache,
     34 }
     35 
     36 impl TexturesCache {
     37     pub fn new(base_dir: PathBuf) -> Self {
     38         Self {
     39             static_image: StaticImgTexCache::new(
     40                 base_dir.join(MediaCache::rel_dir(MediaCacheType::Image)),
     41             ),
     42             blurred: Default::default(),
     43             animated: AnimatedImgTexCache::new(
     44                 base_dir.join(MediaCache::rel_dir(MediaCacheType::Gif)),
     45             ),
     46         }
     47     }
     48 }
     49 
     50 pub enum TextureState<T> {
     51     Pending,
     52     Error(crate::Error),
     53     Loaded(T),
     54 }
     55 
     56 impl<T> std::fmt::Debug for TextureState<T> {
     57     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
     58         match self {
     59             Self::Pending => write!(f, "Pending"),
     60             Self::Error(_) => f.debug_tuple("Error").field(&"").finish(),
     61             Self::Loaded(_) => f.debug_tuple("Loaded").field(&"").finish(),
     62         }
     63     }
     64 }
     65 
     66 pub struct Animation {
     67     pub first_frame: TextureFrame,
     68     pub other_frames: Vec<TextureFrame>,
     69 }
     70 
     71 impl Animation {
     72     pub fn get_frame(&self, index: usize) -> Option<&TextureFrame> {
     73         if index == 0 {
     74             Some(&self.first_frame)
     75         } else {
     76             self.other_frames.get(index - 1)
     77         }
     78     }
     79 
     80     pub fn num_frames(&self) -> usize {
     81         self.other_frames.len() + 1
     82     }
     83 }
     84 
     85 pub struct TextureFrame {
     86     pub delay: Duration,
     87     pub texture: TextureHandle,
     88 }
     89 
     90 pub struct ImageFrame {
     91     pub delay: Duration,
     92     pub image: ColorImage,
     93 }
     94 
     95 pub struct MediaCache {
     96     pub cache_dir: path::PathBuf,
     97     pub cache_type: MediaCacheType,
     98     pub cache_size: Arc<Mutex<Option<u64>>>,
     99 }
    100 
    101 #[derive(Debug, Eq, PartialEq, Clone, Copy)]
    102 pub enum MediaCacheType {
    103     Image,
    104     Gif,
    105 }
    106 
    107 impl MediaCache {
    108     pub fn new(parent_dir: &Path, cache_type: MediaCacheType) -> Self {
    109         let cache_dir = parent_dir.join(Self::rel_dir(cache_type));
    110 
    111         let cache_dir_clone = cache_dir.clone();
    112         let cache_size = Arc::new(Mutex::new(None));
    113         let cache_size_clone = Arc::clone(&cache_size);
    114 
    115         thread::spawn(move || {
    116             let mut last_checked = Instant::now() - Duration::from_secs(999);
    117             loop {
    118                 // check cache folder size every 60 s
    119                 if last_checked.elapsed() >= Duration::from_secs(60) {
    120                     let size = compute_folder_size(&cache_dir_clone);
    121                     *cache_size_clone.lock().unwrap() = Some(size);
    122                     last_checked = Instant::now();
    123                 }
    124                 thread::sleep(Duration::from_secs(5));
    125             }
    126         });
    127 
    128         Self {
    129             cache_dir,
    130             cache_type,
    131             cache_size,
    132         }
    133     }
    134 
    135     pub fn rel_dir(cache_type: MediaCacheType) -> &'static str {
    136         match cache_type {
    137             MediaCacheType::Image => "img",
    138             MediaCacheType::Gif => "gif",
    139         }
    140     }
    141 
    142     pub fn write(cache_dir: &path::Path, url: &str, data: ColorImage) -> Result<()> {
    143         let file = Self::create_file(cache_dir, url)?;
    144         let encoder = image::codecs::webp::WebPEncoder::new_lossless(file);
    145 
    146         encoder.encode(
    147             data.as_raw(),
    148             data.size[0] as u32,
    149             data.size[1] as u32,
    150             image::ColorType::Rgba8.into(),
    151         )?;
    152 
    153         Ok(())
    154     }
    155 
    156     fn create_file(cache_dir: &path::Path, url: &str) -> Result<File> {
    157         let file_path = cache_dir.join(Self::key(url));
    158         if let Some(p) = file_path.parent() {
    159             create_dir_all(p)?;
    160         }
    161         Ok(File::options()
    162             .write(true)
    163             .create(true)
    164             .truncate(true)
    165             .open(file_path)?)
    166     }
    167 
    168     pub fn write_gif(cache_dir: &path::Path, url: &str, data: Vec<ImageFrame>) -> Result<()> {
    169         let file = Self::create_file(cache_dir, url)?;
    170 
    171         let mut encoder = image::codecs::gif::GifEncoder::new(file);
    172         for img in data {
    173             let buf = color_image_to_rgba(img.image);
    174             let frame = Frame::from_parts(buf, 0, 0, Delay::from_saturating_duration(img.delay));
    175             if let Err(e) = encoder.encode_frame(frame) {
    176                 tracing::error!("problem encoding frame: {e}");
    177             }
    178         }
    179 
    180         Ok(())
    181     }
    182 
    183     pub fn key(url: &str) -> String {
    184         let k: String = sha2::Sha256::digest(url.as_bytes()).encode_hex();
    185         PathBuf::from(&k[0..2])
    186             .join(&k[2..4])
    187             .join(k)
    188             .to_string_lossy()
    189             .to_string()
    190     }
    191 
    192     /// Migrate from base32 encoded url to sha256 url + sub-dir structure
    193     pub fn migrate_v0(&self) -> Result<()> {
    194         for file in std::fs::read_dir(&self.cache_dir)? {
    195             let file = if let Ok(f) = file {
    196                 f
    197             } else {
    198                 // not sure how this could fail, skip entry
    199                 continue;
    200             };
    201             if !file.path().is_file() {
    202                 continue;
    203             }
    204             let old_filename = file.file_name().to_string_lossy().to_string();
    205             let old_url = if let Some(u) =
    206                 base32::decode(base32::Alphabet::Crockford, &old_filename)
    207                     .and_then(|s| String::from_utf8(s).ok())
    208             {
    209                 u
    210             } else {
    211                 warn!("Invalid base32 filename: {}", &old_filename);
    212                 continue;
    213             };
    214             let new_path = self.cache_dir.join(Self::key(&old_url));
    215             if let Some(p) = new_path.parent() {
    216                 create_dir_all(p)?;
    217             }
    218 
    219             if let Err(e) = std::fs::rename(file.path(), &new_path) {
    220                 warn!(
    221                     "Failed to migrate file from {} to {}: {:?}",
    222                     file.path().display(),
    223                     new_path.display(),
    224                     e
    225                 );
    226             }
    227         }
    228 
    229         Ok(())
    230     }
    231 
    232     fn clear(&mut self) {
    233         *self.cache_size.try_lock().unwrap() = Some(0);
    234     }
    235 }
    236 
    237 fn color_image_to_rgba(color_image: ColorImage) -> image::RgbaImage {
    238     let width = color_image.width() as u32;
    239     let height = color_image.height() as u32;
    240 
    241     let rgba_pixels: Vec<u8> = color_image
    242         .pixels
    243         .iter()
    244         .flat_map(|color| color.to_array()) // Convert Color32 to `[u8; 4]`
    245         .collect();
    246 
    247     image::RgbaImage::from_raw(width, height, rgba_pixels)
    248         .expect("Failed to create RgbaImage from ColorImage")
    249 }
    250 
    251 fn compute_folder_size<P: AsRef<Path>>(path: P) -> u64 {
    252     fn walk(path: &Path) -> u64 {
    253         let mut size = 0;
    254         if let Ok(entries) = fs::read_dir(path) {
    255             for entry in entries.flatten() {
    256                 let path = entry.path();
    257                 if let Ok(metadata) = entry.metadata() {
    258                     if metadata.is_file() {
    259                         size += metadata.len();
    260                     } else if metadata.is_dir() {
    261                         size += walk(&path);
    262                     }
    263                 }
    264             }
    265         }
    266         size
    267     }
    268     walk(path.as_ref())
    269 }
    270 
    271 pub struct Images {
    272     pub base_path: path::PathBuf,
    273     pub static_imgs: MediaCache,
    274     pub gifs: MediaCache,
    275     pub textures: TexturesCache,
    276     pub urls: UrlMimes,
    277     /// cached imeta data
    278     pub metadata: HashMap<String, ImageMetadata>,
    279     pub gif_states: GifStateMap,
    280 }
    281 
    282 impl Images {
    283     /// path to directory to place [`MediaCache`]s
    284     pub fn new(path: path::PathBuf) -> Self {
    285         Self {
    286             base_path: path.clone(),
    287             static_imgs: MediaCache::new(&path, MediaCacheType::Image),
    288             gifs: MediaCache::new(&path, MediaCacheType::Gif),
    289             urls: UrlMimes::new(UrlCache::new(path.join(UrlCache::rel_dir()))),
    290             gif_states: Default::default(),
    291             metadata: Default::default(),
    292             textures: TexturesCache::new(path.clone()),
    293         }
    294     }
    295 
    296     pub fn migrate_v0(&self) -> Result<()> {
    297         self.static_imgs.migrate_v0()?;
    298         self.gifs.migrate_v0()
    299     }
    300 
    301     pub fn get_renderable_media(&mut self, url: &str) -> Option<RenderableMedia> {
    302         Self::find_renderable_media(&mut self.urls, &self.metadata, url)
    303     }
    304 
    305     pub fn find_renderable_media(
    306         urls: &mut UrlMimes,
    307         imeta: &HashMap<String, ImageMetadata>,
    308         url: &str,
    309     ) -> Option<RenderableMedia> {
    310         let media_type = crate::urls::supported_mime_hosted_at_url(urls, url)?;
    311 
    312         let obfuscation_type = match imeta.get(url) {
    313             Some(blur) => ObfuscationType::Blurhash(blur.clone()),
    314             None => ObfuscationType::Default,
    315         };
    316 
    317         Some(RenderableMedia {
    318             url: url.to_string(),
    319             media_type,
    320             obfuscation_type,
    321         })
    322     }
    323 
    324     pub fn latest_texture<'a>(
    325         &'a mut self,
    326         jobs: &MediaJobSender,
    327         ui: &mut egui::Ui,
    328         url: &str,
    329         img_type: ImageType,
    330         animation_mode: AnimationMode,
    331     ) -> Option<&'a TextureHandle> {
    332         let cache_type = crate::urls::supported_mime_hosted_at_url(&mut self.urls, url)?;
    333 
    334         let mut loader = NoLoadingLatestTex::new(
    335             &self.textures.static_image,
    336             &self.textures.animated,
    337             &mut self.gif_states,
    338         );
    339         loader.latest(jobs, ui.ctx(), url, cache_type, img_type, animation_mode)
    340     }
    341 
    342     pub fn get_cache(&self, cache_type: MediaCacheType) -> &MediaCache {
    343         match cache_type {
    344             MediaCacheType::Image => &self.static_imgs,
    345             MediaCacheType::Gif => &self.gifs,
    346         }
    347     }
    348 
    349     pub fn get_cache_mut(&mut self, cache_type: MediaCacheType) -> &mut MediaCache {
    350         match cache_type {
    351             MediaCacheType::Image => &mut self.static_imgs,
    352             MediaCacheType::Gif => &mut self.gifs,
    353         }
    354     }
    355 
    356     pub fn clear_folder_contents(&mut self) -> io::Result<()> {
    357         for entry in fs::read_dir(self.base_path.clone())? {
    358             let entry = entry?;
    359             let path = entry.path();
    360 
    361             if path.is_dir() {
    362                 fs::remove_dir_all(path)?;
    363             } else {
    364                 fs::remove_file(path)?;
    365             }
    366         }
    367 
    368         self.urls.cache.clear();
    369         self.static_imgs.clear();
    370         self.gifs.clear();
    371         self.gif_states.clear();
    372 
    373         Ok(())
    374     }
    375 
    376     pub fn trusted_texture_loader(&mut self) -> TrustedMediaLatestTex<'_> {
    377         TrustedMediaLatestTex::new(
    378             NoLoadingLatestTex::new(
    379                 &self.textures.static_image,
    380                 &self.textures.animated,
    381                 &mut self.gif_states,
    382             ),
    383             &self.textures.blurred,
    384         )
    385     }
    386 
    387     pub fn untrusted_texture_loader(&mut self) -> UntrustedMediaLatestTex<'_> {
    388         UntrustedMediaLatestTex::new(&self.textures.blurred)
    389     }
    390 
    391     pub fn no_img_loading_tex_loader(&'_ mut self) -> NoLoadingLatestTex<'_> {
    392         NoLoadingLatestTex::new(
    393             &self.textures.static_image,
    394             &self.textures.animated,
    395             &mut self.gif_states,
    396         )
    397     }
    398 
    399     pub fn user_trusts_img(&self, url: &str, media_type: MediaCacheType) -> bool {
    400         match media_type {
    401             MediaCacheType::Image => self.textures.static_image.contains(url),
    402             MediaCacheType::Gif => self.textures.animated.contains(url),
    403         }
    404     }
    405 }
    406 
    407 pub type GifStateMap = HashMap<String, GifState>;
    408 
    409 pub struct GifState {
    410     pub last_frame_rendered: Instant,
    411     pub last_frame_duration: Duration,
    412     pub next_frame_time: Option<SystemTime>,
    413     pub last_frame_index: usize,
    414 }
    415 
    416 pub struct LatestTexture {
    417     pub texture: TextureHandle,
    418     pub request_next_repaint: Option<SystemTime>,
    419 }