1 files changed, 13 insertions, 6 deletions
diff --git a/examples/csrf/src/main.rs b/examples/csrf/src/main.rs
index 1048689..7259abd 100644
--- a/examples/csrf/src/main.rs
+++ b/examples/csrf/src/main.rs
@@ -33,11 +33,13 @@ async fn route(req: &mut Parts, body: Body) -> Result<Response, Error> {
 }
 
 fn get_form(req: &mut Parts) -> Response {
-    let mut response = Builder::new();
-    let csrf_input = req.csrf_token(&mut response).html_input();
-    response.content_type(mime::TEXT_HTML).body(
-        format!("<form method=post>
-        <input name=text>{}<button>Submit</button></form>", csrf_input).into()
+    Builder::new()
+    .content_type(mime::TEXT_HTML)
+    .body(
+        format!(
+            "<form method=post><input name=text>{}<button>Submit</button></form>",
+             req.csrf_token().html_input()
+        ).into()
     ).unwrap()
 }
 
@@ -54,7 +56,12 @@ async fn post_form(req: &mut Parts, body: Body) -> Result<Response, Error> {
 async fn service(req: hyper::Request<hyper::Body>) -> Result<hyper::Response<hyper::Body>, Infallible> {
     let (mut parts, body) = req.into_parts();
     match route(&mut parts, body).await {
-        Ok(res) => Ok(res),
+        Ok(mut res) => {
+            for (k,v) in parts.response_headers().iter() {
+                res.headers_mut().append(k, v.clone());
+            }
+            Ok(res)
+        }
         Err(err) => {
             let (code, message) = render_error(err);
             // you can easily wrap or log errors here