use http::{header::HeaderName, Request, Response};
use pin_project_lite::pin_project;
use std::{
future::Future,
pin::Pin,
sync::Arc,
task::{ready, Context, Poll},
};
use tower_layer::Layer;
use tower_service::Service;
#[derive(Clone, Debug)]
pub struct SetSensitiveHeadersLayer {
headers: Arc<[HeaderName]>,
}
impl SetSensitiveHeadersLayer {
pub fn new<I>(headers: I) -> Self
where
I: IntoIterator<Item = HeaderName>,
{
let headers = headers.into_iter().collect::<Vec<_>>();
Self::from_shared(headers.into())
}
pub fn from_shared(headers: Arc<[HeaderName]>) -> Self {
Self { headers }
}
}
impl<S> Layer<S> for SetSensitiveHeadersLayer {
type Service = SetSensitiveHeaders<S>;
fn layer(&self, inner: S) -> Self::Service {
SetSensitiveRequestHeaders::from_shared(
SetSensitiveResponseHeaders::from_shared(inner, self.headers.clone()),
self.headers.clone(),
)
}
}
pub type SetSensitiveHeaders<S> = SetSensitiveRequestHeaders<SetSensitiveResponseHeaders<S>>;
#[derive(Clone, Debug)]
pub struct SetSensitiveRequestHeadersLayer {
headers: Arc<[HeaderName]>,
}
impl SetSensitiveRequestHeadersLayer {
pub fn new<I>(headers: I) -> Self
where
I: IntoIterator<Item = HeaderName>,
{
let headers = headers.into_iter().collect::<Vec<_>>();
Self::from_shared(headers.into())
}
pub fn from_shared(headers: Arc<[HeaderName]>) -> Self {
Self { headers }
}
}
impl<S> Layer<S> for SetSensitiveRequestHeadersLayer {
type Service = SetSensitiveRequestHeaders<S>;
fn layer(&self, inner: S) -> Self::Service {
SetSensitiveRequestHeaders {
inner,
headers: self.headers.clone(),
}
}
}
#[derive(Clone, Debug)]
pub struct SetSensitiveRequestHeaders<S> {
inner: S,
headers: Arc<[HeaderName]>,
}
impl<S> SetSensitiveRequestHeaders<S> {
pub fn new<I>(inner: S, headers: I) -> Self
where
I: IntoIterator<Item = HeaderName>,
{
let headers = headers.into_iter().collect::<Vec<_>>();
Self::from_shared(inner, headers.into())
}
pub fn from_shared(inner: S, headers: Arc<[HeaderName]>) -> Self {
Self { inner, headers }
}
define_inner_service_accessors!();
pub fn layer<I>(headers: I) -> SetSensitiveRequestHeadersLayer
where
I: IntoIterator<Item = HeaderName>,
{
SetSensitiveRequestHeadersLayer::new(headers)
}
}
impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for SetSensitiveRequestHeaders<S>
where
S: Service<Request<ReqBody>, Response = Response<ResBody>>,
{
type Response = S::Response;
type Error = S::Error;
type Future = S::Future;
#[inline]
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
let headers = req.headers_mut();
for header in &*self.headers {
if let http::header::Entry::Occupied(mut entry) = headers.entry(header) {
for value in entry.iter_mut() {
value.set_sensitive(true);
}
}
}
self.inner.call(req)
}
}
#[derive(Clone, Debug)]
pub struct SetSensitiveResponseHeadersLayer {
headers: Arc<[HeaderName]>,
}
impl SetSensitiveResponseHeadersLayer {
pub fn new<I>(headers: I) -> Self
where
I: IntoIterator<Item = HeaderName>,
{
let headers = headers.into_iter().collect::<Vec<_>>();
Self::from_shared(headers.into())
}
pub fn from_shared(headers: Arc<[HeaderName]>) -> Self {
Self { headers }
}
}
impl<S> Layer<S> for SetSensitiveResponseHeadersLayer {
type Service = SetSensitiveResponseHeaders<S>;
fn layer(&self, inner: S) -> Self::Service {
SetSensitiveResponseHeaders {
inner,
headers: self.headers.clone(),
}
}
}
#[derive(Clone, Debug)]
pub struct SetSensitiveResponseHeaders<S> {
inner: S,
headers: Arc<[HeaderName]>,
}
impl<S> SetSensitiveResponseHeaders<S> {
pub fn new<I>(inner: S, headers: I) -> Self
where
I: IntoIterator<Item = HeaderName>,
{
let headers = headers.into_iter().collect::<Vec<_>>();
Self::from_shared(inner, headers.into())
}
pub fn from_shared(inner: S, headers: Arc<[HeaderName]>) -> Self {
Self { inner, headers }
}
define_inner_service_accessors!();
pub fn layer<I>(headers: I) -> SetSensitiveResponseHeadersLayer
where
I: IntoIterator<Item = HeaderName>,
{
SetSensitiveResponseHeadersLayer::new(headers)
}
}
impl<ReqBody, ResBody, S> Service<Request<ReqBody>> for SetSensitiveResponseHeaders<S>
where
S: Service<Request<ReqBody>, Response = Response<ResBody>>,
{
type Response = S::Response;
type Error = S::Error;
type Future = SetSensitiveResponseHeadersResponseFuture<S::Future>;
#[inline]
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
SetSensitiveResponseHeadersResponseFuture {
future: self.inner.call(req),
headers: self.headers.clone(),
}
}
}
pin_project! {
#[derive(Debug)]
pub struct SetSensitiveResponseHeadersResponseFuture<F> {
#[pin]
future: F,
headers: Arc<[HeaderName]>,
}
}
impl<F, ResBody, E> Future for SetSensitiveResponseHeadersResponseFuture<F>
where
F: Future<Output = Result<Response<ResBody>, E>>,
{
type Output = F::Output;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
let mut res = ready!(this.future.poll(cx)?);
let headers = res.headers_mut();
for header in &**this.headers {
if let http::header::Entry::Occupied(mut entry) = headers.entry(header) {
for value in entry.iter_mut() {
value.set_sensitive(true);
}
}
}
Poll::Ready(Ok(res))
}
}
#[cfg(test)]
mod tests {
#[allow(unused_imports)]
use super::*;
use http::header;
use tower::{ServiceBuilder, ServiceExt};
#[tokio::test]
async fn multiple_value_header() {
async fn response_set_cookie(req: http::Request<()>) -> Result<http::Response<()>, ()> {
let mut iter = req.headers().get_all(header::COOKIE).iter().peekable();
assert!(iter.peek().is_some());
for value in iter {
assert!(value.is_sensitive())
}
let mut resp = http::Response::new(());
resp.headers_mut().append(
header::CONTENT_TYPE,
http::HeaderValue::from_static("text/html"),
);
resp.headers_mut().append(
header::SET_COOKIE,
http::HeaderValue::from_static("cookie-1"),
);
resp.headers_mut().append(
header::SET_COOKIE,
http::HeaderValue::from_static("cookie-2"),
);
resp.headers_mut().append(
header::SET_COOKIE,
http::HeaderValue::from_static("cookie-3"),
);
Ok(resp)
}
let mut service = ServiceBuilder::new()
.layer(SetSensitiveRequestHeadersLayer::new(vec![header::COOKIE]))
.layer(SetSensitiveResponseHeadersLayer::new(vec![
header::SET_COOKIE,
]))
.service_fn(response_set_cookie);
let mut req = http::Request::new(());
req.headers_mut()
.append(header::COOKIE, http::HeaderValue::from_static("cookie+1"));
req.headers_mut()
.append(header::COOKIE, http::HeaderValue::from_static("cookie+2"));
let resp = service.ready().await.unwrap().call(req).await.unwrap();
assert!(!resp
.headers()
.get(header::CONTENT_TYPE)
.unwrap()
.is_sensitive());
let mut iter = resp.headers().get_all(header::SET_COOKIE).iter().peekable();
assert!(iter.peek().is_some());
for value in iter {
assert!(value.is_sensitive())
}
}
}