1mod attempt;
4
5pub use attempt::{AsyncResponseStaging, ResponseAttempt, ResponseCompletion};
6
7use crate::operation::RequestIdPolicy;
8use crate::rate_limit::RateLimit;
9use core::fmt;
10
11use super::cleanup::sanitize_response_storage;
12use super::retained::{ProtectedRequestId, RetainedMetadataError, RetainedResponseMetadata};
13use super::{ResponseContentType, ResponseHeaders, ResponseStorageSanitizer, StatusCode};
14
15#[derive(Debug)]
17pub struct ResponseMetadata {
18 rate_limit: Option<RateLimit>,
19}
20
21impl ResponseMetadata {
22 pub const EMPTY: Self = Self { rate_limit: None };
24
25 #[must_use]
27 pub fn with_rate_limit(mut self, rate_limit: RateLimit) -> Self {
28 self.rate_limit = Some(rate_limit);
29 self
30 }
31
32 pub(crate) const fn rate_limit(&self) -> Option<RateLimit> {
33 self.rate_limit
34 }
35}
36
37#[derive(Clone, Copy, Debug, Eq, PartialEq)]
39pub enum ResponseWriterError {
40 AlreadyCommitted,
42 NotCommitted,
44 InitializedLengthTooLarge,
46}
47
48impl_static_error!(ResponseWriterError,
49 Self::AlreadyCommitted => "response writer is already committed",
50 Self::NotCommitted => "response writer is not committed",
51 Self::InitializedLengthTooLarge => "response length exceeds admitted storage",
52);
53
54struct ResponseCommit {
55 status: StatusCode,
56 initialized_len: usize,
57 metadata: ResponseMetadata,
58}
59
60pub struct ResponseWriter<'buffer> {
72 storage: &'buffer mut [u8],
73 admitted_len: usize,
74 headers: ResponseHeaders<'buffer>,
75 request_id: Option<ProtectedRequestId>,
76 commit: Option<ResponseCommit>,
77}
78
79impl<'buffer> ResponseWriter<'buffer> {
80 pub fn begin_attempt(&mut self) -> Result<ResponseAttempt<'_, 'buffer>, ResponseWriterError> {
85 if self.commit.is_some() {
86 return Err(ResponseWriterError::AlreadyCommitted);
87 }
88 self.clear_uncommitted();
89 Ok(ResponseAttempt {
90 writer: self,
91 completed: false,
92 })
93 }
94
95 #[must_use]
97 pub const fn body_capacity(&self) -> usize {
98 self.admitted_len
99 }
100
101 fn body_mut(&mut self) -> Result<&mut [u8], ResponseWriterError> {
103 if self.commit.is_some() {
104 return Err(ResponseWriterError::AlreadyCommitted);
105 }
106 self.storage
107 .get_mut(..self.admitted_len)
108 .ok_or(ResponseWriterError::InitializedLengthTooLarge)
109 }
110
111 fn headers_mut(&mut self) -> Result<&mut ResponseHeaders<'buffer>, ResponseWriterError> {
113 if self.commit.is_some() {
114 return Err(ResponseWriterError::AlreadyCommitted);
115 }
116 Ok(&mut self.headers)
117 }
118
119 pub const fn headers(&self) -> &ResponseHeaders<'buffer> {
121 &self.headers
122 }
123
124 fn commit(
126 &mut self,
127 status: StatusCode,
128 initialized_len: usize,
129 metadata: ResponseMetadata,
130 ) -> Result<(), ResponseWriterError> {
131 if self.commit.is_some() {
132 return Err(ResponseWriterError::AlreadyCommitted);
133 }
134 if initialized_len > self.admitted_len {
135 return Err(ResponseWriterError::InitializedLengthTooLarge);
136 }
137 self.commit = Some(ResponseCommit {
138 status,
139 initialized_len,
140 metadata,
141 });
142 Ok(())
143 }
144
145 #[must_use]
147 pub const fn is_committed(&self) -> bool {
148 self.commit.is_some()
149 }
150
151 fn response(&self) -> Result<TransportResponse<'_, 'buffer>, ResponseWriterError> {
152 let commit = self
153 .commit
154 .as_ref()
155 .ok_or(ResponseWriterError::NotCommitted)?;
156 let body = self
157 .storage
158 .get(..commit.initialized_len)
159 .ok_or(ResponseWriterError::InitializedLengthTooLarge)?;
160 Ok(TransportResponse::from_commit(
161 commit,
162 body,
163 &self.headers,
164 self.request_id,
165 ))
166 }
167
168 fn apply_request_id_policy(
169 &mut self,
170 policy: RequestIdPolicy,
171 ) -> Result<(), RetainedMetadataError> {
172 let request_id = self.headers.hide_request_id()?;
173 match policy {
174 RequestIdPolicy::Discard => {
175 if let Some(request_id) = request_id {
176 self.headers.clear_protected(request_id);
177 }
178 }
179 RequestIdPolicy::Protected | RequestIdPolicy::Retain => {
180 self.request_id = request_id;
181 }
182 }
183 Ok(())
184 }
185
186 fn request_id(&self) -> Option<&[u8]> {
187 self.request_id
188 .and_then(|request_id| self.headers.protected_value(request_id))
189 }
190
191 fn retain_request_id<'destination>(
192 &mut self,
193 destination: &'destination mut [u8],
194 retention_limit: usize,
195 ) -> Result<RetainedResponseMetadata<'destination>, RetainedMetadataError> {
196 let mut retained = RetainedResponseMetadata::empty_for_core(destination);
197 let Some(request_id) = self.request_id.take() else {
198 return Ok(retained);
199 };
200 let result = {
201 let source = self
202 .headers
203 .protected_value(request_id)
204 .ok_or(RetainedMetadataError::RequestIdTooLong)?;
205 if source.len() > retention_limit {
206 Err(RetainedMetadataError::RetentionLimitExceeded)
207 } else {
208 retained.write_request_id(source)
209 }
210 };
211 self.headers.clear_protected(request_id);
212 result.map(|()| retained)
213 }
214
215 fn initialized_body(&self, initialized_len: usize) -> &[u8] {
216 self.storage.get(..initialized_len).unwrap_or_default()
217 }
218
219 fn clear_uncommitted(&mut self) {
220 if self.commit.is_none() {
221 sanitize_response_storage(self.storage, None);
222 self.headers.clear();
223 self.request_id = None;
224 }
225 }
226
227 fn rollback_attempt(&mut self) {
228 self.commit = None;
229 sanitize_response_storage(self.storage, None);
230 self.headers.clear();
231 self.request_id = None;
232 }
233}
234
235impl fmt::Debug for ResponseWriter<'_> {
236 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
237 formatter
238 .debug_struct("ResponseWriter")
239 .field("storage_capacity", &self.storage.len())
240 .field("admitted_len", &self.admitted_len)
241 .field("committed", &self.commit.is_some())
242 .field("body", &"[redacted]")
243 .finish()
244 }
245}
246
247pub struct ResponseBuffer<'buffer> {
253 writer: ResponseWriter<'buffer>,
254 additive: Option<&'buffer dyn ResponseStorageSanitizer>,
255}
256
257impl<'buffer> ResponseBuffer<'buffer> {
258 #[must_use]
260 pub fn new(
261 storage: &'buffer mut [u8],
262 max_body_bytes: usize,
263 header_storage: &'buffer mut [u8],
264 ) -> Self {
265 Self::construct(storage, max_body_bytes, header_storage, None)
266 }
267
268 #[must_use]
270 pub fn with_additive_sanitizer(
271 storage: &'buffer mut [u8],
272 max_body_bytes: usize,
273 header_storage: &'buffer mut [u8],
274 additive: &'buffer dyn ResponseStorageSanitizer,
275 ) -> Self {
276 Self::construct(storage, max_body_bytes, header_storage, Some(additive))
277 }
278
279 fn construct(
280 storage: &'buffer mut [u8],
281 max_body_bytes: usize,
282 header_storage: &'buffer mut [u8],
283 additive: Option<&'buffer dyn ResponseStorageSanitizer>,
284 ) -> Self {
285 let headers = ResponseHeaders::new(header_storage);
286 sanitize_response_storage(storage, additive);
287 Self {
288 writer: ResponseWriter {
289 admitted_len: core::cmp::min(storage.len(), max_body_bytes),
290 storage,
291 headers,
292 request_id: None,
293 commit: None,
294 },
295 additive,
296 }
297 }
298
299 #[must_use]
301 pub const fn writer(&mut self) -> &mut ResponseWriter<'buffer> {
302 &mut self.writer
303 }
304
305 pub fn with_response<R>(
307 &self,
308 inspect: impl for<'response> FnOnce(TransportResponse<'response, 'buffer>) -> R,
309 ) -> Result<R, ResponseWriterError> {
310 let response = self.writer.response()?;
311 Ok(inspect(response))
312 }
313
314 pub(crate) fn response(&self) -> Result<TransportResponse<'_, 'buffer>, ResponseWriterError> {
315 self.writer.response()
316 }
317
318 pub(crate) const fn headers(&self) -> &ResponseHeaders<'buffer> {
319 self.writer.headers()
320 }
321
322 pub(crate) fn apply_request_id_policy(
323 &mut self,
324 policy: RequestIdPolicy,
325 ) -> Result<(), RetainedMetadataError> {
326 self.writer.apply_request_id_policy(policy)
327 }
328
329 pub(crate) fn request_id(&self) -> Option<&[u8]> {
330 self.writer.request_id()
331 }
332
333 pub(crate) fn has_request_id(&self) -> Result<bool, ResponseWriterError> {
334 if self.request_id().is_some() {
335 return Ok(true);
336 }
337 self.with_response(|response| response.headers().get("x-request-id").is_some())
338 }
339
340 pub(crate) fn retain_request_id<'destination>(
341 &mut self,
342 destination: &'destination mut [u8],
343 retention_limit: usize,
344 ) -> Result<RetainedResponseMetadata<'destination>, RetainedMetadataError> {
345 self.writer.retain_request_id(destination, retention_limit)
346 }
347
348 pub(crate) fn initialized_body(&self, initialized_len: usize) -> &[u8] {
349 self.writer.initialized_body(initialized_len)
350 }
351}
352
353impl fmt::Debug for ResponseBuffer<'_> {
354 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
355 formatter
356 .debug_struct("ResponseBuffer")
357 .field("writer", &self.writer)
358 .field("additive", &self.additive.is_some())
359 .finish()
360 }
361}
362
363impl Drop for ResponseBuffer<'_> {
364 fn drop(&mut self) {
365 sanitize_response_storage(self.writer.storage, self.additive);
366 }
367}
368
369#[derive(Clone, Copy)]
384pub struct TransportResponse<'response, 'storage> {
385 status: StatusCode,
386 body: &'response [u8],
387 metadata: &'response ResponseMetadata,
388 headers: &'response ResponseHeaders<'storage>,
389 request_id: Option<ProtectedRequestId>,
390}
391
392impl<'response, 'storage> TransportResponse<'response, 'storage> {
393 fn from_commit(
394 commit: &'response ResponseCommit,
395 body: &'response [u8],
396 headers: &'response ResponseHeaders<'storage>,
397 request_id: Option<ProtectedRequestId>,
398 ) -> Self {
399 Self {
400 status: commit.status,
401 body,
402 metadata: &commit.metadata,
403 headers,
404 request_id,
405 }
406 }
407
408 #[must_use]
410 pub const fn status(&self) -> StatusCode {
411 self.status
412 }
413
414 #[must_use]
416 pub const fn body(&self) -> &'response [u8] {
417 self.body
418 }
419
420 pub fn content_type(
424 &self,
425 ) -> Result<Option<ResponseContentType<'response>>, super::ContentTypeError> {
426 let Some(header) = self.headers.get("content-type") else {
427 return Ok(None);
428 };
429 let value =
430 core::str::from_utf8(header.value()).map_err(|_| super::ContentTypeError::Invalid)?;
431 ResponseContentType::new(value).map(Some)
432 }
433
434 #[must_use]
436 pub const fn rate_limit(&self) -> Option<RateLimit> {
437 self.metadata.rate_limit()
438 }
439
440 #[must_use]
442 pub const fn headers(&self) -> &'response ResponseHeaders<'storage> {
443 self.headers
444 }
445
446 pub fn with_request_id<R>(&self, inspect: impl FnOnce(Option<&[u8]>) -> R) -> R {
448 inspect(
449 self.request_id
450 .and_then(|request_id| self.headers.protected_value(request_id)),
451 )
452 }
453}
454
455impl fmt::Debug for TransportResponse<'_, '_> {
456 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
457 formatter
458 .debug_struct("TransportResponse")
459 .field("status", &self.status)
460 .field("body_len", &self.body.len())
461 .field("body", &"[redacted]")
462 .field("metadata", &self.metadata)
463 .field("headers", &self.headers)
464 .field("request_id", &"[redacted]")
465 .finish()
466 }
467}