fix(response)!: Preserve non-UTF-8 bytes in binary success responses

3 files changed, 88 insertions, 29 deletions

diff --git a/src/response.rs b/src/response.rs
index 143ebdc..9e74773 100644
--- a/src/response.rs
+++ b/src/response.rs
@@ -15,11 +15,14 @@ macro_rules! response {
 /// The content and response type a handler should reply with.
 #[derive(Clone)]
 pub struct Response {
-  pub status:        i32,
-  pub mime:          Option<String>,
-  pub content:       String,
-  pub character_set: Option<String>,
-  pub languages:     Option<Vec<String>>,
+  pub status:         i32,
+  pub mime:           Option<String>,
+  pub content:        String,
+  /// Raw body for status `21`/`22`; the router emits these bytes verbatim
+  /// instead of `content`.
+  pub binary_content: Option<Vec<u8>>,
+  pub character_set:  Option<String>,
+  pub languages:      Option<Vec<String>>,
 }
 
 impl Response {
@@ -74,7 +77,9 @@ impl Response {
     content: impl AsRef<[u8]>,
     mime: impl Into<String> + AsRef<str>,
   ) -> Self {
-    let mut response = Self::new(21, String::from_utf8_lossy(content.as_ref()));
+    let mut response = Self::new(21, String::new());
+
+    response.binary_content = Some(content.as_ref().to_vec());
 
     response.with_mime(mime);
 
@@ -84,10 +89,12 @@ impl Response {
   #[cfg(feature = "auto-deduce-mime")]
   #[must_use]
   pub fn binary_success_auto(content: &[u8]) -> Self {
-    let mut response = Self::new(22, String::from_utf8_lossy(content));
+    let mut response = Self::new(22, String::new());
 
     response.with_mime(tree_magic_mini::from_u8(content));
 
+    response.binary_content = Some(content.to_vec());
+
     response
   }
 
@@ -97,11 +104,33 @@ impl Response {
       status,
       mime: None,
       content: content.into(),
+      binary_content: None,
       character_set: None,
       languages: None,
     }
   }
 
+  #[doc(hidden)]
+  #[must_use]
+  pub fn serialize_body(self, header: &str, footer: &str) -> Vec<u8> {
+    match self.status {
+      20 => {
+        let mut body = Vec::with_capacity(
+          header.len() + self.content.len() + footer.len() + 1,
+        );
+
+        body.extend_from_slice(header.as_bytes());
+        body.extend_from_slice(self.content.as_bytes());
+        body.push(b'\n');
+        body.extend_from_slice(footer.as_bytes());
+
+        body
+      }
+      21 | 22 => self.binary_content.unwrap_or_default(),
+      _ => Vec::new(),
+    }
+  }
+
   pub fn with_mime(
     &mut self,
     mime: impl Into<String> + AsRef<str>,
diff --git a/src/router.rs b/src/router.rs
index 81ae1a1..b6a960d 100644
--- a/src/router.rs
+++ b/src/router.rs
@@ -301,29 +301,13 @@ impl RequestHandler {
         format!("{} {}", status_code, content.content)
       }
     };
-    let body = match content.status {
-      20 => {
-        let mut body = String::with_capacity(
-          header.len() + content.content.len() + footer.len() + 1,
-        );
-
-        body.push_str(&header);
-        body.push_str(&content.content);
-        body.push('\n');
-        body.push_str(&footer);
-
-        body
-      }
-      21 | 22 => content.content,
-      _ => String::new(),
-    };
-    let mut response =
-      String::with_capacity(status_line.len() + body.len() + 2);
+    let body = content.serialize_body(&header, &footer);
+    let mut response = Vec::with_capacity(status_line.len() + body.len() + 2);
 
-    response.push_str(&status_line);
-    response.push_str("\r\n");
-    response.push_str(&body);
-    stream.write_all(response.as_bytes()).await?;
+    response.extend_from_slice(status_line.as_bytes());
+    response.extend_from_slice(b"\r\n");
+    response.extend_from_slice(&body);
+    stream.write_all(&response).await?;
     #[cfg(feature = "tokio")]
     stream.shutdown().await?;
     #[cfg(feature = "async-std")]
diff --git a/tests/binary_response.rs b/tests/binary_response.rs
new file mode 100644
index 0000000..47d3935
--- /dev/null
+++ b/tests/binary_response.rs
@@ -0,0 +1,46 @@
+use windmark::response::Response;
+
+// PNG magic + bytes invalid as UTF-8 (lone continuation bytes, lone 4-byte
+// lead, reserved C0/C1 leads, DEL).
+const NON_UTF8: &[u8] = &[
+  0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A, 0x80, 0x81, 0xFE, 0xFF, 0xC0,
+  0xC1, 0xF5, 0x00, 0x7F,
+];
+
+#[test]
+fn binary_success_preserves_non_utf8_bytes() {
+  let response = Response::binary_success(NON_UTF8, "image/png");
+
+  assert_eq!(response.status, 21);
+  assert_eq!(response.mime.as_deref(), Some("image/png"));
+  assert_eq!(response.binary_content.as_deref(), Some(NON_UTF8));
+}
+
+#[test]
+fn serialize_body_writes_binary_bytes_verbatim() {
+  let response = Response::binary_success(NON_UTF8, "image/png");
+  let body = response.serialize_body("ignored-header", "ignored-footer");
+
+  assert_eq!(body, NON_UTF8);
+}
+
+#[cfg(feature = "auto-deduce-mime")]
+#[test]
+fn binary_success_auto_preserves_non_utf8_bytes() {
+  let response = Response::binary_success_auto(NON_UTF8);
+
+  assert_eq!(response.status, 22);
+  assert_eq!(response.binary_content.as_deref(), Some(NON_UTF8));
+  assert_eq!(
+    response.serialize_body("ignored-header", "ignored-footer"),
+    NON_UTF8,
+  );
+}
+
+#[test]
+fn serialize_body_text_success_wraps_with_header_and_footer() {
+  let response = Response::success("body");
+  let body = response.serialize_body("HEADER\n", "FOOTER");
+
+  assert_eq!(body, b"HEADER\nbody\nFOOTER");
+}