encryption/integrations/
axum.rs1use std::convert::Infallible;
37
38use axum::extract::FromRequestParts;
39use axum::http::request::Parts;
40
41use super::Guarded;
42use crate::guard::Guard;
43
44pub type GuardExtractor = Guard;
47
48impl<S> FromRequestParts<S> for Guard
49where
50 S: Guarded + Send + Sync,
51{
52 type Rejection = Infallible;
53
54 async fn from_request_parts(_parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
58 Ok(state.guard().clone())
59 }
60}
61
62#[cfg(test)]
63mod tests {
64 use super::*;
65 use axum::Router;
66 use axum::body::Body;
67 use axum::extract::Extension;
68 use axum::http::{Request, StatusCode};
69 use axum::routing::get;
70 use tower::ServiceExt as _;
71
72 fn guard() -> Guard {
73 Guard::from_master_key(&[7u8; 32], "aes-256-gcm").unwrap()
74 }
75
76 #[derive(Clone)]
78 struct BareState {
79 encryption: Guard,
80 }
81 impl Guarded for BareState {
82 fn guard(&self) -> &Guard {
83 &self.encryption
84 }
85 }
86
87 #[derive(Clone)]
89 struct AppState {
90 #[allow(dead_code)]
91 pool: std::sync::Arc<String>,
92 encryption: Guard,
93 }
94 impl Guarded for AppState {
95 fn guard(&self) -> &Guard {
96 &self.encryption
97 }
98 }
99
100 #[tokio::test]
102 async fn extracts_directly_from_parts() {
103 let state = AppState {
104 pool: std::sync::Arc::new("pool".into()),
105 encryption: guard(),
106 };
107 let (mut parts, _) = Request::new(()).into_parts();
108
109 let extracted = Guard::from_request_parts(&mut parts, &state).await.unwrap();
110 assert_eq!(extracted.identifiers().len(), 5);
111 }
112
113 #[tokio::test]
115 async fn handler_can_take_a_guard_argument() {
116 async fn handler(guard: Guard) -> String {
117 let blob = guard.encrypt("13800138000".as_bytes()).unwrap();
118 String::from_utf8(guard.decrypt(&blob).unwrap()).unwrap()
119 }
120
121 let state = AppState {
122 pool: std::sync::Arc::new("pool".into()),
123 encryption: guard(),
124 };
125 let app = Router::new().route("/", get(handler)).with_state(state);
126
127 let response = app
128 .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
129 .await
130 .unwrap();
131
132 assert_eq!(response.status(), StatusCode::OK);
133 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
134 .await
135 .unwrap();
136 assert_eq!(&body[..], b"13800138000");
137 }
138
139 #[tokio::test]
141 async fn handler_and_state_share_one_encrypter() {
142 async fn handler(guard: Guard, Extension(token): Extension<Vec<u8>>) -> String {
143 guard
144 .decrypt(&token)
145 .map(|plain| String::from_utf8_lossy(&plain).into_owned())
146 .unwrap_or_else(|_| "解不开".into())
147 }
148
149 let state = BareState {
150 encryption: guard(),
151 };
152 let token = state.encryption.encrypt("共享".as_bytes()).unwrap();
153
154 let app = Router::new()
155 .route("/", get(handler))
156 .layer(Extension(token))
157 .with_state(state);
158
159 let response = app
160 .oneshot(Request::builder().uri("/").body(Body::empty()).unwrap())
161 .await
162 .unwrap();
163
164 assert_eq!(response.status(), StatusCode::OK);
165 let body = axum::body::to_bytes(response.into_body(), usize::MAX)
166 .await
167 .unwrap();
168 assert_eq!(&body[..], "共享".as_bytes());
169 }
170}