refactor(ffi): Add HeaderCaseMap preserving http1 header casing

seanmonstar · seanmonstar · commit 7d7eeb0ce589 · 2021-01-07T15:50:59.000-08:00
diff --git a/capi/examples/client.c b/capi/examples/client.c
@@ -228,7 +228,7 @@ int main(int argc, char *argv[]) {
                 }
 
                 hyper_headers *req_headers = hyper_request_headers(req);
-                hyper_headers_set(req_headers,  STR_ARG("host"), STR_ARG(host));
+                hyper_headers_set(req_headers,  STR_ARG("Host"), STR_ARG(host));
 
                 // Send it!
                 hyper_task *send = hyper_clientconn_send(client, req);
diff --git a/src/ffi/body.rs b/src/ffi/body.rs
@@ -24,7 +24,7 @@ pub(crate) struct UserBody {
 type hyper_body_foreach_callback = extern "C" fn(*mut c_void, *const hyper_buf) -> c_int;
 
 type hyper_body_data_callback =
-    extern "C" fn(*mut c_void, *mut hyper_context, *mut *mut hyper_buf) -> c_int;
+    extern "C" fn(*mut c_void, *mut hyper_context<'_>, *mut *mut hyper_buf) -> c_int;
 
 ffi_fn! {
     /// Create a new "empty" body.
diff --git a/src/ffi/client.rs b/src/ffi/client.rs
@@ -67,11 +67,15 @@ ffi_fn! {
             return std::ptr::null_mut();
         }
 
-        let req = unsafe { Box::from_raw(req) };
+        let mut req = unsafe { Box::from_raw(req) };
+
+        // Update request with original-case map of headers
+        req.finalize_request();
+
         let fut = unsafe { &mut *conn }.tx.send_request(req.0);
 
         let fut = async move {
-            fut.await.map(hyper_response)
+            fut.await.map(hyper_response::wrap)
         };
 
         Box::into_raw(Task::boxed(fut))
diff --git a/src/ffi/error.rs b/src/ffi/error.rs
@@ -33,7 +33,7 @@ impl hyper_error {
             ErrorKind::IncompleteMessage => hyper_code::HYPERE_UNEXPECTED_EOF,
             ErrorKind::User(User::AbortedByCallback) => hyper_code::HYPERE_ABORTED_BY_CALLBACK,
             // TODO: add more variants
-            _ => hyper_code::HYPERE_ERROR
+            _ => hyper_code::HYPERE_ERROR,
         }
     }
 
diff --git a/src/ffi/http_types.rs b/src/ffi/http_types.rs
@@ -1,3 +1,4 @@
+use bytes::Bytes;
 use libc::{c_int, size_t};
 use std::ffi::c_void;
 
@@ -8,13 +9,21 @@ use super::HYPER_ITER_CONTINUE;
 use crate::header::{HeaderName, HeaderValue};
 use crate::{Body, HeaderMap, Method, Request, Response, Uri};
 
-// ===== impl Request =====
-
 pub struct hyper_request(pub(super) Request<Body>);
 
 pub struct hyper_response(pub(super) Response<Body>);
 
-pub struct hyper_headers(pub(super) HeaderMap);
+#[derive(Default)]
+pub struct hyper_headers {
+    pub(super) headers: HeaderMap,
+    orig_casing: HeaderCaseMap,
+}
+
+// Will probably be moved to `hyper::ext::http1`
+#[derive(Debug, Default)]
+pub(crate) struct HeaderCaseMap(HeaderMap<Bytes>);
+
+// ===== impl hyper_request =====
 
 ffi_fn! {
     /// Construct a new HTTP request.
@@ -96,7 +105,7 @@ ffi_fn! {
     /// This is not an owned reference, so it should not be accessed after the
     /// `hyper_request` has been consumed.
     fn hyper_request_headers(req: *mut hyper_request) -> *mut hyper_headers {
-        hyper_headers::wrap(unsafe { &mut *req }.0.headers_mut())
+        hyper_headers::get_or_default(unsafe { &mut *req }.0.extensions_mut())
     }
 }
 
@@ -114,7 +123,16 @@ ffi_fn! {
     }
 }
 
-// ===== impl Response =====
+impl hyper_request {
+    pub(super) fn finalize_request(&mut self) {
+        if let Some(headers) = self.0.extensions_mut().remove::<hyper_headers>() {
+            *self.0.headers_mut() = headers.headers;
+            self.0.extensions_mut().insert(headers.orig_casing);
+        }
+    }
+}
+
+// ===== impl hyper_response =====
 
 ffi_fn! {
     /// Free an HTTP response after using it.
@@ -159,7 +177,7 @@ ffi_fn! {
     /// This is not an owned reference, so it should not be accessed after the
     /// `hyper_response` has been freed.
     fn hyper_response_headers(resp: *mut hyper_response) -> *mut hyper_headers {
-        hyper_headers::wrap(unsafe { &mut *resp }.0.headers_mut())
+        hyper_headers::get_or_default(unsafe { &mut *resp }.0.extensions_mut())
     }
 }
 
@@ -173,6 +191,22 @@ ffi_fn! {
     }
 }
 
+impl hyper_response {
+    pub(super) fn wrap(mut resp: Response<Body>) -> hyper_response {
+        let headers = std::mem::take(resp.headers_mut());
+        let orig_casing = resp
+            .extensions_mut()
+            .remove::<HeaderCaseMap>()
+            .unwrap_or_default();
+        resp.extensions_mut().insert(hyper_headers {
+            headers,
+            orig_casing,
+        });
+
+        hyper_response(resp)
+    }
+}
+
 unsafe impl AsTaskType for hyper_response {
     fn as_task_type(&self) -> hyper_task_return_type {
         hyper_task_return_type::HYPER_TASK_RESPONSE
@@ -185,9 +219,15 @@ type hyper_headers_foreach_callback =
     extern "C" fn(*mut c_void, *const u8, size_t, *const u8, size_t) -> c_int;
 
 impl hyper_headers {
-    pub(crate) fn wrap(cx: &mut HeaderMap) -> &mut hyper_headers {
-        // A struct with only one field has the same layout as that field.
-        unsafe { std::mem::transmute::<&mut HeaderMap, &mut hyper_headers>(cx) }
+    pub(super) fn get_or_default(ext: &mut http::Extensions) -> &mut hyper_headers {
+        if let None = ext.get_mut::<hyper_headers>() {
+            ext.insert(hyper_headers {
+                headers: Default::default(),
+                orig_casing: Default::default(),
+            });
+        }
+
+        ext.get_mut::<hyper_headers>().unwrap()
     }
 }
 
@@ -199,14 +239,31 @@ ffi_fn! {
     /// The callback should return `HYPER_ITER_CONTINUE` to keep iterating, or
     /// `HYPER_ITER_BREAK` to stop.
     fn hyper_headers_foreach(headers: *const hyper_headers, func: hyper_headers_foreach_callback, userdata: *mut c_void) {
-        for (name, value) in unsafe { &*headers }.0.iter() {
-            let name_ptr = name.as_str().as_bytes().as_ptr();
-            let name_len = name.as_str().as_bytes().len();
-            let val_ptr = value.as_bytes().as_ptr();
-            let val_len = value.as_bytes().len();
-
-            if HYPER_ITER_CONTINUE != func(userdata, name_ptr, name_len, val_ptr, val_len) {
-                break;
+        let headers = unsafe { &*headers };
+        // For each header name/value pair, there may be a value in the casemap
+        // that corresponds to the HeaderValue. So, we iterator all the keys,
+        // and for each one, try to pair the originally cased name with the value.
+        //
+        // TODO: consider adding http::HeaderMap::entries() iterator
+        for name in headers.headers.keys() {
+            let mut names = headers.orig_casing.get_all(name).iter();
+
+            for value in headers.headers.get_all(name) {
+                let (name_ptr, name_len) = if let Some(orig_name) = names.next() {
+                    (orig_name.as_ptr(), orig_name.len())
+                } else {
+                    (
+                        name.as_str().as_bytes().as_ptr(),
+                        name.as_str().as_bytes().len(),
+                    )
+                };
+
+                let val_ptr = value.as_bytes().as_ptr();
+                let val_len = value.as_bytes().len();
+
+                if HYPER_ITER_CONTINUE != func(userdata, name_ptr, name_len, val_ptr, val_len) {
+                    return;
+                }
             }
         }
     }
@@ -219,8 +276,9 @@ ffi_fn! {
     fn hyper_headers_set(headers: *mut hyper_headers, name: *const u8, name_len: size_t, value: *const u8, value_len: size_t) -> hyper_code {
         let headers = unsafe { &mut *headers };
         match unsafe { raw_name_value(name, name_len, value, value_len) } {
-            Ok((name, value)) => {
-                headers.0.insert(name, value);
+            Ok((name, value, orig_name)) => {
+                headers.headers.insert(&name, value);
+                headers.orig_casing.insert(name, orig_name);
                 hyper_code::HYPERE_OK
             }
             Err(code) => code,
@@ -237,8 +295,9 @@ ffi_fn! {
         let headers = unsafe { &mut *headers };
 
         match unsafe { raw_name_value(name, name_len, value, value_len) } {
-            Ok((name, value)) => {
-                headers.0.append(name, value);
+            Ok((name, value, orig_name)) => {
+                headers.headers.append(&name, value);
+                headers.orig_casing.append(name, orig_name);
                 hyper_code::HYPERE_OK
             }
             Err(code) => code,
@@ -251,8 +310,9 @@ unsafe fn raw_name_value(
     name_len: size_t,
     value: *const u8,
     value_len: size_t,
-) -> Result<(HeaderName, HeaderValue), hyper_code> {
+) -> Result<(HeaderName, HeaderValue, Bytes), hyper_code> {
     let name = std::slice::from_raw_parts(name, name_len);
+    let orig_name = Bytes::copy_from_slice(name);
     let name = match HeaderName::from_bytes(name) {
         Ok(name) => name,
         Err(_) => return Err(hyper_code::HYPERE_INVALID_ARG),
@@ -263,5 +323,78 @@ unsafe fn raw_name_value(
         Err(_) => return Err(hyper_code::HYPERE_INVALID_ARG),
     };
 
-    Ok((name, value))
+    Ok((name, value, orig_name))
+}
+
+// ===== impl HeaderCaseMap =====
+
+impl HeaderCaseMap {
+    pub(crate) fn get_all(&self, name: &HeaderName) -> http::header::GetAll<'_, Bytes> {
+        self.0.get_all(name)
+    }
+
+    pub(crate) fn insert(&mut self, name: HeaderName, orig: Bytes) {
+        self.0.insert(name, orig);
+    }
+
+    pub(crate) fn append<N>(&mut self, name: N, orig: Bytes)
+    where
+        N: http::header::IntoHeaderName,
+    {
+        self.0.append(name, orig);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_headers_foreach_cases_preserved() {
+        let mut headers = hyper_headers::default();
+
+        let name1 = b"Set-CookiE";
+        let value1 = b"a=b";
+        hyper_headers_add(
+            &mut headers,
+            name1.as_ptr(),
+            name1.len(),
+            value1.as_ptr(),
+            value1.len(),
+        );
+
+        let name2 = b"SET-COOKIE";
+        let value2 = b"c=d";
+        hyper_headers_add(
+            &mut headers,
+            name2.as_ptr(),
+            name2.len(),
+            value2.as_ptr(),
+            value2.len(),
+        );
+
+        let mut vec = Vec::<u8>::new();
+        hyper_headers_foreach(&headers, concat, &mut vec as *mut _ as *mut c_void);
+
+        assert_eq!(vec, b"Set-CookiE: a=b\r\nSET-COOKIE: c=d\r\n");
+
+        extern "C" fn concat(
+            vec: *mut c_void,
+            name: *const u8,
+            name_len: usize,
+            value: *const u8,
+            value_len: usize,
+        ) -> c_int {
+            unsafe {
+                let vec = &mut *(vec as *mut Vec<u8>);
+                let name = std::slice::from_raw_parts(name, name_len);
+                let value = std::slice::from_raw_parts(value, value_len);
+                vec.extend(name);
+                vec.extend(b": ");
+                vec.extend(value);
+                vec.extend(b"\r\n");
+            }
+            HYPER_ITER_CONTINUE
+        }
+    }
 }
diff --git a/src/ffi/mod.rs b/src/ffi/mod.rs
@@ -28,6 +28,7 @@ mod io;
 mod task;
 
 pub(crate) use self::body::UserBody;
+pub(crate) use self::http_types::HeaderCaseMap;
 
 pub const HYPER_ITER_CONTINUE: libc::c_int = 0;
 #[allow(unused)]
diff --git a/src/proto/h1/conn.rs b/src/proto/h1/conn.rs
@@ -44,6 +44,8 @@ where
                 error: None,
                 keep_alive: KA::Busy,
                 method: None,
+                #[cfg(feature = "ffi")]
+                preserve_header_case: false,
                 title_case_headers: false,
                 notify_read: false,
                 reading: Reading::Init,
@@ -142,6 +144,8 @@ where
             ParseContext {
                 cached_headers: &mut self.state.cached_headers,
                 req_method: &mut self.state.method,
+                #[cfg(feature = "ffi")]
+                preserve_header_case: self.state.preserve_header_case,
             }
         )) {
             Ok(msg) => msg,
@@ -474,6 +478,16 @@ where
 
         self.enforce_version(&mut head);
 
+        // Maybe check if we should preserve header casing on received
+        // message headers...
+        #[cfg(feature = "ffi")]
+        {
+            if T::is_client() && !self.state.preserve_header_case {
+                self.state.preserve_header_case =
+                    head.extensions.get::<crate::ffi::HeaderCaseMap>().is_some();
+            }
+        }
+
         let buf = self.io.headers_buf();
         match super::role::encode_headers::<T>(
             Encode {
@@ -736,6 +750,8 @@ struct State {
     /// This is used to know things such as if the message can include
     /// a body or not.
     method: Option<Method>,
+    #[cfg(feature = "ffi")]
+    preserve_header_case: bool,
     title_case_headers: bool,
     /// Set to true when the Dispatcher should poll read operations
     /// again. See the `maybe_notify` method for more.
diff --git a/src/proto/h1/dispatch.rs b/src/proto/h1/dispatch.rs
@@ -492,7 +492,7 @@ cfg_server! {
                     version: parts.version,
                     subject: parts.status,
                     headers: parts.headers,
-                    extensions: http::Extensions::default(),
+                    extensions: parts.extensions,
                 };
                 Poll::Ready(Some(Ok((head, body))))
             } else {
@@ -576,7 +576,7 @@ cfg_client! {
                                 version: parts.version,
                                 subject: crate::proto::RequestLine(parts.method, parts.uri),
                                 headers: parts.headers,
-                                extensions: http::Extensions::default(),
+                                extensions: parts.extensions,
                             };
                             *this.callback = Some(cb);
                             Poll::Ready(Some(Ok((head, body))))
diff --git a/src/proto/h1/io.rs b/src/proto/h1/io.rs
@@ -159,6 +159,8 @@ where
                 ParseContext {
                     cached_headers: parse_ctx.cached_headers,
                     req_method: parse_ctx.req_method,
+                    #[cfg(feature = "ffi")]
+                    preserve_header_case: parse_ctx.preserve_header_case,
                 },
             )? {
                 Some(msg) => {
@@ -636,6 +638,8 @@ mod tests {
             let parse_ctx = ParseContext {
                 cached_headers: &mut None,
                 req_method: &mut None,
+                #[cfg(feature = "ffi")]
+                preserve_header_case: false,
             };
             assert!(buffered
                 .parse::<ClientTransaction>(cx, parse_ctx)
diff --git a/src/proto/h1/mod.rs b/src/proto/h1/mod.rs
diff --git a/src/proto/h1/role.rs b/src/proto/h1/role.rs
diff --git a/src/proto/mod.rs b/src/proto/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -228,7 +228,7 @@ int main(int argc, char *argv[]) {`
`228`	`228`	`}`
`229`	`229`
`230`	`230`	`hyper_headers *req_headers = hyper_request_headers(req);`
`231`		`- hyper_headers_set(req_headers, STR_ARG("host"), STR_ARG(host));`
	`231`	`+ hyper_headers_set(req_headers, STR_ARG("Host"), STR_ARG(host));`
`232`	`232`
`233`	`233`	`// Send it!`
`234`	`234`	`hyper_task *send = hyper_clientconn_send(client, req);`
Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@ impl hyper_error {`
`33`	`33`	`ErrorKind::IncompleteMessage => hyper_code::HYPERE_UNEXPECTED_EOF,`
`34`	`34`	`ErrorKind::User(User::AbortedByCallback) => hyper_code::HYPERE_ABORTED_BY_CALLBACK,`
`35`	`35`	`// TODO: add more variants`
`36`		`- _ => hyper_code::HYPERE_ERROR`
	`36`	`+ _ => hyper_code::HYPERE_ERROR,`
`37`	`37`	`}`
`38`	`38`	`}`
`39`	`39`