use axum::body::Body;
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use futures::{Stream, StreamExt};
use std::pin::Pin;
use crate::error::RariError;
pub struct StreamingHtmlResponse {
stream: Pin<Box<dyn Stream<Item = Result<Vec<u8>, RariError>> + Send>>,
status_code: StatusCode,
}
impl StreamingHtmlResponse {
pub fn new<S>(stream: S) -> Self
where
S: Stream<Item = Result<Vec<u8>, RariError>> + Send + 'static,
{
Self { stream: Box::pin(stream), status_code: StatusCode::OK }
}
pub fn with_status<S>(stream: S, status_code: StatusCode) -> Self
where
S: Stream<Item = Result<Vec<u8>, RariError>> + Send + 'static,
{
Self { stream: Box::pin(stream), status_code }
}
}
impl IntoResponse for StreamingHtmlResponse {
fn into_response(self) -> Response {
let client_connected = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(true));
let client_connected_clone = client_connected.clone();
let stream = self.stream.map(move |chunk| {
if !client_connected_clone.load(std::sync::atomic::Ordering::Relaxed) {
return Err(std::io::Error::new(
std::io::ErrorKind::BrokenPipe,
"Client disconnected",
));
}
match chunk {
Ok(bytes) => Ok(bytes::Bytes::from(bytes)),
Err(e) => {
if e.to_string().contains("disconnected")
|| e.to_string().contains("broken pipe")
{
client_connected_clone.store(false, std::sync::atomic::Ordering::Relaxed);
}
Err(std::io::Error::other(e.to_string()))
}
}
});
Response::builder()
.status(self.status_code)
.header("content-type", "text/html; charset=utf-8")
.header("transfer-encoding", "chunked")
.header("x-content-type-options", "nosniff")
.header("cache-control", "no-cache")
.header("vary", "Accept")
.body(Body::from_stream(stream))
.expect("Failed to build streaming response")
}
}
#[cfg(test)]
#[allow(clippy::disallowed_methods)]
mod tests {
use super::*;
use async_stream::stream;
use axum::body::to_bytes;
#[tokio::test]
async fn test_streaming_response_basic() {
let html_stream = stream! {
yield Ok(b"<!DOCTYPE html>".to_vec());
yield Ok(b"<html><body>".to_vec());
yield Ok(b"<h1>Hello</h1>".to_vec());
yield Ok(b"</body></html>".to_vec());
};
let response = StreamingHtmlResponse::new(html_stream).into_response();
assert_eq!(response.status(), 200);
assert_eq!(response.headers().get("content-type").unwrap(), "text/html; charset=utf-8");
assert_eq!(response.headers().get("transfer-encoding").unwrap(), "chunked");
assert_eq!(response.headers().get("x-content-type-options").unwrap(), "nosniff");
assert_eq!(response.headers().get("cache-control").unwrap(), "no-cache");
let body_bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
let body_str = String::from_utf8(body_bytes.to_vec()).unwrap();
assert_eq!(body_str, "<!DOCTYPE html><html><body><h1>Hello</h1></body></html>");
}
#[tokio::test]
async fn test_streaming_response_with_error() {
let html_stream = stream! {
yield Ok(b"<!DOCTYPE html>".to_vec());
yield Err(RariError::internal("Stream error"));
};
let response = StreamingHtmlResponse::new(html_stream).into_response();
assert_eq!(response.status(), 200);
let result = to_bytes(response.into_body(), usize::MAX).await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_streaming_response_empty() {
let html_stream = stream! {
if false {
yield Ok::<Vec<u8>, RariError>(vec![]);
}
};
let response = StreamingHtmlResponse::new(html_stream).into_response();
assert_eq!(response.status(), 200);
let body_bytes = to_bytes(response.into_body(), usize::MAX).await.unwrap();
assert_eq!(body_bytes.len(), 0);
}
#[tokio::test]
async fn test_streaming_response_headers() {
let html_stream = stream! {
yield Ok(b"<html></html>".to_vec());
};
let response = StreamingHtmlResponse::new(html_stream).into_response();
let headers = response.headers();
assert!(headers.contains_key("content-type"));
assert!(headers.contains_key("transfer-encoding"));
assert!(headers.contains_key("x-content-type-options"));
assert!(headers.contains_key("cache-control"));
}
}