use super::{Headers, IntoResponse};
use crate::headers::ContentType;
use crate::{Body, Response, StatusCode};
use rama_core::bytes::Bytes;
use rama_core::error::BoxError;
use rama_core::futures::TryStream;
use rama_core::stream::io::ReaderStream;
use rama_http_headers::{ContentDisposition, ContentLength, ContentRange, Error, HeaderMapExt};
use std::ops::RangeBounds;
use std::path::Path;
use tokio::fs::File;
#[derive(Debug, Clone)]
pub struct OctetStream<S> {
stream: S,
filename: Option<String>,
content_size: Option<u64>,
}
impl<S> OctetStream<S> {
pub fn new(stream: S) -> Self {
Self {
stream,
filename: None,
content_size: None,
}
}
rama_utils::macros::generate_set_and_with! {
pub fn file_name(mut self, filename: String) -> Self {
self.filename = Some(filename);
self
}
}
rama_utils::macros::generate_set_and_with! {
pub fn content_size(mut self, content_size: u64) -> Self {
self.content_size = Some(content_size);
self
}
}
pub fn try_into_range_response(self, range: impl RangeBounds<u64>) -> Result<Response, Error>
where
S: TryStream<Ok: Into<Bytes>, Error: Into<BoxError>> + Send + 'static,
{
let body = Body::from_stream(self.stream);
let content_range =
ContentRange::bytes(range, self.content_size).map_err(|_e| Error::invalid())?;
let mut response = (
StatusCode::PARTIAL_CONTENT,
Headers((ContentType::octet_stream(), content_range)),
body,
)
.into_response();
if let Some(filename) = self.filename {
response
.headers_mut()
.typed_insert(ContentDisposition::attachment(filename.as_str()));
}
Ok(response)
}
async fn open_file_with_metadata(
path: &Path,
) -> std::io::Result<(File, Option<u64>, Option<String>)> {
let file = rama_utils::fs::safe_open(path).await?;
let metadata = file.metadata().await.ok();
let content_size = metadata.as_ref().map(|m| m.len());
let filename = path
.file_name()
.and_then(|name| name.to_str())
.map(|s| s.to_owned());
Ok((file, content_size, filename))
}
pub async fn try_from_path(
path: impl AsRef<Path>,
) -> std::io::Result<OctetStream<ReaderStream<File>>> {
let (file, content_size, filename) = Self::open_file_with_metadata(path.as_ref()).await?;
let stream = ReaderStream::new(file);
Ok(OctetStream {
stream,
filename,
content_size,
})
}
pub async fn try_range_response_from_path(
path: impl AsRef<Path>,
start: u64,
end: u64,
) -> std::io::Result<Response> {
use tokio::io::{AsyncReadExt, AsyncSeekExt};
let (mut file, content_size, filename) =
Self::open_file_with_metadata(path.as_ref()).await?;
file.seek(std::io::SeekFrom::Start(start)).await?;
let stream = ReaderStream::new(file.take(end - start));
let octet_stream = OctetStream {
stream,
filename,
content_size,
};
octet_stream
.try_into_range_response(start..end)
.map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
}
}
impl<S> IntoResponse for OctetStream<S>
where
S: TryStream<Ok: Into<Bytes>, Error: Into<BoxError>> + Send + 'static,
{
fn into_response(self) -> Response {
let body = Body::from_stream(self.stream);
let mut response = (Headers::single(ContentType::octet_stream()), body).into_response();
if let Some(filename) = self.filename {
response
.headers_mut()
.typed_insert(ContentDisposition::attachment(filename.as_str()));
}
if let Some(size) = self.content_size {
response.headers_mut().typed_insert(ContentLength(size));
}
response
}
}
impl<S> From<S> for OctetStream<S> {
fn from(stream: S) -> Self {
Self::new(stream)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_octet_stream() {
let data = vec![1u8, 2, 3, 4, 5];
let stream = OctetStream::new(data.clone());
assert_eq!(stream.stream, data);
assert_eq!(stream.filename, None);
assert_eq!(stream.content_size, None);
}
#[test]
fn test_with_file_name() {
let data = vec![1u8, 2, 3];
let stream = OctetStream::new(data).with_file_name("test.bin".to_owned());
assert_eq!(stream.filename, Some("test.bin".to_owned()));
}
#[test]
fn test_with_content_size() {
let data = vec![1u8, 2, 3];
let stream = OctetStream::new(data).with_content_size(1024);
assert_eq!(stream.content_size, Some(1024));
}
#[test]
fn test_chained_setters() {
let data = vec![1u8, 2, 3];
let stream = OctetStream::new(data)
.with_file_name("test.bin".to_owned())
.with_content_size(1024);
assert_eq!(stream.filename, Some("test.bin".to_owned()));
assert_eq!(stream.content_size, Some(1024));
}
#[tokio::test]
async fn test_into_response() {
use crate::header::{CONTENT_DISPOSITION, CONTENT_LENGTH, CONTENT_TYPE};
let cursor = std::io::Cursor::new(b"hello");
let stream = ReaderStream::new(cursor);
let octet_stream = OctetStream::new(stream)
.with_file_name("test.bin".to_owned())
.with_content_size(5);
let response = octet_stream.into_response();
assert_eq!(response.status(), StatusCode::OK);
assert_eq!(
response.headers().get(CONTENT_TYPE).unwrap(),
"application/octet-stream"
);
assert_eq!(response.headers().get(CONTENT_LENGTH).unwrap(), "5");
assert_eq!(
response.headers().get(CONTENT_DISPOSITION).unwrap(),
"attachment; filename=test.bin"
);
}
#[tokio::test]
async fn test_try_into_range_response() {
use crate::header::{CONTENT_DISPOSITION, CONTENT_RANGE, CONTENT_TYPE};
let cursor = std::io::Cursor::new(b"hello");
let stream = ReaderStream::new(cursor);
let octet_stream = OctetStream::new(stream)
.with_file_name("test.bin".to_owned())
.with_content_size(13);
let response = octet_stream.try_into_range_response(0..5).unwrap();
assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
assert_eq!(
response.headers().get(CONTENT_TYPE).unwrap(),
"application/octet-stream"
);
assert_eq!(
response.headers().get(CONTENT_RANGE).unwrap(),
"bytes 0-4/13"
);
assert_eq!(
response.headers().get(CONTENT_DISPOSITION).unwrap(),
"attachment; filename=test.bin"
);
}
#[tokio::test]
async fn test_try_from_path() {
let file_path = std::fs::canonicalize("../test-files/hello.txt").unwrap();
let stream = OctetStream::<ReaderStream<File>>::try_from_path(&file_path)
.await
.unwrap();
assert_eq!(stream.filename, Some("hello.txt".to_owned()));
#[cfg(target_os = "windows")]
assert_eq!(stream.content_size, Some(15)); #[cfg(not(target_os = "windows"))]
assert_eq!(stream.content_size, Some(14)); }
#[tokio::test]
async fn test_try_range_response_from_path() {
use crate::header::{CONTENT_DISPOSITION, CONTENT_RANGE};
let file_path = std::fs::canonicalize("../test-files/hello.txt").unwrap();
let response =
OctetStream::<ReaderStream<File>>::try_range_response_from_path(&file_path, 0, 5)
.await
.unwrap();
assert_eq!(response.status(), StatusCode::PARTIAL_CONTENT);
let value = response.headers().get(CONTENT_RANGE).unwrap();
#[cfg(target_os = "windows")]
assert_eq!(value, "bytes 0-4/15"); #[cfg(not(target_os = "windows"))]
assert_eq!(value, "bytes 0-4/14");
assert_eq!(
response.headers().get(CONTENT_DISPOSITION).unwrap(),
"attachment; filename=hello.txt"
);
}
}