From d34a8689e55dac141b43ab30430404aba5d51ad0 Mon Sep 17 00:00:00 2001
From: fakeshadow <24548779@qq.com>
Date: Tue, 12 Jan 2021 22:38:53 +0800
Subject: [PATCH] Refactor h1 encoder (#1900)

---
 actix-http/src/h1/encoder.rs | 48 +++++++++++++++++++++---------------
 1 file changed, 28 insertions(+), 20 deletions(-)

diff --git a/actix-http/src/h1/encoder.rs b/actix-http/src/h1/encoder.rs
index 4427174ec..bb89905fb 100644
--- a/actix-http/src/h1/encoder.rs
+++ b/actix-http/src/h1/encoder.rs
@@ -8,7 +8,7 @@ use bytes::{BufMut, BytesMut};
 
 use crate::body::BodySize;
 use crate::config::ServiceConfig;
-use crate::header::map;
+use crate::header::{map::Value, HeaderName};
 use crate::helpers;
 use crate::http::header::{CONNECTION, CONTENT_LENGTH, DATE, TRANSFER_ENCODING};
 use crate::http::{HeaderMap, StatusCode, Version};
@@ -121,21 +121,11 @@ pub(crate) trait MessageType: Sized {
             _ => {}
         }
 
-        // merging headers from head and extra headers. HeaderMap::new() does not allocate.
-        let empty_headers = HeaderMap::new();
-        let extra_headers = self.extra_headers().unwrap_or(&empty_headers);
-        let headers = self
-            .headers()
-            .inner
-            .iter()
-            .filter(|(name, _)| !extra_headers.contains_key(*name))
-            .chain(extra_headers.inner.iter());
-
         // write headers
 
         let mut has_date = false;
 
-        let mut buf = dst.chunk_mut().as_mut_ptr() as *mut u8;
+        let mut buf = dst.chunk_mut().as_mut_ptr();
         let mut remaining = dst.capacity() - dst.len();
 
         // tracks bytes written since last buffer resize
@@ -143,10 +133,10 @@ pub(crate) trait MessageType: Sized {
         // container's knowledge, this is used to sync the containers cursor after data is written
         let mut pos = 0;
 
-        for (key, value) in headers {
+        self.write_headers(|key, value| {
             match *key {
-                CONNECTION => continue,
-                TRANSFER_ENCODING | CONTENT_LENGTH if skip_len => continue,
+                CONNECTION => return,
+                TRANSFER_ENCODING | CONTENT_LENGTH if skip_len => return,
                 DATE => has_date = true,
                 _ => {}
             }
@@ -155,7 +145,7 @@ pub(crate) trait MessageType: Sized {
             let k_len = k.len();
 
             match value {
-                map::Value::One(ref val) => {
+                Value::One(ref val) => {
                     let v = val.as_ref();
                     let v_len = v.len();
 
@@ -177,7 +167,7 @@ pub(crate) trait MessageType: Sized {
 
                         // re-assign buf raw pointer since it's possible that the buffer was
                         // reallocated and/or resized
-                        buf = dst.chunk_mut().as_mut_ptr() as *mut u8;
+                        buf = dst.chunk_mut().as_mut_ptr();
                     }
 
                     // SAFETY: on each write, it is enough to ensure that the advancement of the
@@ -206,7 +196,7 @@ pub(crate) trait MessageType: Sized {
                     remaining -= len;
                 }
 
-                map::Value::Multi(ref vec) => {
+                Value::Multi(ref vec) => {
                     for val in vec {
                         let v = val.as_ref();
                         let v_len = v.len();
@@ -224,7 +214,7 @@ pub(crate) trait MessageType: Sized {
 
                             // re-assign buf raw pointer since it's possible that the buffer was
                             // reallocated and/or resized
-                            buf = dst.chunk_mut().as_mut_ptr() as *mut u8;
+                            buf = dst.chunk_mut().as_mut_ptr();
                         }
 
                         // SAFETY: on each write, it is enough to ensure that the advancement of
@@ -253,7 +243,7 @@ pub(crate) trait MessageType: Sized {
                     }
                 }
             }
-        }
+        });
 
         // final cursor synchronization with the bytes container
         //
@@ -273,6 +263,24 @@ pub(crate) trait MessageType: Sized {
 
         Ok(())
     }
+
+    fn write_headers<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&HeaderName, &Value),
+    {
+        match self.extra_headers() {
+            Some(headers) => {
+                // merging headers from head and extra headers.
+                self.headers()
+                    .inner
+                    .iter()
+                    .filter(|(name, _)| !headers.contains_key(*name))
+                    .chain(headers.inner.iter())
+                    .for_each(|(k, v)| f(k, v))
+            }
+            None => self.headers().inner.iter().for_each(|(k, v)| f(k, v)),
+        }
+    }
 }
 
 impl MessageType for Response<()> {