set_header 一族

pub trait MakeHeaderValue<T> {
    fn make_header_value(&mut self, message: &T) -> Option<HeaderValue>;
}

impl<F, T> MakeHeaderValue<T> for F
where
    F: FnMut(&T) -> Option<HeaderValue>,
// ...
impl<T> MakeHeaderValue<T> for HeaderValue {
    fn make_header_value(&mut self, _message: &T) -> Option<HeaderValue> {
        Some(self.clone())
    }
}

            InsertHeaderMode::Override => {
                if let Some(value) = make.make_header_value(target) {
                    target.headers_mut().insert(header_name.clone(), value);
                }
            }
            InsertHeaderMode::IfNotPresent => {
                if !target.headers().contains_key(header_name) {
                    if let Some(value) = make.make_header_value(target) {
                        target.headers_mut().insert(header_name.clone(), value);
                    }
                }
            }
            InsertHeaderMode::Append => {
                if let Some(value) = make.make_header_value(target) {
                    target.headers_mut().append(header_name.clone(), value);
                }
            }

    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
        self.mode.apply(&self.header_name, &mut req, &mut self.make);
        self.inner.call(req)
    }

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        ResponseFuture {
            future: self.inner.call(req),
            header_name: self.header_name.clone(),
            make: self.make.clone(),
            mode: self.mode,
        }
    }

tower-http = { version = "=0.7.1", features = ["set-header"] }
tower = { version = "=0.5.3", features = ["util"] }
tokio = { version = "=1.54.1", features = ["macros", "rt"] }
http = "1"
http-body-util = "0.1"
bytes = "1"

use bytes::Bytes;
use http::{header, HeaderName, HeaderValue, Request, Response};
use http_body_util::Full;
use std::convert::Infallible;
use tower::{Service, ServiceBuilder, ServiceExt};
use tower_http::set_header::{
  response::SetMultipleResponseHeadersLayer, SetRequestHeaderLayer, SetResponseHeaderLayer,
};

const SEQ: HeaderName = HeaderName::from_static("x-seq");

// 内层：把收到的 x-seq 抄进响应，再自己带一个 cache-control
async fn handler(req: Request<Full<Bytes>>) -> Result<Response<Full<Bytes>>, Infallible> {
  let seen = req.headers().get(&SEQ).cloned().unwrap_or(HeaderValue::from_static("-"));
  Ok(
    Response::builder()
      .header("x-seen-seq", seen)
      .header(header::CACHE_CONTROL, "no-store")
      .body(Full::default())
      .unwrap(),
  )
}

// 一个带状态的闭包：每调一次加一
fn counter() -> impl FnMut(&Request<Full<Bytes>>) -> Option<HeaderValue> + Clone {
  let mut n = 0;
  move |_| {
    n += 1;
    Some(HeaderValue::from(n))
  }
}

fn counter_res() -> impl FnMut(&Response<Full<Bytes>>) -> Option<HeaderValue> + Clone {
  let mut n = 0;
  move |_| {
    n += 1;
    Some(HeaderValue::from(n))
  }
}

#[tokio::main(flavor = "current_thread")]
async fn main() {
  println!("stateful closure, same service, 3 calls");
  let mut svc = ServiceBuilder::new()
    .layer(SetRequestHeaderLayer::overriding(SEQ, counter()))
    .layer(SetResponseHeaderLayer::overriding(HeaderName::from_static("x-res-seq"), counter_res()))
    .service_fn(handler);
  for _ in 0..3 {
    let res = svc.ready().await.unwrap().call(Request::new(Full::default())).await.unwrap();
    println!("  request side x-seq={:?}  response side x-res-seq={:?}", res.headers()["x-seen-seq"], res.headers()["x-res-seq"]);
  }
  println!("same layers, svc.clone().oneshot() 3 times");
  for _ in 0..3 {
    let res = svc.clone().oneshot(Request::new(Full::default())).await.unwrap();
    println!("  request side x-seq={:?}  response side x-res-seq={:?}", res.headers()["x-seen-seq"], res.headers()["x-res-seq"]);
  }

  println!("three modes on a header the handler already set");
  for (name, layer) in [
    ("overriding", SetResponseHeaderLayer::overriding(header::CACHE_CONTROL, HeaderValue::from_static("max-age=60"))),
    ("appending", SetResponseHeaderLayer::appending(header::CACHE_CONTROL, HeaderValue::from_static("max-age=60"))),
    ("if_not_present", SetResponseHeaderLayer::if_not_present(header::CACHE_CONTROL, HeaderValue::from_static("max-age=60"))),
  ] {
    let res = ServiceBuilder::new().layer(layer).service_fn(handler).oneshot(Request::new(Full::default())).await.unwrap();
    let all: Vec<_> = res.headers().get_all(header::CACHE_CONTROL).iter().collect();
    println!("  {name:<15} {all:?}");
  }

  println!("closure returning None, and multiple headers");
  let layer = SetMultipleResponseHeadersLayer::overriding(vec![
    (header::X_FRAME_OPTIONS, HeaderValue::from_static("DENY")).into(),
    (header::CACHE_CONTROL, |res: &Response<Full<Bytes>>| {
      // 只在 handler 没说 no-store 时改
      (res.headers()[header::CACHE_CONTROL] != "no-store").then(|| HeaderValue::from_static("max-age=60"))
    })
      .into(),
  ]);
  let res = ServiceBuilder::new().layer(layer.clone()).service_fn(handler).oneshot(Request::new(Full::default())).await.unwrap();
  println!("  x-frame-options={:?} cache-control={:?}", res.headers()[header::X_FRAME_OPTIONS], res.headers()[header::CACHE_CONTROL]);
  println!("  layer debug: {:?}", layer);
}

stateful closure, same service, 3 calls
  request side x-seq="1"  response side x-res-seq="1"
  request side x-seq="2"  response side x-res-seq="1"
  request side x-seq="3"  response side x-res-seq="1"
same layers, svc.clone().oneshot() 3 times
  request side x-seq="4"  response side x-res-seq="1"
  request side x-seq="4"  response side x-res-seq="1"
  request side x-seq="4"  response side x-res-seq="1"
three modes on a header the handler already set
  overriding      ["max-age=60"]
  appending       ["no-store", "max-age=60"]
  if_not_present  ["no-store"]
closure returning None, and multiple headers
  x-frame-options="DENY" cache-control="no-store"
  layer debug: SetMultipleResponseHeadersLayer { headers: [HeaderInsertionConfig { header_name: "x-frame-options", mode: Override, make: "BoxedMakeHeaderValue" }, HeaderInsertionConfig { header_name: "cache-control", mode: Override, make: "BoxedMakeHeaderValue" }] }

    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
        ResponseFuture {
            future: self.inner.call(req),
            headers: self.headers.clone(),
        }
    }