datatrash/src/download.rs

use actix_files::NamedFile;
use actix_web::{
    error,
    http::header::{Charset, ContentDisposition, DispositionParam, DispositionType, ExtendedValue},
    web, Error, HttpRequest, HttpResponse,
};
use async_std::{fs, path::Path};
use futures::TryStreamExt;
use mime::Mime;
use sqlx::{
    postgres::{PgPool, PgRow},
    Row,
};

use crate::config::Config;
use crate::deleter;

const VIEW_HTML: &str = include_str!("../template/view.html");

pub async fn download(
    req: HttpRequest,
    db: web::Data<PgPool>,
    config: web::Data<Config>,
) -> Result<HttpResponse, Error> {
    let id = req.match_info().query("id");
    let mut rows =
        sqlx::query("SELECT file_id, file_name, delete_on_download from files WHERE file_id = $1")
            .bind(id)
            .fetch(db.as_ref());
    let row: PgRow = rows
        .try_next()
        .await
        .map_err(|_| error::ErrorInternalServerError("could not run select statement"))?
        .ok_or_else(|| error::ErrorNotFound("file does not exist or has expired"))?;

    let file_id: String = row.get("file_id");
    let file_name: String = row.get("file_name");
    let delete_on_download: bool = row.get("delete_on_download");
    let mut path = config.files_dir.clone();
    path.push(&file_id);

    let download = req.query_string().contains("dl");
    let (content_type, mut content_disposition) = get_content_types(&path, &file_name);
    let response = if content_type.type_() == mime::TEXT && !download {
        let content = fs::read_to_string(path).await.map_err(|_| {
            error::ErrorInternalServerError("this file should be here but could not be found")
        })?;
        let encoded = htmlescape::encode_minimal(&content);
        let view_html = VIEW_HTML.replace("{text}", &encoded);
        let response = HttpResponse::Ok().content_type("text/html").body(view_html);
        Ok(response)
    } else {
        if download {
            content_disposition.disposition = DispositionType::Attachment;
        }
        let file = NamedFile::open(path)
            .map_err(|_| {
                error::ErrorInternalServerError("this file should be here but could not be found")
            })?
            .set_content_type(content_type)
            .set_content_disposition(content_disposition);
        file.into_response(&req)
    };
    if delete_on_download {
        deleter::delete_by_id(&db, &file_id, &config.files_dir)
            .await
            .map_err(|_| error::ErrorInternalServerError("could not delete file"))?;
    }
    response
}

fn get_content_types(path: &Path, filename: &str) -> (Mime, ContentDisposition) {
    let std_path = std::path::Path::new(path.as_os_str());
    let ct = tree_magic_mini::from_filepath(std_path)
        .unwrap_or("application/octet-stream")
        .parse::<Mime>()
        .expect("tree_magic_mini should not produce invalid mime");

    let disposition = match ct.type_() {
        mime::IMAGE | mime::TEXT | mime::AUDIO | mime::VIDEO => DispositionType::Inline,
        _ => DispositionType::Attachment,
    };

    let cd = ContentDisposition {
        disposition,
        parameters: get_disposition_params(filename),
    };

    (ct, cd)
}

fn get_disposition_params(filename: &str) -> Vec<DispositionParam> {
    let mut parameters = vec![DispositionParam::Filename(filename.to_owned())];
    if !filename.is_ascii() {
        parameters.push(DispositionParam::FilenameExt(ExtendedValue {
            charset: Charset::Ext(String::from("UTF-8")),
            language_tag: None,
            value: filename.to_owned().into_bytes(),
        }))
    }
    parameters
}