use std::convert::Infallible;
use crate::{
Cookie, Error, Headers, InvalidHeader, ResponseStream, ValidationErrors,
error::PendingError,
openapi::Operation,
schemaval::{SchemaKind, SchemaMetadata},
};
#[cfg(feature = "websocket")]
use crate::websocket::WebSocketPlan;
pub enum ResponseBody {
Buffered(Vec<u8>),
Streaming(Box<dyn ResponseStream>),
#[cfg(feature = "websocket")]
WebSocket(WebSocketPlan),
}
impl ResponseBody {
pub fn buffered(&self) -> Option<&[u8]> {
match self {
Self::Buffered(bytes) => Some(bytes),
Self::Streaming(_) => None,
#[cfg(feature = "websocket")]
Self::WebSocket(_) => None,
}
}
pub fn is_streaming(&self) -> bool {
matches!(self, Self::Streaming(_))
}
}
pub struct Response {
status: u16,
headers: Headers,
body: ResponseBody,
error: Option<PendingError>,
}
impl Response {
pub fn new(status: u16) -> Self {
Self {
status,
headers: Headers::new(),
body: ResponseBody::Buffered(Vec::new()),
error: None,
}
}
pub fn empty() -> Self {
Self::new(204)
}
pub fn text(status: u16, text: impl Into<String>) -> Self {
let mut response = Self {
status,
headers: Headers::new(),
body: ResponseBody::Buffered(text.into().into_bytes()),
error: None,
};
response
.headers
.set_unchecked("Content-Type", "text/plain; charset=utf-8");
response
}
pub fn bytes(status: u16, bytes: impl Into<Vec<u8>>) -> Self {
let mut response = Self {
status,
headers: Headers::new(),
body: ResponseBody::Buffered(bytes.into()),
error: None,
};
response
.headers
.set_unchecked("Content-Type", "application/octet-stream");
response
}
pub fn stream(status: u16, stream: impl ResponseStream + 'static) -> Self {
Self {
status,
headers: Headers::new(),
body: ResponseBody::Streaming(Box::new(stream)),
error: None,
}
}
#[cfg(feature = "websocket")]
pub(crate) fn websocket(plan: WebSocketPlan) -> Self {
let selected_protocol = plan.selected_protocol().map(str::to_owned);
let mut response = Self {
status: 101,
headers: Headers::new(),
body: ResponseBody::WebSocket(plan),
error: None,
};
if let Some(protocol) = selected_protocol {
response
.headers
.set_unchecked("Sec-WebSocket-Protocol", protocol);
}
response
}
pub(crate) fn pending_error(error: Error) -> Self {
Self {
status: error.status(),
headers: Headers::new(),
body: ResponseBody::Buffered(Vec::new()),
error: Some(PendingError::new(error)),
}
}
pub(crate) fn pending_validation(error: Error, validation: ValidationErrors) -> Self {
Self {
status: error.status(),
headers: Headers::new(),
body: ResponseBody::Buffered(Vec::new()),
error: Some(PendingError::validation(error, validation)),
}
}
pub fn status(&self) -> u16 {
self.status
}
pub fn content_type(&self) -> Option<&str> {
self.headers
.get("content-type")
.and_then(|value| std::str::from_utf8(value).ok())
}
pub fn body(&self) -> &[u8] {
self.body.buffered().unwrap_or_default()
}
pub fn is_streaming(&self) -> bool {
self.body.is_streaming()
}
pub fn headers(&mut self) -> &mut Headers {
&mut self.headers
}
pub fn set_cookie(&mut self, cookie: Cookie) -> Result<(), InvalidHeader> {
self.headers.append("Set-Cookie", cookie.header_value())
}
pub(crate) fn set_header(&mut self, name: impl Into<String>, value: impl Into<Vec<u8>>) {
self.headers.set_unchecked(name, value);
}
pub(crate) fn without_body(mut self) -> Self {
match &mut self.body {
ResponseBody::Buffered(body) => {
let content_length = body.len().to_string();
body.clear();
self.set_header("Content-Length", content_length);
}
ResponseBody::Streaming(_) => {
self.body = ResponseBody::Buffered(Vec::new());
}
#[cfg(feature = "websocket")]
ResponseBody::WebSocket(_) => {}
}
self
}
pub(crate) fn take_error(&mut self) -> Option<PendingError> {
self.error.take()
}
pub(crate) fn set_status(&mut self, status: u16) {
self.status = status;
}
pub(crate) fn take_headers(&mut self) -> Headers {
std::mem::take(&mut self.headers)
}
pub(crate) fn merge_headers(&mut self, headers: Headers) {
self.headers.merge_from(headers);
}
pub fn into_parts(self) -> (u16, Headers, ResponseBody) {
(self.status, self.headers, self.body)
}
}
pub trait IntoResponse {
fn into_response(self) -> Response;
#[doc(hidden)]
fn openapi(operation: &mut Operation) {
operation.response(200, "Success", None, None);
}
}
impl IntoResponse for Response {
fn into_response(self) -> Response {
self
}
}
impl IntoResponse for () {
fn into_response(self) -> Response {
Response::empty()
}
fn openapi(operation: &mut Operation) {
operation.response(204, "No Content", None, None);
}
}
impl IntoResponse for String {
fn into_response(self) -> Response {
Response::text(200, self)
}
fn openapi(operation: &mut Operation) {
operation.response(
200,
"Success",
Some("text/plain; charset=utf-8"),
Some(SchemaMetadata::new(SchemaKind::String)),
);
}
}
impl IntoResponse for &str {
fn into_response(self) -> Response {
Response::text(200, self)
}
fn openapi(operation: &mut Operation) {
String::openapi(operation);
}
}
impl IntoResponse for Vec<u8> {
fn into_response(self) -> Response {
Response::bytes(200, self)
}
fn openapi(operation: &mut Operation) {
operation.response(
200,
"Success",
Some("application/octet-stream"),
Some(SchemaMetadata::new(SchemaKind::Bytes)),
);
}
}
impl IntoResponse for Infallible {
fn into_response(self) -> Response {
match self {}
}
}
impl<T: IntoResponse, E: IntoResponse> IntoResponse for Result<T, E> {
fn into_response(self) -> Response {
match self {
Ok(value) => value.into_response(),
Err(error) => error.into_response(),
}
}
fn openapi(operation: &mut Operation) {
T::openapi(operation);
E::openapi(operation);
}
}
#[cfg(test)]
mod tests {
use std::task::{Context, Poll, Waker};
use crate::{
Chunk, Cookie, IntoResponse, Redirect, Response, ResponseBody, ResponseStream, StreamError,
};
struct OneChunk {
chunk: Vec<u8>,
sent: bool,
}
impl ResponseStream for OneChunk {
fn poll_next(
&mut self,
_context: &mut Context<'_>,
) -> Poll<Option<Result<Chunk, StreamError>>> {
if self.sent {
Poll::Ready(None)
} else {
self.sent = true;
Poll::Ready(Some(Ok(Chunk::from(std::mem::take(&mut self.chunk)))))
}
}
}
#[test]
fn headers_set_append_and_remove_case_insensitively() {
let mut response = Response::text(200, "ok");
response.headers().set("X-Value", "first").unwrap();
response.headers().append("x-value", "second").unwrap();
assert_eq!(
response.headers().get_all("X-VALUE").collect::<Vec<_>>(),
vec![b"first".as_slice(), b"second".as_slice()],
);
response.headers().set("X-VALUE", "replacement").unwrap();
assert_eq!(
response.headers().get_all("x-value").collect::<Vec<_>>(),
vec![b"replacement".as_slice()],
);
response.headers().remove("x-VaLuE");
assert!(!response.headers().contains("x-value"));
}
#[test]
fn content_type_uses_the_shared_header_collection() {
let mut response = Response::text(200, "ok");
response
.headers()
.set("Content-Type", "application/custom")
.unwrap();
assert_eq!(response.content_type(), Some("application/custom"));
}
#[test]
fn appends_set_cookie_headers() {
let mut response = Response::new(200);
response.set_cookie(Cookie::new("a", "1")).unwrap();
response.set_cookie(Cookie::new("b", "2")).unwrap();
assert_eq!(response.headers().get_all("set-cookie").count(), 2);
}
#[test]
fn creates_streaming_responses_and_redirects() {
let response = Response::stream(
200,
OneChunk {
chunk: b"chunk".to_vec(),
sent: false,
},
);
assert!(response.is_streaming());
let (_, _, ResponseBody::Streaming(mut stream)) = response.into_parts() else {
panic!("expected a streaming response");
};
let mut context = Context::from_waker(Waker::noop());
let Poll::Ready(Some(Ok(chunk))) = stream.poll_next(&mut context) else {
panic!("expected a response chunk");
};
assert_eq!(chunk.bytes(), b"chunk");
assert!(matches!(stream.poll_next(&mut context), Poll::Ready(None)));
let mut redirect = Redirect::temporary("/next").into_response();
assert_eq!(redirect.status(), 307);
assert_eq!(
redirect.headers().get("location"),
Some(b"/next".as_slice())
);
}
}