Line data    Source code

       1              : //! # Media files
       2              : 
       3              : use bytes::Bytes;
       4              : use rusqlite::{Connection, params};
       5              : use sha3::{Digest as _, Sha3_256};
       6              : use std::fs::write;
       7              : use std::path::Path;
       8              : 
       9              : /// Compute the message digest, using Sha256, of file
      10            2 : pub fn digest(data: &Bytes) -> Vec<u8> {
      11            2 :     let mut hasher = Sha3_256::new();
      12            2 :     hasher.update(data);
      13            2 :     let digest = hasher.finalize();
      14            2 :     digest.to_vec()
      15            2 : }
      16              : 
      17            2 : pub fn tohex(data: &Vec<u8>) -> String {
      18            2 :     let mut ret = String::new();
      19           66 :     for byte in data {
      20           64 :         ret.push_str(&format!("{:02X}", byte));
      21           64 :     }
      22            2 :     ret
      23            2 : }
      24              : 
      25              : // TODO: Consider a FileField struct we can convert from to perform validation
      26              : // and do the heavy lifting for error reporting.
      27              : 
      28              : // Should we receive the body as a stream?
      29              : // TODO: We want to avoid reading all the file into memory. We want to read the
      30              : // field by chunks instead.
      31              : // Cf. https://docs.rs/axum/latest/axum/extract/multipart/struct.Field.html
      32              : /// Store the file in SQLite
      33            2 : pub(crate) fn store_file<P: AsRef<Path>>(
      34            2 :     db: &Connection,
      35            2 :     upload_dir: P,
      36            2 :     bytes: Bytes,
      37            2 :     filename: String,
      38            2 :     content_type: String,
      39            2 : ) -> anyhow::Result<i64> {
      40            2 :     let digest = digest(&bytes);
      41            2 :     let digest_hex = tohex(&digest);
      42              :     // TODO: Infer file category from content_type
      43            2 :     let mut stmt = db.prepare(
      44            2 :         "INSERT INTO files (filename, digest, mimetype) VALUES (?, ?, ?) RETURNING file_id",
      45            0 :     )?;
      46            2 :     let file_id: i64 =
      47            2 :         stmt.query_row(params![filename, digest_hex, content_type], |row| row.get("file_id"))?;
      48            2 :     let data = bytes.to_vec();
      49            2 :     let path = upload_dir.as_ref().join(digest_hex);
      50            2 :     write(path, data)?;
      51              :     // TODO: Use a tx. Commit here. so that we only record the upload if the
      52              :     // file is saved on disk.
      53            2 :     Ok(file_id)
      54            2 : }
      55              : 
      56              : pub(crate) struct MediaFile {
      57              :     pub digest: String,
      58              :     pub content_type: String,
      59              : }
      60              : 
      61            1 : pub(crate) fn find_by_id(db: &Connection, file_id: i64) -> anyhow::Result<MediaFile> {
      62            1 :     let mut stmt = db.prepare("SELECT digest, mimetype FROM files where file_id = ?")?;
      63            1 :     let (digest, content_type): (String, String) =
      64            1 :         stmt.query_row([file_id], |row| Ok((row.get("digest")?, row.get("mimetype")?)))?;
      65              : 
      66            1 :     Ok(MediaFile { content_type, digest })
      67            1 : }
      68              : 
      69              : #[cfg(test)]
      70              : mod tests {
      71              :     use super::*;
      72              :     use crate::db::setup_db;
      73              : 
      74              :     #[test]
      75            1 :     fn test_store_file() -> anyhow::Result<()> {
      76            1 :         let db = setup_db(None);
      77            1 :         let buf = std::fs::read("public/favicon.ico")?;
      78            1 :         let file_id = store_file(
      79            1 :             &db,
      80              :             "/tmp/",
      81            1 :             buf.into(),
      82            1 :             "favicon.ico".to_string(),
      83            1 :             "image/x-icon".to_string(),
      84            0 :         )?;
      85              : 
      86            1 :         let mut stmt = db.prepare("SELECT file_id, digest FROM files WHERE file_id = ?")?;
      87            1 :         let digest: String = stmt.query_row([file_id], |row| row.get("digest"))?;
      88            1 :         assert_eq!(
      89              :             digest,
      90            1 :             "C9D25FC03F47AA1062C344781A66E680D558F0563CFB5522FF57863F027637BA".to_string()
      91              :         );
      92            1 :         let got =
      93            1 :             std::fs::read("/tmp/C9D25FC03F47AA1062C344781A66E680D558F0563CFB5522FF57863F027637BA")?;
      94            1 :         let want = std::fs::read("public/favicon.ico")?;
      95            1 :         assert_eq!(want, got);
      96            1 :         Ok(())
      97            1 :     }
      98              : }