1use crate::operation::RequestIdPolicy;
4use crate::rate_limit::RateLimit;
5use core::fmt;
6
7use super::cleanup::sanitize_response_storage;
8use super::retained::{ProtectedRequestId, RetainedMetadataError, RetainedResponseMetadata};
9use super::{ResponseContentType, ResponseHeaders, ResponseStorageSanitizer, StatusCode};
10
11#[derive(Debug)]
13pub struct ResponseMetadata {
14 rate_limit: Option<RateLimit>,
15}
16
17impl ResponseMetadata {
18 pub const EMPTY: Self = Self { rate_limit: None };
20
21 #[must_use]
23 pub fn with_rate_limit(mut self, rate_limit: RateLimit) -> Self {
24 self.rate_limit = Some(rate_limit);
25 self
26 }
27
28 pub(crate) const fn rate_limit(&self) -> Option<RateLimit> {
29 self.rate_limit
30 }
31}
32
33#[derive(Clone, Copy, Debug, Eq, PartialEq)]
35pub enum ResponseWriterError {
36 AlreadyCommitted,
38 NotCommitted,
40 InitializedLengthTooLarge,
42}
43
44impl_static_error!(ResponseWriterError,
45 Self::AlreadyCommitted => "response writer is already committed",
46 Self::NotCommitted => "response writer is not committed",
47 Self::InitializedLengthTooLarge => "response length exceeds admitted storage",
48);
49
50struct ResponseCommit {
51 status: StatusCode,
52 initialized_len: usize,
53 metadata: ResponseMetadata,
54}
55
56pub struct ResponseWriter<'buffer> {
68 storage: &'buffer mut [u8],
69 admitted_len: usize,
70 headers: ResponseHeaders<'buffer>,
71 request_id: Option<ProtectedRequestId>,
72 commit: Option<ResponseCommit>,
73}
74
75impl<'buffer> ResponseWriter<'buffer> {
76 pub fn begin_attempt(&mut self) -> Result<ResponseAttempt<'_, 'buffer>, ResponseWriterError> {
81 if self.commit.is_some() {
82 return Err(ResponseWriterError::AlreadyCommitted);
83 }
84 self.clear_uncommitted();
85 Ok(ResponseAttempt {
86 writer: self,
87 completed: false,
88 })
89 }
90
91 #[must_use]
93 pub const fn body_capacity(&self) -> usize {
94 self.admitted_len
95 }
96
97 fn body_mut(&mut self) -> Result<&mut [u8], ResponseWriterError> {
99 if self.commit.is_some() {
100 return Err(ResponseWriterError::AlreadyCommitted);
101 }
102 self.storage
103 .get_mut(..self.admitted_len)
104 .ok_or(ResponseWriterError::InitializedLengthTooLarge)
105 }
106
107 fn headers_mut(&mut self) -> Result<&mut ResponseHeaders<'buffer>, ResponseWriterError> {
109 if self.commit.is_some() {
110 return Err(ResponseWriterError::AlreadyCommitted);
111 }
112 Ok(&mut self.headers)
113 }
114
115 pub const fn headers(&self) -> &ResponseHeaders<'buffer> {
117 &self.headers
118 }
119
120 fn commit(
122 &mut self,
123 status: StatusCode,
124 initialized_len: usize,
125 metadata: ResponseMetadata,
126 ) -> Result<(), ResponseWriterError> {
127 if self.commit.is_some() {
128 return Err(ResponseWriterError::AlreadyCommitted);
129 }
130 if initialized_len > self.admitted_len {
131 return Err(ResponseWriterError::InitializedLengthTooLarge);
132 }
133 self.commit = Some(ResponseCommit {
134 status,
135 initialized_len,
136 metadata,
137 });
138 Ok(())
139 }
140
141 #[must_use]
143 pub const fn is_committed(&self) -> bool {
144 self.commit.is_some()
145 }
146
147 fn response(&self) -> Result<TransportResponse<'_, 'buffer>, ResponseWriterError> {
148 let commit = self
149 .commit
150 .as_ref()
151 .ok_or(ResponseWriterError::NotCommitted)?;
152 let body = self
153 .storage
154 .get(..commit.initialized_len)
155 .ok_or(ResponseWriterError::InitializedLengthTooLarge)?;
156 Ok(TransportResponse::from_commit(
157 commit,
158 body,
159 &self.headers,
160 self.request_id,
161 ))
162 }
163
164 fn apply_request_id_policy(
165 &mut self,
166 policy: RequestIdPolicy,
167 ) -> Result<(), RetainedMetadataError> {
168 let request_id = self.headers.hide_request_id()?;
169 match policy {
170 RequestIdPolicy::Discard => {
171 if let Some(request_id) = request_id {
172 self.headers.clear_protected(request_id);
173 }
174 }
175 RequestIdPolicy::Protected | RequestIdPolicy::Retain => {
176 self.request_id = request_id;
177 }
178 }
179 Ok(())
180 }
181
182 fn request_id(&self) -> Option<&[u8]> {
183 self.request_id
184 .and_then(|request_id| self.headers.protected_value(request_id))
185 }
186
187 fn retain_request_id<'destination>(
188 &mut self,
189 destination: &'destination mut [u8],
190 retention_limit: usize,
191 ) -> Result<RetainedResponseMetadata<'destination>, RetainedMetadataError> {
192 let mut retained = RetainedResponseMetadata::empty_for_core(destination);
193 let Some(request_id) = self.request_id.take() else {
194 return Ok(retained);
195 };
196 let result = {
197 let source = self
198 .headers
199 .protected_value(request_id)
200 .ok_or(RetainedMetadataError::RequestIdTooLong)?;
201 if source.len() > retention_limit {
202 Err(RetainedMetadataError::RetentionLimitExceeded)
203 } else {
204 retained.write_request_id(source)
205 }
206 };
207 self.headers.clear_protected(request_id);
208 result.map(|()| retained)
209 }
210
211 fn initialized_body(&self, initialized_len: usize) -> &[u8] {
212 self.storage.get(..initialized_len).unwrap_or_default()
213 }
214
215 fn clear_uncommitted(&mut self) {
216 if self.commit.is_none() {
217 sanitize_response_storage(self.storage, None);
218 self.headers.clear();
219 self.request_id = None;
220 }
221 }
222}
223
224impl fmt::Debug for ResponseWriter<'_> {
225 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
226 formatter
227 .debug_struct("ResponseWriter")
228 .field("storage_capacity", &self.storage.len())
229 .field("admitted_len", &self.admitted_len)
230 .field("committed", &self.commit.is_some())
231 .field("body", &"[redacted]")
232 .finish()
233 }
234}
235
236pub struct ResponseAttempt<'writer, 'buffer> {
242 writer: &'writer mut ResponseWriter<'buffer>,
243 completed: bool,
244}
245
246impl<'buffer> ResponseAttempt<'_, 'buffer> {
247 #[must_use]
249 pub const fn body_capacity(&self) -> usize {
250 self.writer.body_capacity()
251 }
252
253 pub fn body_mut(&mut self) -> Result<&mut [u8], ResponseWriterError> {
255 self.writer.body_mut()
256 }
257
258 pub fn headers_mut(&mut self) -> Result<&mut ResponseHeaders<'buffer>, ResponseWriterError> {
260 self.writer.headers_mut()
261 }
262
263 #[must_use]
265 pub const fn headers(&self) -> &ResponseHeaders<'buffer> {
266 self.writer.headers()
267 }
268
269 pub fn commit(
271 &mut self,
272 status: StatusCode,
273 initialized_len: usize,
274 metadata: ResponseMetadata,
275 ) -> Result<(), ResponseWriterError> {
276 self.writer.commit(status, initialized_len, metadata)?;
277 self.completed = true;
278 Ok(())
279 }
280}
281
282impl Drop for ResponseAttempt<'_, '_> {
283 fn drop(&mut self) {
284 if !self.completed {
285 self.writer.clear_uncommitted();
286 }
287 }
288}
289
290pub struct ResponseBuffer<'buffer> {
296 writer: ResponseWriter<'buffer>,
297 additive: Option<&'buffer dyn ResponseStorageSanitizer>,
298}
299
300impl<'buffer> ResponseBuffer<'buffer> {
301 #[must_use]
303 pub fn new(
304 storage: &'buffer mut [u8],
305 max_body_bytes: usize,
306 header_storage: &'buffer mut [u8],
307 ) -> Self {
308 Self::construct(storage, max_body_bytes, header_storage, None)
309 }
310
311 #[must_use]
313 pub fn with_additive_sanitizer(
314 storage: &'buffer mut [u8],
315 max_body_bytes: usize,
316 header_storage: &'buffer mut [u8],
317 additive: &'buffer dyn ResponseStorageSanitizer,
318 ) -> Self {
319 Self::construct(storage, max_body_bytes, header_storage, Some(additive))
320 }
321
322 fn construct(
323 storage: &'buffer mut [u8],
324 max_body_bytes: usize,
325 header_storage: &'buffer mut [u8],
326 additive: Option<&'buffer dyn ResponseStorageSanitizer>,
327 ) -> Self {
328 let headers = ResponseHeaders::new(header_storage);
329 sanitize_response_storage(storage, additive);
330 Self {
331 writer: ResponseWriter {
332 admitted_len: core::cmp::min(storage.len(), max_body_bytes),
333 storage,
334 headers,
335 request_id: None,
336 commit: None,
337 },
338 additive,
339 }
340 }
341
342 #[must_use]
344 pub const fn writer(&mut self) -> &mut ResponseWriter<'buffer> {
345 &mut self.writer
346 }
347
348 pub fn with_response<R>(
350 &self,
351 inspect: impl for<'response> FnOnce(TransportResponse<'response, 'buffer>) -> R,
352 ) -> Result<R, ResponseWriterError> {
353 let response = self.writer.response()?;
354 Ok(inspect(response))
355 }
356
357 pub(crate) fn response(&self) -> Result<TransportResponse<'_, 'buffer>, ResponseWriterError> {
358 self.writer.response()
359 }
360
361 pub(crate) fn apply_request_id_policy(
362 &mut self,
363 policy: RequestIdPolicy,
364 ) -> Result<(), RetainedMetadataError> {
365 self.writer.apply_request_id_policy(policy)
366 }
367
368 pub(crate) fn request_id(&self) -> Option<&[u8]> {
369 self.writer.request_id()
370 }
371
372 pub(crate) fn retain_request_id<'destination>(
373 &mut self,
374 destination: &'destination mut [u8],
375 retention_limit: usize,
376 ) -> Result<RetainedResponseMetadata<'destination>, RetainedMetadataError> {
377 self.writer.retain_request_id(destination, retention_limit)
378 }
379
380 pub(crate) fn initialized_body(&self, initialized_len: usize) -> &[u8] {
381 self.writer.initialized_body(initialized_len)
382 }
383}
384
385impl fmt::Debug for ResponseBuffer<'_> {
386 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
387 formatter
388 .debug_struct("ResponseBuffer")
389 .field("writer", &self.writer)
390 .field("additive", &self.additive.is_some())
391 .finish()
392 }
393}
394
395impl Drop for ResponseBuffer<'_> {
396 fn drop(&mut self) {
397 sanitize_response_storage(self.writer.storage, self.additive);
398 }
399}
400
401#[derive(Clone, Copy)]
416pub struct TransportResponse<'response, 'storage> {
417 status: StatusCode,
418 body: &'response [u8],
419 metadata: &'response ResponseMetadata,
420 headers: &'response ResponseHeaders<'storage>,
421 request_id: Option<ProtectedRequestId>,
422}
423
424impl<'response, 'storage> TransportResponse<'response, 'storage> {
425 fn from_commit(
426 commit: &'response ResponseCommit,
427 body: &'response [u8],
428 headers: &'response ResponseHeaders<'storage>,
429 request_id: Option<ProtectedRequestId>,
430 ) -> Self {
431 Self {
432 status: commit.status,
433 body,
434 metadata: &commit.metadata,
435 headers,
436 request_id,
437 }
438 }
439
440 #[must_use]
442 pub const fn status(&self) -> StatusCode {
443 self.status
444 }
445
446 #[must_use]
448 pub const fn body(&self) -> &'response [u8] {
449 self.body
450 }
451
452 pub fn content_type(
456 &self,
457 ) -> Result<Option<ResponseContentType<'response>>, super::ContentTypeError> {
458 let Some(header) = self.headers.get("content-type") else {
459 return Ok(None);
460 };
461 let value =
462 core::str::from_utf8(header.value()).map_err(|_| super::ContentTypeError::Invalid)?;
463 ResponseContentType::new(value).map(Some)
464 }
465
466 #[must_use]
468 pub const fn rate_limit(&self) -> Option<RateLimit> {
469 self.metadata.rate_limit()
470 }
471
472 #[must_use]
474 pub const fn headers(&self) -> &'response ResponseHeaders<'storage> {
475 self.headers
476 }
477
478 pub fn with_request_id<R>(&self, inspect: impl FnOnce(Option<&[u8]>) -> R) -> R {
480 inspect(
481 self.request_id
482 .and_then(|request_id| self.headers.protected_value(request_id)),
483 )
484 }
485}
486
487impl fmt::Debug for TransportResponse<'_, '_> {
488 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
489 formatter
490 .debug_struct("TransportResponse")
491 .field("status", &self.status)
492 .field("body_len", &self.body.len())
493 .field("body", &"[redacted]")
494 .field("metadata", &self.metadata)
495 .field("headers", &self.headers)
496 .field("request_id", &"[redacted]")
497 .finish()
498 }
499}