use std::fmt;
use http::HeaderMap;
use http::HeaderName;
use http::HeaderValue;
use http::header::USER_AGENT;
use serde::Deserialize;
use serde::Serialize;
const SENSITIVE_HEADERS: &[&str] = &[
"authorization",
"proxy-authorization",
"x-api-key",
"api-key",
"x-goog-api-key",
"cookie",
"set-cookie",
];
const MASKED: &str = "***";
#[derive(Clone, Default, PartialEq, Eq)]
pub struct Headers(HeaderMap);
impl Headers {
#[must_use]
pub fn new() -> Self {
Self(HeaderMap::new())
}
#[must_use]
pub fn from_map(map: HeaderMap) -> Self {
Self(map)
}
#[must_use]
pub fn as_map(&self) -> &HeaderMap {
&self.0
}
pub fn as_map_mut(&mut self) -> &mut HeaderMap {
&mut self.0
}
#[must_use]
pub fn into_map(self) -> HeaderMap {
self.0
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
#[must_use]
pub fn len(&self) -> usize {
self.0.len()
}
#[must_use]
pub fn get_str(&self, name: &str) -> Option<&str> {
self.0.get(name).and_then(|value| value.to_str().ok())
}
#[must_use]
pub fn contains(&self, name: &str) -> bool {
self.0.contains_key(name)
}
pub fn insert(&mut self, name: &str, value: &str) -> Result<(), InvalidHeader> {
let name = HeaderName::from_bytes(name.as_bytes()).map_err(|_| InvalidHeader {
name: name.to_owned(),
})?;
let value = HeaderValue::from_str(value).map_err(|_| InvalidHeader {
name: name.to_string(),
})?;
self.0.insert(name, value);
Ok(())
}
#[must_use]
pub fn with(mut self, name: &str, value: &str) -> Self {
let _ = self.insert(name, value);
self
}
pub fn remove(&mut self, name: &str) {
self.0.remove(name);
}
pub fn merge(&mut self, other: &Headers) {
let mut current: Option<HeaderName> = None;
for (name, value) in &other.0 {
if current.as_ref() != Some(name) {
self.0.remove(name);
current = Some(name.clone());
}
self.0.append(name.clone(), value.clone());
}
}
#[must_use]
pub fn merged(mut self, other: &Headers) -> Self {
self.merge(other);
self
}
pub fn merge_pairs<'a, I>(&mut self, pairs: I)
where
I: IntoIterator<Item = (&'a str, Option<&'a str>)>,
{
for (name, value) in pairs {
if let Some(value) = value {
let _ = self.insert(name, value);
}
}
}
pub fn with_user_agent_suffix<'a, I>(mut self, suffixes: I) -> Self
where
I: IntoIterator<Item = &'a str>,
{
let mut parts: Vec<String> = Vec::new();
if let Some(existing) = self.get_str(USER_AGENT.as_str()) {
let existing = existing.trim();
if !existing.is_empty() {
parts.push(existing.to_owned());
}
}
parts.extend(
suffixes
.into_iter()
.map(str::trim)
.filter(|s| !s.is_empty())
.map(str::to_owned),
);
let joined = parts.join(" ");
if let Ok(value) = HeaderValue::from_str(&joined) {
self.0.insert(USER_AGENT, value);
}
self
}
pub fn iter_str(&self) -> impl Iterator<Item = (&str, String)> + '_ {
self.0.iter().map(|(name, value)| {
(
name.as_str(),
String::from_utf8_lossy(value.as_bytes()).into_owned(),
)
})
}
#[must_use]
pub fn masked(&self) -> Headers {
let mut masked = HeaderMap::with_capacity(self.0.len());
for (name, value) in &self.0 {
if is_sensitive(name) {
masked.append(name.clone(), HeaderValue::from_static(MASKED));
} else {
masked.append(name.clone(), value.clone());
}
}
Headers(masked)
}
}
fn is_sensitive(name: &HeaderName) -> bool {
SENSITIVE_HEADERS.contains(&name.as_str())
}
impl fmt::Debug for Headers {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut map = f.debug_map();
for (name, value) in &self.0 {
if is_sensitive(name) {
map.entry(&name.as_str(), &MASKED);
} else {
map.entry(&name.as_str(), &String::from_utf8_lossy(value.as_bytes()));
}
}
map.finish()
}
}
impl From<HeaderMap> for Headers {
fn from(map: HeaderMap) -> Self {
Self(map)
}
}
impl From<Headers> for HeaderMap {
fn from(headers: Headers) -> Self {
headers.0
}
}
impl<'a> IntoIterator for &'a Headers {
type Item = (&'a HeaderName, &'a HeaderValue);
type IntoIter = http::header::Iter<'a, HeaderValue>;
fn into_iter(self) -> Self::IntoIter {
self.0.iter()
}
}
impl Serialize for Headers {
fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
use serde::ser::SerializeMap;
let mut entries: Vec<(&str, String)> = Vec::with_capacity(self.0.len());
for name in self.0.keys() {
let rendered = if is_sensitive(name) {
MASKED.to_owned()
} else {
self.0
.get_all(name)
.iter()
.map(|v| String::from_utf8_lossy(v.as_bytes()).into_owned())
.collect::<Vec<_>>()
.join(", ")
};
entries.push((name.as_str(), rendered));
}
let mut map = serializer.serialize_map(Some(entries.len()))?;
for (name, value) in entries {
map.serialize_entry(name, &value)?;
}
map.end()
}
}
impl<'de> Deserialize<'de> for Headers {
fn deserialize<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
struct HeadersVisitor;
impl<'de> serde::de::Visitor<'de> for HeadersVisitor {
type Value = Headers;
fn expecting(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("a map of header names to string values")
}
fn visit_map<A: serde::de::MapAccess<'de>>(
self,
mut map: A,
) -> Result<Headers, A::Error> {
let mut headers = Headers::new();
while let Some((name, value)) = map.next_entry::<String, String>()? {
headers
.insert(&name, &value)
.map_err(serde::de::Error::custom)?;
}
Ok(headers)
}
}
deserializer.deserialize_map(HeadersVisitor)
}
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
#[error("invalid header `{name}`")]
pub struct InvalidHeader {
pub name: String,
}