From 4d70f63cbf7b69a0125dd42877d57e198e5b9977 Mon Sep 17 00:00:00 2001
From: Austaras <austaras@outlook.com>
Date: Thu, 31 Dec 2020 00:35:10 +0800
Subject: [PATCH] Add a custom response builder (#158)

* add custom reponse builder

* change names and shorten implemtation

* Re-format the code

Co-authored-by: Daniel Abramov <dabramov@snapview.de>
---
 src/client.rs           |  6 ++----
 src/handshake/server.rs | 22 ++++++++++++++++++----
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 9b13be9..70c157c 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -119,10 +119,8 @@ pub fn connect_with_config<Req: IntoClientRequest>(
     }
 
     fn create_request(parts: &Parts, uri: &Uri) -> Request {
-        let mut builder = Request::builder()
-            .uri(uri.clone())
-            .method(parts.method.clone())
-            .version(parts.version);
+        let mut builder =
+            Request::builder().uri(uri.clone()).method(parts.method.clone()).version(parts.version);
         *builder.headers_mut().expect("Failed to create `Request`") = parts.headers.clone();
         builder.body(()).expect("Failed to create `Request`")
     }
diff --git a/src/handshake/server.rs b/src/handshake/server.rs
index 1b6eed8..e60acb5 100644
--- a/src/handshake/server.rs
+++ b/src/handshake/server.rs
@@ -6,7 +6,9 @@ use std::{
     result::Result as StdResult,
 };
 
-use http::{HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode};
+use http::{
+    response::Builder, HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode,
+};
 use httparse::Status;
 use log::*;
 
@@ -30,8 +32,7 @@ pub type Response = HttpResponse<()>;
 /// Server error response type.
 pub type ErrorResponse = HttpResponse<Option<String>>;
 
-/// Create a response for the request.
-pub fn create_response(request: &Request) -> Result<Response> {
+fn create_parts<T>(request: &HttpRequest<T>) -> Result<Builder> {
     if request.method() != http::Method::GET {
         return Err(Error::Protocol("Method is not GET".into()));
     }
@@ -76,7 +77,20 @@ pub fn create_response(request: &Request) -> Result<Response> {
         .header("Upgrade", "websocket")
         .header("Sec-WebSocket-Accept", convert_key(key.as_bytes())?);
 
-    Ok(builder.body(())?)
+    Ok(builder)
+}
+
+/// Create a response for the request.
+pub fn create_response(request: &Request) -> Result<Response> {
+    Ok(create_parts(&request)?.body(())?)
+}
+
+/// Create a response for the request with a custom body.
+pub fn create_response_with_body<T>(
+    request: &HttpRequest<T>,
+    generate_body: impl FnOnce() -> T,
+) -> Result<HttpResponse<T>> {
+    Ok(create_parts(&request)?.body(generate_body())?)
 }
 
 // Assumes that this is a valid response