1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114

/*!
# File Center Download Response on MongoDB for Rocket Framework

This crate provides a response struct used for client downloading from the File Center on MongoDB.

See `examples`.
*/

extern crate tokio_util;

pub extern crate mongo_file_center;

extern crate rocket;
extern crate url_escape;

use std::io::Cursor;

use tokio_util::io::StreamReader;

use mongo_file_center::bson::oid::ObjectId;
use mongo_file_center::{FileCenter, FileCenterError, FileData, FileItem};

use rocket::request::Request;
use rocket::response::{self, Responder, Response};

/// The response struct used for responding raw data from the File Center on MongoDB with **Etag** cache.
#[derive(Debug)]
pub struct FileCenterDownloadResponse {
    file_name: Option<String>,
    file_item: FileItem,
}

impl FileCenterDownloadResponse {
    /// Create a `FileCenterDownloadResponse` instance from a file item.
    #[inline]
    pub fn from_file_item<S: Into<String>>(
        file_item: FileItem,
        file_name: Option<S>,
    ) -> FileCenterDownloadResponse {
        let file_name = file_name.map(|file_name| file_name.into());

        FileCenterDownloadResponse {
            file_name,
            file_item,
        }
    }

    /// Create a `FileCenterDownloadResponse` instance from the object ID.
    pub async fn from_object_id<S: Into<String>>(
        file_center: &FileCenter,
        id: ObjectId,
        file_name: Option<S>,
    ) -> Result<Option<FileCenterDownloadResponse>, FileCenterError> {
        let file_item = file_center.get_file_item_by_id(id).await?;

        match file_item {
            Some(file_item) => Ok(Some(Self::from_file_item(file_item, file_name))),
            None => Ok(None),
        }
    }

    /// Create a `FileCenterDownloadResponse` instance from an ID token.
    #[inline]
    pub async fn from_id_token<I: AsRef<str> + Into<String>, S: Into<String>>(
        file_center: &FileCenter,
        id_token: I,
        file_name: Option<S>,
    ) -> Result<Option<FileCenterDownloadResponse>, FileCenterError> {
        let id = file_center.decrypt_id_token(id_token.as_ref())?;

        Self::from_object_id(file_center, id, file_name).await
    }

    #[inline]
    /// Check if the file item is temporary.
    pub fn is_temporary(&self) -> bool {
        self.file_item.get_expiration_time().is_some()
    }
}

impl<'r, 'o: 'r> Responder<'r, 'o> for FileCenterDownloadResponse {
    fn respond_to(self, _: &'r Request<'_>) -> response::Result<'o> {
        let mut response = Response::build();

        let file_name = self.file_name.as_deref().unwrap_or_else(|| self.file_item.get_file_name());

        if file_name.is_empty() {
            response.raw_header("Content-Disposition", "attachment");
        } else {
            let mut v = String::from("attachment; filename*=UTF-8''");

            url_escape::encode_component_to_string(file_name, &mut v);

            response.raw_header("Content-Disposition", v);
        }

        response.raw_header("Content-Type", self.file_item.get_mime_type().to_string());

        let file_size = self.file_item.get_file_size();

        match self.file_item.into_file_data() {
            FileData::Buffer(v) => {
                response.sized_body(v.len(), Cursor::new(v));
            }
            FileData::Stream(g) => {
                response.raw_header("Content-Length", file_size.to_string());

                response.streamed_body(StreamReader::new(g));
            }
        }

        response.ok()
    }
}