encryption/integrations/
ecat.rs1use std::task::{Context, Poll};
37
38use http::Request;
39use tower::{Layer, Service};
40
41use crate::guard::Guard;
42
43#[derive(Clone, Debug)]
45pub struct GuardLayer {
46 guard: Guard,
47}
48
49impl GuardLayer {
50 pub fn new(guard: Guard) -> Self {
52 Self { guard }
53 }
54}
55
56impl<S> Layer<S> for GuardLayer {
57 type Service = GuardService<S>;
58
59 fn layer(&self, inner: S) -> Self::Service {
60 GuardService {
61 inner,
62 guard: self.guard.clone(),
63 }
64 }
65}
66
67#[derive(Clone, Debug)]
69pub struct GuardService<S> {
70 inner: S,
71 guard: Guard,
72}
73
74impl<S, B> Service<Request<B>> for GuardService<S>
75where
76 S: Service<Request<B>>,
77{
78 type Response = S::Response;
79 type Error = S::Error;
80 type Future = S::Future;
82
83 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
84 self.inner.poll_ready(cx)
85 }
86
87 fn call(&mut self, mut req: Request<B>) -> Self::Future {
88 req.extensions_mut().insert(self.guard.clone());
90 self.inner.call(req)
91 }
92}
93
94#[cfg(test)]
95mod tests {
96 use super::*;
97 use std::convert::Infallible;
98 use tower::service_fn;
99
100 fn guard() -> Guard {
101 Guard::from_master_key(&[7u8; 32], "aes-256-gcm").unwrap()
102 }
103
104 async fn echo_identifier_count(req: Request<()>) -> Result<String, Infallible> {
107 let g = req
108 .extensions()
109 .get::<Guard>()
110 .expect("守卫应当已被注入")
111 .clone();
112 Ok(g.identifiers().len().to_string())
113 }
114
115 #[tokio::test]
116 async fn injects_the_guard_into_extensions() {
117 let mut svc = GuardLayer::new(guard()).layer(service_fn(echo_identifier_count));
118 let res = svc.call(Request::new(())).await.unwrap();
119 assert_eq!(res, "5");
120 }
121
122 #[tokio::test]
124 async fn injected_guard_shares_the_key() {
125 let state = guard();
126 let token = state.encrypt("共享".as_bytes()).unwrap();
127
128 let inner = service_fn(move |req: Request<()>| {
129 let token = token.clone();
130 async move {
131 let g = req.extensions().get::<Guard>().unwrap().clone();
132 Ok::<_, Infallible>(
133 g.decrypt(&token)
134 .map(|plain| String::from_utf8_lossy(&plain).into_owned())
135 .unwrap_or_else(|_| "解不开".into()),
136 )
137 }
138 });
139
140 let mut svc = GuardLayer::new(state).layer(inner);
141 let res = svc.call(Request::new(())).await.unwrap();
142 assert_eq!(res, "共享");
143 }
144
145 #[tokio::test]
147 async fn works_across_multiple_requests() {
148 let mut svc = GuardLayer::new(guard()).layer(service_fn(echo_identifier_count));
149 for _ in 0..3 {
150 assert_eq!(svc.call(Request::new(())).await.unwrap(), "5");
151 }
152 }
153
154 #[tokio::test]
156 async fn composes_with_ecat_service_builder() {
157 use tower::ServiceBuilder;
158
159 let mut svc = ServiceBuilder::new()
160 .layer(GuardLayer::new(guard()))
161 .service(service_fn(echo_identifier_count));
162
163 assert_eq!(svc.call(Request::new(())).await.unwrap(), "5");
164 }
165}