diff --git a/moli-fetch/src/lib.rs b/moli-fetch/src/lib.rs index 0914dd20f..cca1f8655 100644 --- a/moli-fetch/src/lib.rs +++ b/moli-fetch/src/lib.rs @@ -63,7 +63,7 @@ pub use request::{ pub use request_policy::{is_bad_port, should_request_be_blocked_due_to_bad_port}; pub use response::{ NegotiatedHttpVersion, NetworkRequestExtraInfo, NetworkResponseExtraInfo, RawResponse, - RedirectInfo, Response, ResponseBody, ResponseHead, + RedirectInfo, Response, ResponseBody, ResponseHead, SharedResponseBodyBytes, }; pub use runtime::PendingStreamingRawResponse; pub use runtime::{ diff --git a/moli-fetch/src/response.rs b/moli-fetch/src/response.rs index a7f7b7343..dc618f308 100644 --- a/moli-fetch/src/response.rs +++ b/moli-fetch/src/response.rs @@ -1,5 +1,6 @@ use anyhow::{Result, anyhow}; use moli_cookie_jar::{StoredCookieQueryReport, StoredCookieSetReport}; +use std::sync::Arc; use url::Url; use crate::{StreamingHtmlResponse, StreamingRawResponse}; @@ -50,21 +51,100 @@ pub struct ResponseHead { pub negotiated_http_version: Option, } +#[derive(Clone, Debug)] +pub struct SharedResponseBodyBytes { + backing: SharedResponseBodyBytesBacking, +} + +#[derive(Clone, Debug)] +enum SharedResponseBodyBytesBacking { + Text(Arc), + Bytes(Arc>), +} + +impl SharedResponseBodyBytes { + pub fn from_bytes(bytes: Vec) -> Self { + Self { + backing: SharedResponseBodyBytesBacking::Bytes(Arc::new(bytes)), + } + } + + fn from_text(text: Arc) -> Self { + Self { + backing: SharedResponseBodyBytesBacking::Text(text), + } + } + + fn from_shared_bytes(bytes: Arc>) -> Self { + Self { + backing: SharedResponseBodyBytesBacking::Bytes(bytes), + } + } + + pub fn as_slice(&self) -> &[u8] { + match &self.backing { + SharedResponseBodyBytesBacking::Text(text) => text.as_bytes(), + SharedResponseBodyBytesBacking::Bytes(bytes) => bytes, + } + } + + pub fn capacity(&self) -> usize { + match &self.backing { + SharedResponseBodyBytesBacking::Text(text) => text.capacity(), + SharedResponseBodyBytesBacking::Bytes(bytes) => bytes.capacity(), + } + } +} + +fn unwrap_shared_string(value: Arc) -> String { + Arc::try_unwrap(value).unwrap_or_else(|shared| (*shared).clone()) +} + +fn unwrap_shared_bytes(value: Arc>) -> Vec { + Arc::try_unwrap(value).unwrap_or_else(|shared| (*shared).clone()) +} + #[derive(Debug)] pub enum ResponseBody { - MaterializedText { text: String, bytes: Vec }, - MaterializedBytes(Vec), + MaterializedText { + text: Arc, + exact_bytes: Option>>, + }, + MaterializedBytes(Arc>), StreamingText(Box), StreamingBytes(Box), } impl ResponseBody { pub fn materialized_text(text: String, bytes: Vec) -> Self { - Self::MaterializedText { text, bytes } + let exact_bytes = (!bytes.is_empty()).then(|| Arc::new(bytes)); + Self::MaterializedText { + text: Arc::new(text), + exact_bytes, + } + } + + /// Builds a text body from exact bytes without retaining a second copy + /// when the bytes are already valid UTF-8. + pub fn lossy_text_from_bytes(bytes: Vec) -> Self { + match String::from_utf8(bytes) { + Ok(text) => Self::MaterializedText { + text: Arc::new(text), + exact_bytes: None, + }, + Err(error) => { + let bytes = error.into_bytes(); + let text = String::from_utf8_lossy(&bytes).into_owned(); + Self::MaterializedText { + text: Arc::new(text), + exact_bytes: Some(Arc::new(bytes)), + } + } + } } pub fn materialized_bytes(bytes: Vec) -> Self { - Self::MaterializedBytes(bytes) + Self::MaterializedBytes(Arc::new(bytes)) } pub fn is_streaming(&self) -> bool { @@ -73,14 +153,11 @@ impl ResponseBody { pub fn try_into_materialized_bytes(self) -> std::result::Result, Self> { match self { - Self::MaterializedBytes(bytes) => Ok(bytes), - Self::MaterializedText { text, bytes } => { - if bytes.is_empty() && !text.is_empty() { - Ok(text.into_bytes()) - } else { - Ok(bytes) - } - } + Self::MaterializedBytes(bytes) => Ok(unwrap_shared_bytes(bytes)), + Self::MaterializedText { text, exact_bytes } => Ok(exact_bytes.map_or_else( + || unwrap_shared_string(text).into_bytes(), + unwrap_shared_bytes, + )), Self::StreamingText(_) | Self::StreamingBytes(_) => Err(self), } } @@ -88,13 +165,11 @@ impl ResponseBody { pub fn as_materialized_bytes(&self) -> Option<&[u8]> { match self { Self::MaterializedBytes(bytes) => Some(bytes), - Self::MaterializedText { text, bytes } => { - if bytes.is_empty() && !text.is_empty() { - Some(text.as_bytes()) - } else { - Some(bytes) - } - } + Self::MaterializedText { text, exact_bytes } => Some( + exact_bytes + .as_deref() + .map_or(text.as_bytes(), Vec::as_slice), + ), Self::StreamingText(_) | Self::StreamingBytes(_) => None, } } @@ -102,9 +177,9 @@ impl ResponseBody { pub fn clone_materialized(&self) -> Option { match self { Self::MaterializedBytes(bytes) => Some(Self::MaterializedBytes(bytes.clone())), - Self::MaterializedText { text, bytes } => Some(Self::MaterializedText { + Self::MaterializedText { text, exact_bytes } => Some(Self::MaterializedText { text: text.clone(), - bytes: bytes.clone(), + exact_bytes: exact_bytes.clone(), }), Self::StreamingText(_) | Self::StreamingBytes(_) => None, } @@ -119,8 +194,14 @@ impl ResponseBody { pub fn try_into_lossy_materialized_text(self) -> std::result::Result<(String, Vec), Self> { match self { - Self::MaterializedText { text, bytes } => Ok((text, bytes)), + Self::MaterializedText { text, exact_bytes } => { + let bytes = + exact_bytes.map_or_else(|| text.as_bytes().to_vec(), unwrap_shared_bytes); + let text = unwrap_shared_string(text); + Ok((text, bytes)) + } Self::MaterializedBytes(bytes) => { + let bytes = unwrap_shared_bytes(bytes); let text = String::from_utf8_lossy(&bytes).into_owned(); Ok((text, bytes)) } @@ -128,20 +209,28 @@ impl ResponseBody { } } + /// Converts an already-materialized body to its text representation. + pub fn try_into_materialized_text_body(self) -> std::result::Result { + match self { + Self::MaterializedText { .. } => Ok(self), + Self::MaterializedBytes(bytes) => { + Ok(Self::lossy_text_from_bytes(unwrap_shared_bytes(bytes))) + } + Self::StreamingText(_) | Self::StreamingBytes(_) => Err(self), + } + } + /// Drains any streaming source and returns a complete byte body. /// /// This is an explicit compatibility boundary. Prefer chunked consumption /// when the caller can avoid building a full response body in memory. pub async fn into_materialized_bytes(self) -> Result> { match self { - Self::MaterializedBytes(bytes) => Ok(bytes), - Self::MaterializedText { text, bytes } => { - if bytes.is_empty() && !text.is_empty() { - Ok(text.into_bytes()) - } else { - Ok(bytes) - } - } + Self::MaterializedBytes(bytes) => Ok(unwrap_shared_bytes(bytes)), + Self::MaterializedText { text, exact_bytes } => Ok(exact_bytes.map_or_else( + || unwrap_shared_string(text).into_bytes(), + unwrap_shared_bytes, + )), Self::StreamingText(response) => { let mut response = *response; let mut bytes = Vec::new(); @@ -167,8 +256,14 @@ impl ResponseBody { /// the exact bytes used to derive it. pub async fn into_lossy_materialized_text(self) -> Result<(String, Vec)> { match self { - Self::MaterializedText { text, bytes } => Ok((text, bytes)), + Self::MaterializedText { text, exact_bytes } => { + let bytes = + exact_bytes.map_or_else(|| text.as_bytes().to_vec(), unwrap_shared_bytes); + let text = unwrap_shared_string(text); + Ok((text, bytes)) + } Self::MaterializedBytes(bytes) => { + let bytes = unwrap_shared_bytes(bytes); let text = String::from_utf8_lossy(&bytes).into_owned(); Ok((text, bytes)) } @@ -191,6 +286,47 @@ impl ResponseBody { } } } + + /// Drains a body source and returns a materialized text body. + pub async fn into_materialized_text_body(self) -> Result { + match self { + Self::MaterializedText { .. } => Ok(self), + Self::MaterializedBytes(bytes) => { + Ok(Self::lossy_text_from_bytes(unwrap_shared_bytes(bytes))) + } + Self::StreamingText(response) => { + let mut response = *response; + let mut text = String::new(); + while let Some(chunk) = response.next_chunk().await { + text.push_str(&chunk); + } + response.finish().await?; + Ok(Self::MaterializedText { + text: Arc::new(text), + exact_bytes: None, + }) + } + Self::StreamingBytes(response) => { + let bytes = Self::StreamingBytes(response) + .into_materialized_bytes() + .await?; + Ok(Self::lossy_text_from_bytes(bytes)) + } + } + } + + pub fn shared_materialized_bytes(&self) -> Option { + match self { + Self::MaterializedText { text, exact_bytes } => Some(exact_bytes.as_ref().map_or_else( + || SharedResponseBodyBytes::from_text(Arc::clone(text)), + |bytes| SharedResponseBodyBytes::from_shared_bytes(Arc::clone(bytes)), + )), + Self::MaterializedBytes(bytes) => Some(SharedResponseBodyBytes::from_shared_bytes( + Arc::clone(bytes), + )), + Self::StreamingText(_) | Self::StreamingBytes(_) => None, + } + } } #[derive(Debug)] @@ -260,6 +396,12 @@ impl Response { self.body_bytes().to_vec() } + pub fn shared_body_bytes(&self) -> SharedResponseBodyBytes { + self.body + .shared_materialized_bytes() + .expect("Response body should remain materialized") + } + pub fn materialized_body(&self) -> ResponseBody { self.body .clone_materialized() @@ -289,26 +431,25 @@ impl Response { } pub fn from_head_and_text_body(head: ResponseHead, body: String) -> Self { - let body_bytes = body.as_bytes().to_vec(); - Self::from_head_and_body(head, body, body_bytes) + Self::from_head_and_body(head, body, Vec::new()) } /// Builds a materialized text response from exact bytes by deriving the /// compatibility text view with UTF-8 replacement semantics. pub fn from_head_and_lossy_body_bytes(head: ResponseHead, body_bytes: Vec) -> Self { - let body = String::from_utf8_lossy(&body_bytes).into_owned(); - Self::from_head_and_body(head, body, body_bytes) + Self::from_head_and_materialized_body(head, ResponseBody::lossy_text_from_bytes(body_bytes)) + .expect("materialized text body should build a Response") } pub fn from_head_and_materialized_body(head: ResponseHead, body: ResponseBody) -> Result { - let (body, body_bytes) = body.try_into_lossy_materialized_text().map_err(|_| { + let body = body.try_into_materialized_text_body().map_err(|_| { anyhow!("cannot build materialized Response from streaming response body") })?; Ok(Self { final_url: head.final_url, status: head.status, headers: head.headers, - body: ResponseBody::materialized_text(body, body_bytes), + body, request_cookie_report: head.request_cookie_report, cookie_set_reports: head.cookie_set_reports, redirected: head.redirected, @@ -320,12 +461,12 @@ impl Response { } pub async fn from_head_and_body_source(head: ResponseHead, body: ResponseBody) -> Result { - let (body, body_bytes) = body.into_lossy_materialized_text().await?; + let body = body.into_materialized_text_body().await?; Ok(Self { final_url: head.final_url, status: head.status, headers: head.headers, - body: ResponseBody::materialized_text(body, body_bytes), + body, request_cookie_report: head.request_cookie_report, cookie_set_reports: head.cookie_set_reports, redirected: head.redirected, diff --git a/moli-fetch/src/tests/mod.rs b/moli-fetch/src/tests/mod.rs index c469fdeee..f63759784 100644 --- a/moli-fetch/src/tests/mod.rs +++ b/moli-fetch/src/tests/mod.rs @@ -247,6 +247,56 @@ fn response_body_marks_materialized_and_streaming_shapes() { assert!(streaming_body.try_into_materialized_bytes().is_err()); } +#[test] +fn response_body_reuses_valid_utf8_bytes_as_text_storage() { + let body = ResponseBody::lossy_text_from_bytes(b"hello".to_vec()); + assert_eq!(body.as_materialized_text(), Some("hello")); + assert_eq!(body.as_materialized_bytes(), Some(b"hello".as_slice())); + let storage = body.as_materialized_bytes().unwrap().as_ptr(); + let shared = body.shared_materialized_bytes().unwrap(); + assert_eq!(shared.as_slice().as_ptr(), storage); + let cloned = body.clone_materialized().unwrap(); + assert_eq!(cloned.as_materialized_bytes().unwrap().as_ptr(), storage); + assert!(matches!( + body, + ResponseBody::MaterializedText { + exact_bytes: None, + .. + } + )); + + let (text, bytes) = body + .try_into_lossy_materialized_text() + .expect("materialized text should expose exact parts"); + assert_eq!(text, "hello"); + assert_eq!(bytes, b"hello"); +} + +#[test] +fn response_body_keeps_invalid_utf8_bytes_beside_lossy_text() { + let body = ResponseBody::lossy_text_from_bytes(vec![b'a', 0xff, b'b']); + assert_eq!(body.as_materialized_text(), Some("a\u{fffd}b")); + assert_eq!( + body.as_materialized_bytes(), + Some([b'a', 0xff, b'b'].as_slice()) + ); + let storage = body.as_materialized_bytes().unwrap().as_ptr(); + assert_eq!( + body.shared_materialized_bytes() + .unwrap() + .as_slice() + .as_ptr(), + storage + ); + assert!(matches!( + body, + ResponseBody::MaterializedText { + exact_bytes: Some(ref bytes), + .. + } if bytes.as_slice() == [b'a', 0xff, b'b'] + )); +} + #[tokio::test] async fn response_body_materializes_streaming_text_source() { let (body_tx, body_rx) = mpsc::unbounded_channel(); diff --git a/moli-page-types/src/lib.rs b/moli-page-types/src/lib.rs index b9feeab3c..d7046f3cf 100644 --- a/moli-page-types/src/lib.rs +++ b/moli-page-types/src/lib.rs @@ -37,6 +37,7 @@ use moli_dom::NodeId; use moli_fetch::{ NegotiatedHttpVersion, NetworkRequestExtraInfo, NetworkResponseExtraInfo, RedirectInfo, RequestAuth, RequestAuthScheme, RequestAuthTarget, Response, ResponseBody, ResponseHead, + SharedResponseBodyBytes, }; use moli_web_mime::is_json_module_mime; @@ -357,6 +358,12 @@ impl NavigationResponse { self.body_bytes().to_vec() } + fn shared_body_bytes(&self) -> SharedResponseBodyBytes { + self.body + .shared_materialized_bytes() + .expect("NavigationResponse body should remain materialized") + } + pub fn materialized_body(&self) -> ResponseBody { self.body .clone_materialized() @@ -383,19 +390,10 @@ impl NavigationResponse { } pub fn from_head_and_body(head: ResponseHead, body: String, body_bytes: Vec) -> Self { - Self { - final_url: head.final_url, - status: head.status, - headers: head.headers, - body: ResponseBody::materialized_text(body, body_bytes), - request_cookie_report: head.request_cookie_report, - cookie_set_reports: head.cookie_set_reports, - redirected: head.redirected, - redirect_chain: head.redirect_chain.into_iter().map(Into::into).collect(), - from_cache: head.from_cache, - negotiated_http_version: head.negotiated_http_version, - network_request_headers: None, - } + Self::from_head_and_materialized_body( + head, + ResponseBody::materialized_text(body, body_bytes), + ) } /// Headers configured on the HTTP transfer that produced this response. @@ -414,15 +412,26 @@ impl NavigationResponse { } pub fn from_head_and_materialized_body(head: ResponseHead, body: ResponseBody) -> Self { - let (body, body_bytes) = body - .try_into_lossy_materialized_text() + let body = body + .try_into_materialized_text_body() .expect("NavigationResponse body should remain materialized text"); - Self::from_head_and_body(head, body, body_bytes) + Self { + final_url: head.final_url, + status: head.status, + headers: head.headers, + body, + request_cookie_report: head.request_cookie_report, + cookie_set_reports: head.cookie_set_reports, + redirected: head.redirected, + redirect_chain: head.redirect_chain.into_iter().map(Into::into).collect(), + from_cache: head.from_cache, + negotiated_http_version: head.negotiated_http_version, + network_request_headers: None, + } } pub fn from_head_and_text_body(head: ResponseHead, body: String) -> Self { - let body_bytes = body.as_bytes().to_vec(); - Self::from_head_and_body(head, body, body_bytes) + Self::from_head_and_body(head, body, Vec::new()) } pub fn from_text_body( @@ -1459,14 +1468,14 @@ pub struct SubresourceResponseBody { #[derive(Debug)] enum SubresourceResponseBodyInner { - Memory(Vec), + Memory(SharedResponseBodyBytes), File { path: PathBuf, len: usize }, } impl SubresourceResponseBodyInner { fn in_memory_bytes(&self) -> Option<&[u8]> { match self { - Self::Memory(bytes) => Some(bytes), + Self::Memory(bytes) => Some(bytes.as_slice()), Self::File { .. } => None, } } @@ -1635,20 +1644,30 @@ impl Drop for SubresourceResponseBodyWriter { impl SubresourceResponseBody { pub fn from_bytes(bytes: Vec) -> Self { Self { - inner: Arc::new(SubresourceResponseBodyInner::Memory(bytes)), + inner: Arc::new(SubresourceResponseBodyInner::Memory( + SharedResponseBodyBytes::from_bytes(bytes), + )), } } - /// Copies the exact bytes from a materialized fetch response into the + /// Shares the exact bytes from a materialized fetch response with the /// renderer-neutral subresource body carrier. pub fn from_fetch_response(response: &Response) -> Self { - Self::from_bytes(response.body_bytes().to_vec()) + Self { + inner: Arc::new(SubresourceResponseBodyInner::Memory( + response.shared_body_bytes(), + )), + } } - /// Copies the exact bytes from a materialized navigation response into the + /// Shares the exact bytes from a materialized navigation response with the /// renderer-neutral subresource body carrier. pub fn from_navigation_response(response: &NavigationResponse) -> Self { - Self::from_bytes(response.body_bytes().to_vec()) + Self { + inner: Arc::new(SubresourceResponseBodyInner::Memory( + response.shared_body_bytes(), + )), + } } pub fn bytes(&self) -> Cow<'_, [u8]> { @@ -1663,7 +1682,7 @@ impl SubresourceResponseBody { pub fn try_bytes(&self) -> io::Result> { match self.inner.as_ref() { - SubresourceResponseBodyInner::Memory(bytes) => Ok(Cow::Borrowed(bytes)), + SubresourceResponseBodyInner::Memory(bytes) => Ok(Cow::Borrowed(bytes.as_slice())), SubresourceResponseBodyInner::File { .. } => self.materialize_bytes().map(Cow::Owned), } } @@ -1710,9 +1729,11 @@ impl SubresourceResponseBody { pub fn materialize_bytes_from(&self, offset: usize) -> io::Result> { match self.inner.as_ref() { - SubresourceResponseBodyInner::Memory(bytes) => { - Ok(bytes.get(offset..).map(<[u8]>::to_vec).unwrap_or_default()) - } + SubresourceResponseBodyInner::Memory(bytes) => Ok(bytes + .as_slice() + .get(offset..) + .map(<[u8]>::to_vec) + .unwrap_or_default()), SubresourceResponseBodyInner::File { path, len, .. } => { if offset >= *len { return Ok(Vec::new()); @@ -1732,7 +1753,7 @@ impl SubresourceResponseBody { } match self.inner.as_ref() { SubresourceResponseBodyInner::Memory(bytes) => { - let Some(remaining) = bytes.get(offset..) else { + let Some(remaining) = bytes.as_slice().get(offset..) else { return Ok(Vec::new()); }; let len = remaining.len().min(max_len); @@ -1754,7 +1775,7 @@ impl SubresourceResponseBody { pub fn write_bytes_to(&self, writer: &mut W) -> io::Result<()> { match self.inner.as_ref() { - SubresourceResponseBodyInner::Memory(bytes) => writer.write_all(bytes), + SubresourceResponseBodyInner::Memory(bytes) => writer.write_all(bytes.as_slice()), SubresourceResponseBodyInner::File { path, .. } => { let mut file = File::open(path)?; let mut buffer = [0; 64 * 1024]; @@ -1772,7 +1793,7 @@ impl SubresourceResponseBody { pub fn len(&self) -> usize { match self.inner.as_ref() { - SubresourceResponseBodyInner::Memory(bytes) => bytes.len(), + SubresourceResponseBodyInner::Memory(bytes) => bytes.as_slice().len(), SubresourceResponseBodyInner::File { len, .. } => *len, } } @@ -3875,12 +3896,14 @@ mod tests { "hello".to_owned(), b"hello".to_vec(), ); + let response_body_storage = response.body_bytes().as_ptr(); let body = SubresourceResponseBody::from_fetch_response(&response); let SubresourceResponseBodyInner::Memory(bytes) = body.inner.as_ref() else { panic!("fetch response should use in-memory byte storage"); }; - assert_eq!(bytes, b"hello"); + assert_eq!(bytes.as_slice(), b"hello"); + assert_eq!(bytes.as_slice().as_ptr(), response_body_storage); assert_eq!(body.try_bytes().unwrap().as_ref(), b"hello"); }