revolt_files/
lib.rs

1use std::io::{BufRead, Read, Seek, Write};
2
3use aes_gcm::{
4    aead::{AeadCore, AeadMutInPlace, OsRng},
5    Aes256Gcm, Key, KeyInit, Nonce,
6};
7use image::{DynamicImage, ImageBuffer};
8use revolt_config::{config, report_internal_error, FilesS3};
9use revolt_result::{create_error, Result};
10
11use aws_sdk_s3::{
12    config::{Credentials, Region},
13    Client, Config,
14};
15
16use base64::prelude::*;
17use tempfile::NamedTempFile;
18use tiny_skia::Pixmap;
19
20/// Size of the authentication tag in the buffer
21pub const AUTHENTICATION_TAG_SIZE_BYTES: usize = 16;
22
23/// Create an S3 client
24pub fn create_client(s3_config: FilesS3) -> Client {
25    let provider_name = "my-creds";
26    let creds = Credentials::new(
27        s3_config.access_key_id,
28        s3_config.secret_access_key,
29        None,
30        None,
31        provider_name,
32    );
33
34    let config = Config::builder()
35        .region(Region::new(s3_config.region))
36        .endpoint_url(s3_config.endpoint)
37        .force_path_style(s3_config.path_style_buckets)
38        .credentials_provider(creds)
39        .build();
40
41    Client::from_conf(config)
42}
43
44/// Create an AES-256-GCM cipher
45pub fn create_cipher(key: &str) -> Aes256Gcm {
46    let key = &BASE64_STANDARD.decode(key).expect("valid base64 string")[..];
47    let key: &Key<Aes256Gcm> = key.into();
48    Aes256Gcm::new(key)
49}
50
51/// Fetch a file from S3 (and decrypt it)
52pub async fn fetch_from_s3(bucket_id: &str, path: &str, nonce: &str) -> Result<Vec<u8>> {
53    let config = config().await;
54    let client = create_client(config.files.s3);
55
56    // Send a request for the file
57    let mut obj =
58        report_internal_error!(client.get_object().bucket(bucket_id).key(path).send().await)?;
59
60    // Read the file from remote
61    let mut buf = vec![];
62    while let Some(bytes) = obj.body.next().await {
63        let data = report_internal_error!(bytes)?;
64        report_internal_error!(buf.write_all(&data))?;
65        // is there a more efficient way to do this?
66        // we just want the Vec<u8>
67    }
68
69    // File is not encrypted
70    if nonce.is_empty() {
71        return Ok(buf);
72    }
73
74    // Recover nonce as bytes
75    let nonce = &BASE64_STANDARD.decode(nonce).unwrap()[..];
76    let nonce: &Nonce<typenum::consts::U12> = nonce.into();
77
78    // Decrypt the file
79    create_cipher(&config.files.encryption_key)
80        .decrypt_in_place(nonce, b"", &mut buf)
81        .map_err(|_| create_error!(InternalError))?;
82
83    Ok(buf)
84}
85
86/// Encrypt and upload a file to S3 (returning its nonce/IV)
87pub async fn upload_to_s3(bucket_id: &str, path: &str, buf: &[u8]) -> Result<String> {
88    let config = config().await;
89    let client = create_client(config.files.s3);
90
91    // Generate a nonce
92    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
93
94    // Extend the buffer for in-place encryption
95    let mut buf = [buf, &[0; AUTHENTICATION_TAG_SIZE_BYTES]].concat();
96
97    // Encrypt the file in place
98    create_cipher(&config.files.encryption_key)
99        .encrypt_in_place(&nonce, b"", &mut buf)
100        .map_err(|_| create_error!(InternalError))?;
101
102    // Upload the file to remote
103    report_internal_error!(
104        client
105            .put_object()
106            .bucket(bucket_id)
107            .key(path)
108            .body(buf.into())
109            .send()
110            .await
111    )?;
112
113    Ok(BASE64_STANDARD.encode(nonce))
114}
115
116/// Delete a file from S3 by path
117pub async fn delete_from_s3(bucket_id: &str, path: &str) -> Result<()> {
118    let config = config().await;
119    let client = create_client(config.files.s3);
120
121    report_internal_error!(
122        client
123            .delete_object()
124            .bucket(bucket_id)
125            .key(path)
126            .send()
127            .await
128    )?;
129
130    Ok(())
131}
132
133/// Determine size of image at temp file
134pub fn image_size(f: &NamedTempFile) -> Option<(usize, usize)> {
135    if let Ok(size) = imagesize::size(f.path())
136        .inspect_err(|err| tracing::error!("Failed to generate image size! {err:?}"))
137    {
138        Some((size.width, size.height))
139    } else {
140        None
141    }
142}
143
144/// Determine size of image with buffer
145pub fn image_size_vec(v: &[u8], mime: &str) -> Option<(usize, usize)> {
146    match mime {
147        "image/svg+xml" => {
148            let tree =
149                report_internal_error!(usvg::Tree::from_data(v, &Default::default())).ok()?;
150
151            let size = tree.size();
152            Some((size.width() as usize, size.height() as usize))
153        }
154        _ => {
155            if let Ok(size) = imagesize::blob_size(v)
156                .inspect_err(|err| tracing::error!("Failed to generate image size! {err:?}"))
157            {
158                Some((size.width, size.height))
159            } else {
160                None
161            }
162        }
163    }
164}
165
166/// Determine size of video at temp file
167pub fn video_size(f: &NamedTempFile) -> Option<(i64, i64)> {
168    if let Ok(data) = ffprobe::ffprobe(f.path())
169        .inspect_err(|err| tracing::error!("Failed to ffprobe file! {err:?}"))
170    {
171        // Use first valid stream
172        for stream in data.streams {
173            if let (Some(w), Some(h)) = (stream.width, stream.height) {
174                return Some((w, h));
175            }
176        }
177
178        None
179    } else {
180        None
181    }
182}
183
184/// Decode image from reader
185pub fn decode_image<R: Read + BufRead + Seek>(reader: &mut R, mime: &str) -> Result<DynamicImage> {
186    match mime {
187        // Read image using jxl-oxide crate
188        "image/jxl" => {
189            let jxl_image = report_internal_error!(jxl_oxide::JxlImage::builder().read(reader))?;
190            if let Ok(frame) = jxl_image.render_frame(0) {
191                match frame.color_channels().len() {
192                    3 => Ok(DynamicImage::ImageRgb8(
193                        DynamicImage::ImageRgb32F(
194                            ImageBuffer::from_vec(
195                                jxl_image.width(),
196                                jxl_image.height(),
197                                frame.image().buf().to_vec(),
198                            )
199                            .ok_or_else(|| create_error!(ImageProcessingFailed))?,
200                        )
201                        .to_rgb8(),
202                    )),
203                    4 => Ok(DynamicImage::ImageRgba8(
204                        DynamicImage::ImageRgba32F(
205                            ImageBuffer::from_vec(
206                                jxl_image.width(),
207                                jxl_image.height(),
208                                frame.image().buf().to_vec(),
209                            )
210                            .ok_or_else(|| create_error!(ImageProcessingFailed))?,
211                        )
212                        .to_rgba8(),
213                    )),
214                    _ => Err(create_error!(ImageProcessingFailed)),
215                }
216            } else {
217                Err(create_error!(ImageProcessingFailed))
218            }
219        }
220        // Read image using resvg
221        "image/svg+xml" => {
222            // usvg doesn't support Read trait so copy to buffer
223            let mut buf = Vec::new();
224            report_internal_error!(reader.read_to_end(&mut buf))?;
225
226            let tree = report_internal_error!(usvg::Tree::from_data(&buf, &Default::default()))?;
227            let size = tree.size();
228            let mut pixmap = Pixmap::new(size.width() as u32, size.height() as u32)
229                .ok_or_else(|| create_error!(ImageProcessingFailed))?;
230
231            let mut pixmap_mut = pixmap.as_mut();
232            resvg::render(&tree, Default::default(), &mut pixmap_mut);
233
234            Ok(DynamicImage::ImageRgba8(
235                ImageBuffer::from_vec(
236                    size.width() as u32,
237                    size.height() as u32,
238                    pixmap.data().to_vec(),
239                )
240                .ok_or_else(|| create_error!(ImageProcessingFailed))?,
241            ))
242        }
243        // Check if we can read using image-rs crate
244        _ => report_internal_error!(report_internal_error!(
245            image::ImageReader::new(reader).with_guessed_format()
246        )?
247        .decode()),
248    }
249}
250
251/// Check whether given reader has a valid image
252pub fn is_valid_image<R: Read + BufRead + Seek>(reader: &mut R, mime: &str) -> bool {
253    match mime {
254        // Check if we can read using jxl-oxide crate
255        "image/jxl" => jxl_oxide::JxlImage::builder()
256            .read(reader)
257            .inspect_err(|err| tracing::error!("Failed to read JXL! {err:?}"))
258            .is_ok(),
259        // Check if we can read using image-rs crate
260        _ => !matches!(
261            image::ImageReader::new(reader)
262                .with_guessed_format()
263                .inspect_err(|err| tracing::error!("Failed to read image! {err:?}"))
264                .map(|f| f.decode()),
265            Err(_) | Ok(Err(_))
266        ),
267    }
268}
269
270/// Create thumbnail from given image
271pub async fn create_thumbnail(image: DynamicImage, tag: &str) -> Vec<u8> {
272    // Load configuration
273    let config = config().await;
274    let [w, h] = config.files.preview.get(tag).unwrap();
275
276    // Create thumbnail
277    //.resize(width as u32, height as u32, image::imageops::FilterType::Gaussian)
278    // resize is about 2.5x slower,
279    // thumbnail doesn't have terrible quality
280    // so we use thumbnail
281    let image = image.thumbnail(image.width().min(*w as u32), image.height().min(*h as u32));
282
283    // Encode it into WEBP
284    let encoder = webp::Encoder::from_image(&image).expect("Could not create encoder.");
285    if config.files.webp_quality != 100.0 {
286        encoder.encode(config.files.webp_quality).to_vec()
287    } else {
288        encoder.encode_lossless().to_vec()
289    }
290}
revolt_files/lib.rs

revolt_files/
lib.rs
