use core::fmt;
use crate::rate_limit::RateLimit;
use super::{ResponseContentType, ResponseHeaders, ResponseStorageSanitizer, StatusCode};
#[derive(Clone, Copy, Debug)]
pub struct ResponseMetadata {
content_type: Option<ResponseContentType>,
rate_limit: Option<RateLimit>,
headers: ResponseHeaders,
}
impl ResponseMetadata {
pub const EMPTY: Self = Self {
content_type: None,
rate_limit: None,
headers: ResponseHeaders::new(),
};
#[must_use]
pub const fn with_content_type(mut self, content_type: ResponseContentType) -> Self {
self.content_type = Some(content_type);
self
}
#[must_use]
pub const fn with_rate_limit(mut self, rate_limit: RateLimit) -> Self {
self.rate_limit = Some(rate_limit);
self
}
#[must_use]
pub const fn with_headers(mut self, headers: ResponseHeaders) -> Self {
self.headers = headers;
self
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ResponseWriterError {
AlreadyCommitted,
NotCommitted,
InitializedLengthTooLarge,
}
impl_static_error!(ResponseWriterError,
Self::AlreadyCommitted => "response writer is already committed",
Self::NotCommitted => "response writer is not committed",
Self::InitializedLengthTooLarge => "response length exceeds admitted storage",
);
#[derive(Clone, Copy, Debug)]
struct ResponseCommit {
status: StatusCode,
initialized_len: usize,
metadata: ResponseMetadata,
}
pub struct ResponseWriter<'buffer> {
storage: &'buffer mut [u8],
admitted_len: usize,
commit: Option<ResponseCommit>,
}
impl ResponseWriter<'_> {
#[must_use]
pub const fn body_capacity(&self) -> usize {
self.admitted_len
}
pub fn body_mut(&mut self) -> Result<&mut [u8], ResponseWriterError> {
if self.commit.is_some() {
return Err(ResponseWriterError::AlreadyCommitted);
}
self.storage
.get_mut(..self.admitted_len)
.ok_or(ResponseWriterError::InitializedLengthTooLarge)
}
pub fn commit(
&mut self,
status: StatusCode,
initialized_len: usize,
metadata: ResponseMetadata,
) -> Result<(), ResponseWriterError> {
if self.commit.is_some() {
return Err(ResponseWriterError::AlreadyCommitted);
}
if initialized_len > self.admitted_len {
return Err(ResponseWriterError::InitializedLengthTooLarge);
}
self.commit = Some(ResponseCommit {
status,
initialized_len,
metadata,
});
Ok(())
}
#[must_use]
pub const fn is_committed(&self) -> bool {
self.commit.is_some()
}
fn response(&self) -> Result<TransportResponse<'_>, ResponseWriterError> {
let commit = self.commit.ok_or(ResponseWriterError::NotCommitted)?;
let body = self
.storage
.get(..commit.initialized_len)
.ok_or(ResponseWriterError::InitializedLengthTooLarge)?;
Ok(TransportResponse::from_commit(commit, body))
}
fn initialized_body(&self, initialized_len: usize) -> &[u8] {
self.storage.get(..initialized_len).unwrap_or_default()
}
}
impl fmt::Debug for ResponseWriter<'_> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ResponseWriter")
.field("storage_capacity", &self.storage.len())
.field("admitted_len", &self.admitted_len)
.field("committed", &self.commit.is_some())
.field("body", &"[redacted]")
.finish()
}
}
pub struct ResponseBuffer<'buffer, 'sanitizer, S: ResponseStorageSanitizer + ?Sized> {
writer: ResponseWriter<'buffer>,
sanitizer: &'sanitizer S,
}
impl<'buffer, 'sanitizer, S> ResponseBuffer<'buffer, 'sanitizer, S>
where
S: ResponseStorageSanitizer + ?Sized,
{
#[must_use]
pub fn new(
storage: &'buffer mut [u8],
max_body_bytes: usize,
sanitizer: &'sanitizer S,
) -> Self {
sanitizer.sanitize_response_storage(storage);
Self {
writer: ResponseWriter {
admitted_len: core::cmp::min(storage.len(), max_body_bytes),
storage,
commit: None,
},
sanitizer,
}
}
#[must_use]
pub const fn writer(&mut self) -> &mut ResponseWriter<'buffer> {
&mut self.writer
}
pub fn with_response<R>(
&self,
inspect: impl for<'response> FnOnce(TransportResponse<'response>) -> R,
) -> Result<R, ResponseWriterError> {
let response = self.writer.response()?;
Ok(inspect(response))
}
pub(crate) fn response(&self) -> Result<TransportResponse<'_>, ResponseWriterError> {
self.writer.response()
}
pub(crate) fn initialized_body(&self, initialized_len: usize) -> &[u8] {
self.writer.initialized_body(initialized_len)
}
}
impl<S> fmt::Debug for ResponseBuffer<'_, '_, S>
where
S: ResponseStorageSanitizer + ?Sized,
{
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ResponseBuffer")
.field("writer", &self.writer)
.finish()
}
}
impl<S> Drop for ResponseBuffer<'_, '_, S>
where
S: ResponseStorageSanitizer + ?Sized,
{
fn drop(&mut self) {
self.sanitizer
.sanitize_response_storage(self.writer.storage);
}
}
#[derive(Clone, Copy)]
pub struct TransportResponse<'buffer> {
status: StatusCode,
body: &'buffer [u8],
content_type: Option<ResponseContentType>,
rate_limit: Option<RateLimit>,
headers: ResponseHeaders,
}
impl<'buffer> TransportResponse<'buffer> {
const fn from_commit(commit: ResponseCommit, body: &'buffer [u8]) -> Self {
Self {
status: commit.status,
body,
content_type: commit.metadata.content_type,
rate_limit: commit.metadata.rate_limit,
headers: commit.metadata.headers,
}
}
#[must_use]
pub const fn status(&self) -> StatusCode {
self.status
}
#[must_use]
pub const fn body(&self) -> &'buffer [u8] {
self.body
}
#[must_use]
pub const fn content_type(&self) -> Option<ResponseContentType> {
self.content_type
}
#[must_use]
pub const fn rate_limit(&self) -> Option<RateLimit> {
self.rate_limit
}
#[must_use]
pub const fn headers(&self) -> &ResponseHeaders {
&self.headers
}
}
impl fmt::Debug for TransportResponse<'_> {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("TransportResponse")
.field("status", &self.status)
.field("body_len", &self.body.len())
.field("body", &"[redacted]")
.field("content_type", &self.content_type)
.field("rate_limit", &self.rate_limit)
.field("headers", &self.headers)
.finish()
}
}