use std::task::{Context, Poll};
use http::Request;
use tower::{Layer, Service};
use crate::guard::Guard;
#[derive(Clone, Debug)]
pub struct GuardLayer {
guard: Guard,
}
impl GuardLayer {
pub fn new(guard: Guard) -> Self {
Self { guard }
}
}
impl<S> Layer<S> for GuardLayer {
type Service = GuardService<S>;
fn layer(&self, inner: S) -> Self::Service {
GuardService {
inner,
guard: self.guard.clone(),
}
}
}
#[derive(Clone, Debug)]
pub struct GuardService<S> {
inner: S,
guard: Guard,
}
impl<S, B> Service<Request<B>> for GuardService<S>
where
S: Service<Request<B>>,
{
type Response = S::Response;
type Error = S::Error;
type Future = S::Future;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: Request<B>) -> Self::Future {
req.extensions_mut().insert(self.guard.clone());
self.inner.call(req)
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::convert::Infallible;
use tower::service_fn;
fn guard() -> Guard {
Guard::from_master_key(&[7u8; 32], "aes-256-gcm").unwrap()
}
async fn echo_identifier_count(req: Request<()>) -> Result<String, Infallible> {
let g = req
.extensions()
.get::<Guard>()
.expect("守卫应当已被注入")
.clone();
Ok(g.identifiers().len().to_string())
}
#[tokio::test]
async fn injects_the_guard_into_extensions() {
let mut svc = GuardLayer::new(guard()).layer(service_fn(echo_identifier_count));
let res = svc.call(Request::new(())).await.unwrap();
assert_eq!(res, "5");
}
#[tokio::test]
async fn injected_guard_shares_the_key() {
let state = guard();
let token = state.encrypt("共享".as_bytes()).unwrap();
let inner = service_fn(move |req: Request<()>| {
let token = token.clone();
async move {
let g = req.extensions().get::<Guard>().unwrap().clone();
Ok::<_, Infallible>(
g.decrypt(&token)
.map(|plain| String::from_utf8_lossy(&plain).into_owned())
.unwrap_or_else(|_| "解不开".into()),
)
}
});
let mut svc = GuardLayer::new(state).layer(inner);
let res = svc.call(Request::new(())).await.unwrap();
assert_eq!(res, "共享");
}
#[tokio::test]
async fn works_across_multiple_requests() {
let mut svc = GuardLayer::new(guard()).layer(service_fn(echo_identifier_count));
for _ in 0..3 {
assert_eq!(svc.call(Request::new(())).await.unwrap(), "5");
}
}
#[tokio::test]
async fn composes_with_ecat_service_builder() {
use tower::ServiceBuilder;
let mut svc = ServiceBuilder::new()
.layer(GuardLayer::new(guard()))
.service(service_fn(echo_identifier_count));
assert_eq!(svc.call(Request::new(())).await.unwrap(), "5");
}
}