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    // Remove the authentication tag bytes that were added during encryption
84    buf.truncate(buf.len() - AUTHENTICATION_TAG_SIZE_BYTES);
85
86    Ok(buf)
87}
88
89/// Encrypt and upload a file to S3 (returning its nonce/IV)
90pub async fn upload_to_s3(bucket_id: &str, path: &str, buf: &[u8]) -> Result<String> {
91    let config = config().await;
92    let client = create_client(config.files.s3);
93
94    // Generate a nonce
95    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
96
97    // Extend the buffer for in-place encryption
98    let mut buf = [buf, &[0; AUTHENTICATION_TAG_SIZE_BYTES]].concat();
99
100    // Encrypt the file in place
101    create_cipher(&config.files.encryption_key)
102        .encrypt_in_place(&nonce, b"", &mut buf)
103        .map_err(|_| create_error!(InternalError))?;
104
105    // Upload the file to remote
106    report_internal_error!(
107        client
108            .put_object()
109            .bucket(bucket_id)
110            .key(path)
111            .body(buf.into())
112            .send()
113            .await
114    )?;
115
116    Ok(BASE64_STANDARD.encode(nonce))
117}
118
119/// Delete a file from S3 by path
120pub async fn delete_from_s3(bucket_id: &str, path: &str) -> Result<()> {
121    let config = config().await;
122    let client = create_client(config.files.s3);
123
124    report_internal_error!(
125        client
126            .delete_object()
127            .bucket(bucket_id)
128            .key(path)
129            .send()
130            .await
131    )?;
132
133    Ok(())
134}
135
136/// Determine size of image at temp file
137pub fn image_size(f: &NamedTempFile) -> Option<(usize, usize)> {
138    if let Ok(size) = imagesize::size(f.path())
139        .inspect_err(|err| tracing::error!("Failed to generate image size! {err:?}"))
140    {
141        Some((size.width, size.height))
142    } else {
143        None
144    }
145}
146
147/// Determine size of image with buffer
148pub fn image_size_vec(v: &[u8], mime: &str) -> Option<(usize, usize)> {
149    match mime {
150        "image/svg+xml" => {
151            let tree =
152                report_internal_error!(usvg::Tree::from_data(v, &Default::default())).ok()?;
153
154            let size = tree.size();
155            Some((size.width() as usize, size.height() as usize))
156        }
157        _ => {
158            if let Ok(size) = imagesize::blob_size(v)
159                .inspect_err(|err| tracing::error!("Failed to generate image size! {err:?}"))
160            {
161                Some((size.width, size.height))
162            } else {
163                None
164            }
165        }
166    }
167}
168
169/// Determine size of video at temp file
170pub fn video_size(f: &NamedTempFile) -> Option<(i64, i64)> {
171    if let Ok(data) = ffprobe::ffprobe(f.path())
172        .inspect_err(|err| tracing::error!("Failed to ffprobe file! {err:?}"))
173    {
174        // Use first valid stream
175        for stream in data.streams {
176            if let (Some(w), Some(h)) = (stream.width, stream.height) {
177                return Some((w, h));
178            }
179        }
180
181        None
182    } else {
183        None
184    }
185}
186
187/// Decode image from reader
188pub fn decode_image<R: Read + BufRead + Seek>(reader: &mut R, mime: &str) -> Result<DynamicImage> {
189    match mime {
190        // Read image using jxl-oxide crate
191        "image/jxl" => {
192            let jxl_image = report_internal_error!(jxl_oxide::JxlImage::builder().read(reader))?;
193            if let Ok(frame) = jxl_image.render_frame(0) {
194                match frame.color_channels().len() {
195                    3 => Ok(DynamicImage::ImageRgb8(
196                        DynamicImage::ImageRgb32F(
197                            ImageBuffer::from_vec(
198                                jxl_image.width(),
199                                jxl_image.height(),
200                                frame.image().buf().to_vec(),
201                            )
202                            .ok_or_else(|| create_error!(ImageProcessingFailed))?,
203                        )
204                        .to_rgb8(),
205                    )),
206                    4 => Ok(DynamicImage::ImageRgba8(
207                        DynamicImage::ImageRgba32F(
208                            ImageBuffer::from_vec(
209                                jxl_image.width(),
210                                jxl_image.height(),
211                                frame.image().buf().to_vec(),
212                            )
213                            .ok_or_else(|| create_error!(ImageProcessingFailed))?,
214                        )
215                        .to_rgba8(),
216                    )),
217                    _ => Err(create_error!(ImageProcessingFailed)),
218                }
219            } else {
220                Err(create_error!(ImageProcessingFailed))
221            }
222        }
223        // Read image using resvg
224        "image/svg+xml" => {
225            // usvg doesn't support Read trait so copy to buffer
226            let mut buf = Vec::new();
227            report_internal_error!(reader.read_to_end(&mut buf))?;
228
229            let tree = report_internal_error!(usvg::Tree::from_data(&buf, &Default::default()))?;
230            let size = tree.size();
231            let mut pixmap = Pixmap::new(size.width() as u32, size.height() as u32)
232                .ok_or_else(|| create_error!(ImageProcessingFailed))?;
233
234            let mut pixmap_mut = pixmap.as_mut();
235            resvg::render(&tree, Default::default(), &mut pixmap_mut);
236
237            Ok(DynamicImage::ImageRgba8(
238                ImageBuffer::from_vec(
239                    size.width() as u32,
240                    size.height() as u32,
241                    pixmap.data().to_vec(),
242                )
243                .ok_or_else(|| create_error!(ImageProcessingFailed))?,
244            ))
245        }
246        // Check if we can read using image-rs crate
247        _ => report_internal_error!(report_internal_error!(
248            image::ImageReader::new(reader).with_guessed_format()
249        )?
250        .decode()),
251    }
252}
253
254/// Check whether given reader has a valid image
255pub fn is_valid_image<R: Read + BufRead + Seek>(reader: &mut R, mime: &str) -> bool {
256    match mime {
257        // Check if we can read using jxl-oxide crate
258        "image/jxl" => jxl_oxide::JxlImage::builder()
259            .read(reader)
260            .inspect_err(|err| tracing::error!("Failed to read JXL! {err:?}"))
261            .is_ok(),
262        // Check if we can read using image-rs crate
263        _ => !matches!(
264            image::ImageReader::new(reader)
265                .with_guessed_format()
266                .inspect_err(|err| tracing::error!("Failed to read image! {err:?}"))
267                .map(|f| f.decode()),
268            Err(_) | Ok(Err(_))
269        ),
270    }
271}
272
273/// Create thumbnail from given image
274pub async fn create_thumbnail(image: DynamicImage, tag: &str) -> Vec<u8> {
275    // Load configuration
276    let config = config().await;
277    let [w, h] = config.files.preview.get(tag).unwrap();
278
279    // Create thumbnail
280    //.resize(width as u32, height as u32, image::imageops::FilterType::Gaussian)
281    // resize is about 2.5x slower,
282    // thumbnail doesn't have terrible quality
283    // so we use thumbnail
284    let image = image.thumbnail(image.width().min(*w as u32), image.height().min(*h as u32));
285
286    // Encode it into WEBP
287    let encoder = webp::Encoder::from_image(&image).expect("Could not create encoder.");
288    if config.files.webp_quality != 100.0 {
289        encoder.encode(config.files.webp_quality).to_vec()
290    } else {
291        encoder.encode_lossless().to_vec()
292    }
293}
revolt_files/lib.rs

revolt_files/
lib.rs
