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) fn apply_request_id_policy(
319 &mut self,
320 policy: RequestIdPolicy,
321 ) -> Result<(), RetainedMetadataError> {
322 self.writer.apply_request_id_policy(policy)
323 }
324
325 pub(crate) fn request_id(&self) -> Option<&[u8]> {
326 self.writer.request_id()
327 }
328
329 pub(crate) fn retain_request_id<'destination>(
330 &mut self,
331 destination: &'destination mut [u8],
332 retention_limit: usize,
333 ) -> Result<RetainedResponseMetadata<'destination>, RetainedMetadataError> {
334 self.writer.retain_request_id(destination, retention_limit)
335 }
336
337 pub(crate) fn initialized_body(&self, initialized_len: usize) -> &[u8] {
338 self.writer.initialized_body(initialized_len)
339 }
340}
341
342impl fmt::Debug for ResponseBuffer<'_> {
343 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
344 formatter
345 .debug_struct("ResponseBuffer")
346 .field("writer", &self.writer)
347 .field("additive", &self.additive.is_some())
348 .finish()
349 }
350}
351
352impl Drop for ResponseBuffer<'_> {
353 fn drop(&mut self) {
354 sanitize_response_storage(self.writer.storage, self.additive);
355 }
356}
357
358#[derive(Clone, Copy)]
373pub struct TransportResponse<'response, 'storage> {
374 status: StatusCode,
375 body: &'response [u8],
376 metadata: &'response ResponseMetadata,
377 headers: &'response ResponseHeaders<'storage>,
378 request_id: Option<ProtectedRequestId>,
379}
380
381impl<'response, 'storage> TransportResponse<'response, 'storage> {
382 fn from_commit(
383 commit: &'response ResponseCommit,
384 body: &'response [u8],
385 headers: &'response ResponseHeaders<'storage>,
386 request_id: Option<ProtectedRequestId>,
387 ) -> Self {
388 Self {
389 status: commit.status,
390 body,
391 metadata: &commit.metadata,
392 headers,
393 request_id,
394 }
395 }
396
397 #[must_use]
399 pub const fn status(&self) -> StatusCode {
400 self.status
401 }
402
403 #[must_use]
405 pub const fn body(&self) -> &'response [u8] {
406 self.body
407 }
408
409 pub fn content_type(
413 &self,
414 ) -> Result<Option<ResponseContentType<'response>>, super::ContentTypeError> {
415 let Some(header) = self.headers.get("content-type") else {
416 return Ok(None);
417 };
418 let value =
419 core::str::from_utf8(header.value()).map_err(|_| super::ContentTypeError::Invalid)?;
420 ResponseContentType::new(value).map(Some)
421 }
422
423 #[must_use]
425 pub const fn rate_limit(&self) -> Option<RateLimit> {
426 self.metadata.rate_limit()
427 }
428
429 #[must_use]
431 pub const fn headers(&self) -> &'response ResponseHeaders<'storage> {
432 self.headers
433 }
434
435 pub fn with_request_id<R>(&self, inspect: impl FnOnce(Option<&[u8]>) -> R) -> R {
437 inspect(
438 self.request_id
439 .and_then(|request_id| self.headers.protected_value(request_id)),
440 )
441 }
442}
443
444impl fmt::Debug for TransportResponse<'_, '_> {
445 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
446 formatter
447 .debug_struct("TransportResponse")
448 .field("status", &self.status)
449 .field("body_len", &self.body.len())
450 .field("body", &"[redacted]")
451 .field("metadata", &self.metadata)
452 .field("headers", &self.headers)
453 .field("request_id", &"[redacted]")
454 .finish()
455 }
456}