From 36ab77005952a1f6aecc0b1f41c7e0cbf0282172 Mon Sep 17 00:00:00 2001
From: Josiah Bull <josiah.bull7@gmail.com>
Date: Tue, 30 Aug 2022 09:48:36 +1200
Subject: [PATCH] feat: error type returns `Vec<u8>` instead of
 `Option<String>`

---
 examples/callback-error.rs |  2 +-
 src/error.rs               |  2 +-
 src/handshake/client.rs    | 17 ++++++++++++-----
 src/handshake/server.rs    |  6 ++----
 4 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/examples/callback-error.rs b/examples/callback-error.rs
index cf78a2e..ee51136 100644
--- a/examples/callback-error.rs
+++ b/examples/callback-error.rs
@@ -13,7 +13,7 @@ fn main() {
             let callback = |_req: &Request, _resp| {
                 let resp = Response::builder()
                     .status(StatusCode::FORBIDDEN)
-                    .body(Some("Access denied".into()))
+                    .body("Access denied".into())
                     .unwrap();
                 Err(resp)
             };
diff --git a/src/error.rs b/src/error.rs
index 3e15c7c..9d6576e 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -65,7 +65,7 @@ pub enum Error {
     /// HTTP error.
     #[error("HTTP error: {}", .0.status())]
     #[cfg(feature = "handshake")]
-    Http(Response<Option<String>>),
+    Http(Response<Vec<u8>>),
     /// HTTP format error.
     #[error("HTTP format error: {0}")]
     #[cfg(feature = "handshake")]
diff --git a/src/handshake/client.rs b/src/handshake/client.rs
index 32a3c40..87305f4 100644
--- a/src/handshake/client.rs
+++ b/src/handshake/client.rs
@@ -26,7 +26,7 @@ use crate::{
 pub type Request = HttpRequest<()>;
 
 /// Client response type.
-pub type Response = HttpResponse<Option<String>>;
+pub type Response = HttpResponse<Vec<u8>>;
 
 /// Client handshake role.
 #[derive(Debug)]
@@ -83,7 +83,15 @@ impl<S: Read + Write> HandshakeRole for ClientHandshake<S> {
                 ProcessingResult::Continue(HandshakeMachine::start_read(stream))
             }
             StageResult::DoneReading { stream, result, tail } => {
-                let result = self.verify_data.verify_response(result, &tail)?;
+                let result = match self.verify_data.verify_response(result) {
+                    Ok(r) => r,
+                    Err(Error::Http(mut e)) => {
+                        *e.body_mut() = tail;
+                        return Err(Error::Http(e))
+                    },
+                    Err(e) => return Err(e),
+                };
+
                 debug!("Client handshake done.");
                 let websocket =
                     WebSocket::from_partially_read(stream, tail, Role::Client, self.config);
@@ -178,11 +186,10 @@ struct VerifyData {
 }
 
 impl VerifyData {
-    pub fn verify_response(&self, mut response: Response, tail: &[u8]) -> Result<Response> {
+    pub fn verify_response(&self, response: Response) -> Result<Response> {
         // 1. If the status code received from the server is not 101, the
         // client handles the response per HTTP [RFC2616] procedures. (RFC 6455)
         if response.status() != StatusCode::SWITCHING_PROTOCOLS {
-            *response.body_mut() = Some(String::from_utf8_lossy(tail).to_string());
             return Err(Error::Http(response));
         }
 
@@ -256,7 +263,7 @@ impl<'h, 'b: 'h> FromHttparse<httparse::Response<'h, 'b>> for Response {
 
         let headers = HeaderMap::from_httparse(raw.headers)?;
 
-        let mut response = Response::new(None);
+        let mut response = Response::new(vec![]);
         *response.status_mut() = StatusCode::from_u16(raw.code.expect("Bug: no HTTP status code"))?;
         *response.headers_mut() = headers;
         // TODO: httparse only supports HTTP 0.9/1.0/1.1 but not HTTP 2.0
diff --git a/src/handshake/server.rs b/src/handshake/server.rs
index 2f095ce..7e3a57f 100644
--- a/src/handshake/server.rs
+++ b/src/handshake/server.rs
@@ -30,7 +30,7 @@ pub type Request = HttpRequest<()>;
 pub type Response = HttpResponse<()>;
 
 /// Server error response type.
-pub type ErrorResponse = HttpResponse<Option<String>>;
+pub type ErrorResponse = HttpResponse<Vec<u8>>;
 
 fn create_parts<T>(request: &HttpRequest<T>) -> Result<Builder> {
     if request.method() != http::Method::GET {
@@ -265,9 +265,7 @@ impl<S: Read + Write, C: Callback> HandshakeRole for ServerHandshake<S, C> {
                         let mut output = vec![];
                         write_response(&mut output, resp)?;
 
-                        if let Some(body) = resp.body() {
-                            output.extend_from_slice(body.as_bytes());
-                        }
+                        output.extend_from_slice(resp.body());
 
                         ProcessingResult::Continue(HandshakeMachine::start_write(stream, output))
                     }