From d2f5065c8f14f42076537039e10bf22c68f9fe38 Mon Sep 17 00:00:00 2001
From: Florin Lipan <florinlipan@gmail.com>
Date: Wed, 1 Nov 2023 20:28:28 +0100
Subject: [PATCH] Introduce a local IntoHeaderName trait to avoid exposing
 hyper::header::HeaderName

---
 src/error.rs |  3 +++
 src/lib.rs   |  2 +-
 src/mock.rs  | 48 +++++++++++++++++++++++++++++++++++++++++++++---
 tests/lib.rs |  3 +--
 4 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/error.rs b/src/error.rs
index f1daeb5..ee0955b 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -62,6 +62,8 @@ pub enum ErrorKind {
     ResponseBodyFailure,
     /// File not found
     FileNotFound,
+    /// Invalid header name
+    InvalidHeaderName,
 }
 
 impl ErrorKind {
@@ -75,6 +77,7 @@ impl ErrorKind {
             ErrorKind::RequestBodyFailure => "failed to read the request body",
             ErrorKind::ResponseBodyFailure => "failed to write the response body",
             ErrorKind::FileNotFound => "file not found",
+            ErrorKind::InvalidHeaderName => "invalid header name",
         }
     }
 }
diff --git a/src/lib.rs b/src/lib.rs
index d639cab..efc0912 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -672,7 +672,7 @@
 pub use error::{Error, ErrorKind};
 #[allow(deprecated)]
 pub use matcher::Matcher;
-pub use mock::Mock;
+pub use mock::{IntoHeaderName, Mock};
 pub use request::Request;
 pub use server::Server;
 pub use server_pool::ServerGuard;
diff --git a/src/mock.rs b/src/mock.rs
index fb1629d..ff4dd92 100644
--- a/src/mock.rs
+++ b/src/mock.rs
@@ -5,7 +5,7 @@ use crate::server::RemoteMock;
 use crate::server::State;
 use crate::Request;
 use crate::{Error, ErrorKind};
-use hyper::header::IntoHeaderName;
+use hyper::header::HeaderName;
 use hyper::HeaderMap;
 use hyper::StatusCode;
 use rand::distributions::Alphanumeric;
@@ -19,6 +19,42 @@ use std::string::ToString;
 use std::sync::Arc;
 use std::sync::RwLock;
 
+#[allow(missing_docs)]
+pub trait IntoHeaderName {
+    #[track_caller]
+    fn into_header_name(self) -> HeaderName;
+}
+
+impl IntoHeaderName for String {
+    fn into_header_name(self) -> HeaderName {
+        HeaderName::try_from(self)
+            .map_err(|_| Error::new(ErrorKind::InvalidHeaderName))
+            .unwrap()
+    }
+}
+
+impl IntoHeaderName for &String {
+    fn into_header_name(self) -> HeaderName {
+        HeaderName::try_from(self)
+            .map_err(|_| Error::new(ErrorKind::InvalidHeaderName))
+            .unwrap()
+    }
+}
+
+impl IntoHeaderName for &str {
+    fn into_header_name(self) -> HeaderName {
+        HeaderName::try_from(self)
+            .map_err(|_| Error::new(ErrorKind::InvalidHeaderName))
+            .unwrap()
+    }
+}
+
+impl IntoHeaderName for HeaderName {
+    fn into_header_name(self) -> HeaderName {
+        self
+    }
+}
+
 #[derive(Clone, Debug)]
 pub struct InnerMock {
     pub(crate) id: String,
@@ -202,8 +238,11 @@ impl Mock {
     ///   .match_header("authorization", "password");
     /// ```
     ///
+    #[track_caller]
     pub fn match_header<T: IntoHeaderName, M: Into<Matcher>>(mut self, field: T, value: M) -> Self {
-        self.inner.headers.append(field, value.into());
+        self.inner
+            .headers
+            .append(field.into_header_name(), value.into());
 
         self
     }
@@ -284,7 +323,10 @@ impl Mock {
     /// ```
     ///
     pub fn with_header<T: IntoHeaderName>(mut self, field: T, value: &str) -> Self {
-        self.inner.response.headers.append(field, value.to_owned());
+        self.inner
+            .response
+            .headers
+            .append(field.into_header_name(), value.to_owned());
 
         self
     }
diff --git a/tests/lib.rs b/tests/lib.rs
index c371812..a0d2e19 100644
--- a/tests/lib.rs
+++ b/tests/lib.rs
@@ -1,7 +1,6 @@
 #[macro_use]
 extern crate serde_json;
 
-use hyper::header::HeaderName;
 use mockito::{Matcher, Server};
 use rand::distributions::Alphanumeric;
 use rand::Rng;
@@ -700,7 +699,7 @@ fn test_mock_preserves_header_order() {
 
     // Add a large number of headers so getting the same order accidentally is unlikely.
     for i in 0..100 {
-        let field: HeaderName = format!("x-custom-header-{}", i).try_into().unwrap();
+        let field = format!("x-custom-header-{}", i);
         let value = "test";
         mock = mock.with_header(&field, value);
         expected_headers.push(format!("{}: {}", field, value));