diff --git a/moli-fetch/src/response.rs b/moli-fetch/src/response.rs index a7f7b73439..3956999799 100644 --- a/moli-fetch/src/response.rs +++ b/moli-fetch/src/response.rs @@ -52,7 +52,10 @@ pub struct ResponseHead { #[derive(Debug)] pub enum ResponseBody { - MaterializedText { text: String, bytes: Vec }, + MaterializedText { + text: String, + exact_bytes: Option>, + }, MaterializedBytes(Vec), StreamingText(Box), StreamingBytes(Box), @@ -60,7 +63,27 @@ pub enum ResponseBody { impl ResponseBody { pub fn materialized_text(text: String, bytes: Vec) -> Self { - Self::MaterializedText { text, bytes } + let exact_bytes = (!bytes.is_empty()).then_some(bytes); + Self::MaterializedText { 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, + exact_bytes: None, + }, + Err(error) => { + let bytes = error.into_bytes(); + let text = String::from_utf8_lossy(&bytes).into_owned(); + Self::MaterializedText { + text, + exact_bytes: Some(bytes), + } + } + } } pub fn materialized_bytes(bytes: Vec) -> Self { @@ -74,12 +97,8 @@ 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::MaterializedText { text, exact_bytes } => { + Ok(exact_bytes.unwrap_or_else(|| text.into_bytes())) } Self::StreamingText(_) | Self::StreamingBytes(_) => Err(self), } @@ -88,12 +107,8 @@ 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().unwrap_or(text.as_bytes())) } Self::StreamingText(_) | Self::StreamingBytes(_) => None, } @@ -102,9 +117,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,7 +134,10 @@ 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.unwrap_or_else(|| text.as_bytes().to_vec()); + Ok((text, bytes)) + } Self::MaterializedBytes(bytes) => { let text = String::from_utf8_lossy(&bytes).into_owned(); Ok((text, bytes)) @@ -128,6 +146,15 @@ 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(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 @@ -135,12 +162,8 @@ impl ResponseBody { 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::MaterializedText { text, exact_bytes } => { + Ok(exact_bytes.unwrap_or_else(|| text.into_bytes())) } Self::StreamingText(response) => { let mut response = *response; @@ -167,7 +190,10 @@ 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.unwrap_or_else(|| text.as_bytes().to_vec()); + Ok((text, bytes)) + } Self::MaterializedBytes(bytes) => { let text = String::from_utf8_lossy(&bytes).into_owned(); Ok((text, bytes)) @@ -191,6 +217,32 @@ 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(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, + exact_bytes: None, + }) + } + Self::StreamingBytes(response) => { + let bytes = Self::StreamingBytes(response) + .into_materialized_bytes() + .await?; + Ok(Self::lossy_text_from_bytes(bytes)) + } + } + } } #[derive(Debug)] @@ -289,26 +341,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 +371,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 c469fdeee2..9724c1bcb1 100644 --- a/moli-fetch/src/tests/mod.rs +++ b/moli-fetch/src/tests/mod.rs @@ -247,6 +247,43 @@ 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())); + 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()) + ); + assert!(matches!( + body, + ResponseBody::MaterializedText { + exact_bytes: Some(ref bytes), + .. + } if bytes == &[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 b9feeab3cd..687fe3a75b 100644 --- a/moli-page-types/src/lib.rs +++ b/moli-page-types/src/lib.rs @@ -383,19 +383,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 +405,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(