From 9095d8349dc476cb104562d23f790399ecb5afeb Mon Sep 17 00:00:00 2001
From: austaras <austaras@outlook.com>
Date: Tue, 24 Nov 2020 19:19:55 +0800
Subject: [PATCH] add custom reponse builder

---
 src/handshake/server.rs | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/handshake/server.rs b/src/handshake/server.rs
index 1b6eed8..298832e 100644
--- a/src/handshake/server.rs
+++ b/src/handshake/server.rs
@@ -7,6 +7,7 @@ use std::{
 };
 
 use http::{HeaderMap, Request as HttpRequest, Response as HttpResponse, StatusCode};
+use http::response::Builder;
 use httparse::Status;
 use log::*;
 
@@ -30,8 +31,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()));
     }
@@ -75,8 +75,24 @@ pub fn create_response(request: &Request) -> Result<Response> {
         .header("Connection", "Upgrade")
         .header("Upgrade", "websocket")
         .header("Sec-WebSocket-Accept", convert_key(key.as_bytes())?);
+    
+    Ok(builder)
+}
 
-    Ok(builder.body(())?)
+/// Create a response for the request.
+pub fn create_response(request: &Request) -> Result<Response> {
+    match create_parts(&request) {
+        Ok(builder) => Ok(builder.body(())?),
+        Err(e) => Err(e)
+    }
+}
+
+/// Create a response for the request, with a custom body builder
+pub fn create_response_t<T>(request: &HttpRequest<T>, empty: impl FnOnce() -> T) -> Result<HttpResponse<T>> {
+    match create_parts(&request) {
+        Ok(builder) => Ok(builder.body(empty())?),
+        Err(e) => Err(e)
+    }
 }
 
 // Assumes that this is a valid response