datatrash/src/download.rs

use std::str::FromStr;

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 mime::Mime;
use sqlx::postgres::PgPool;
use url::Url;

use crate::deleter;
use crate::{config::Config, file_kind::FileKind};

const TEXT_VIEW_HTML: &str = include_str!("../template/text-view.html");
const URL_VIEW_HTML: &str = include_str!("../template/url-view.html");

const TEXT_VIEW_SIZE_LIMIT: u64 = 512 * 1024; // 512KiB

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 (file_id, file_name, file_kind, delete_on_download) = load_file_info(id, &db).await?;
    let mut path = config.files_dir.clone();
    path.push(&file_id);

    let download = delete_on_download || req.query_string().contains("dl");
    let content_type = get_content_type(&path);
    let response = if use_text_view(&file_kind, &content_type, &path, download).await {
        build_text_response(&path).await
    } else {
        build_file_response(download, &file_name, path, content_type, req)
    };
    if delete_on_download {
        deleter::delete_by_id(&db, &file_id, &config.files_dir)
            .await
            .map_err(|db_err| {
                log::error!("could not delete file {:?}", db_err);
                error::ErrorInternalServerError("could not delete file")
            })?;
    }
    response
}

async fn load_file_info(
    id: &str,
    db: &web::Data<sqlx::Pool<sqlx::Postgres>>,
) -> Result<(String, String, String, bool), Error> {
    sqlx::query_as(
        "SELECT file_id, file_name, kind, delete_on_download from files WHERE file_id = $1",
    )
    .bind(id)
    .fetch_optional(db.as_ref())
    .await
    .map_err(|db_err| {
        log::error!("could not run select statement {:?}", db_err);
        error::ErrorInternalServerError("could not run select statement")
    })?
    .ok_or_else(|| error::ErrorNotFound("file does not exist or has expired"))
}

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

async fn use_text_view(
    file_kind: &str,
    content_type: &Mime,
    file_path: &Path,
    download: bool,
) -> bool {
    let is_text =
        FileKind::from_str(file_kind) == Ok(FileKind::Text) || content_type.type_() == mime::TEXT;
    let is_not_large = get_file_size(file_path).await < TEXT_VIEW_SIZE_LIMIT;
    is_text && is_not_large && !download
}

async fn get_file_size(file_path: &Path) -> u64 {
    fs::metadata(file_path)
        .await
        .map(|metadata| metadata.len())
        .unwrap_or(0)
}

async fn build_text_response(path: &Path) -> Result<HttpResponse, Error> {
    let content = fs::read_to_string(path).await.map_err(|file_err| {
        log::error!("file could not be read {:?}", file_err);
        error::ErrorInternalServerError("this file should be here but could not be found")
    })?;
    let encoded = htmlescape::encode_minimal(&content);
    let html = if !content.contains(&['\n', '\r'][..]) && Url::from_str(&content).is_ok() {
        let attribute_encoded = htmlescape::encode_attribute(&content);
        URL_VIEW_HTML
            .replace("{link_content}", &encoded)
            .replace("{link_attribute}", &attribute_encoded)
    } else {
        TEXT_VIEW_HTML.replace("{text}", &encoded)
    };
    let response = HttpResponse::Ok().content_type("text/html").body(html);
    Ok(response)
}

fn build_file_response(
    download: bool,
    file_name: &str,
    path: async_std::path::PathBuf,
    content_type: Mime,
    req: HttpRequest,
) -> Result<HttpResponse, Error> {
    let content_disposition = ContentDisposition {
        disposition: if download {
            DispositionType::Attachment
        } else {
            DispositionType::Inline
        },
        parameters: get_disposition_params(file_name),
    };
    let file = NamedFile::open(path)
        .map_err(|file_err| {
            log::error!("file could not be read {:?}", file_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)
}

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
}