unitycatalog_storage_proxy/
auth.rs1use std::task::{Context, Poll};
24
25use axum::extract::Request;
26use axum::response::{IntoResponse, Response};
27use futures::{FutureExt, future::BoxFuture};
28use tower::{Layer, Service};
29
30use crate::error::ProxyError;
31
32pub const DEFAULT_FORWARDED_USER_HEADER: &str = "x-forwarded-user";
34
35#[derive(Clone, Debug, PartialEq, Eq)]
42pub struct ForwardedIdentity(pub Option<String>);
43
44impl ForwardedIdentity {
45 pub fn anonymous() -> Self {
47 Self(None)
48 }
49
50 pub fn user(name: impl Into<String>) -> Self {
52 Self(Some(name.into()))
53 }
54}
55
56impl crate::backend::ForwardedUser for ForwardedIdentity {
57 fn forwarded_user(&self) -> Option<&str> {
58 self.0.as_deref()
59 }
60}
61
62#[derive(Clone, Debug, PartialEq, Eq)]
64pub enum AuthMode {
65 Anonymous,
67 ReverseProxy {
69 header: String,
72 allow_missing: bool,
75 },
76}
77
78impl AuthMode {
79 fn authenticate(&self, request: &Request) -> Result<ForwardedIdentity, ProxyError> {
81 match self {
82 AuthMode::Anonymous => Ok(ForwardedIdentity::anonymous()),
83 AuthMode::ReverseProxy {
84 header,
85 allow_missing,
86 } => match request.headers().get(header) {
87 Some(value) => {
88 let name = value.to_str().map_err(|_| {
89 ProxyError::Unauthenticated(format!("`{header}` header is not valid UTF-8"))
90 })?;
91 let name = name.trim();
92 if name.is_empty() {
93 return Err(ProxyError::Unauthenticated(format!(
94 "`{header}` header is empty"
95 )));
96 }
97 Ok(ForwardedIdentity::user(name))
98 }
99 None if *allow_missing => Ok(ForwardedIdentity::anonymous()),
100 None => Err(ProxyError::Unauthenticated(format!(
101 "missing `{header}` header"
102 ))),
103 },
104 }
105 }
106}
107
108#[derive(Clone)]
112pub struct AuthLayer {
113 mode: AuthMode,
114}
115
116impl AuthLayer {
117 pub fn new(mode: AuthMode) -> Self {
119 Self { mode }
120 }
121}
122
123impl<S> Layer<S> for AuthLayer {
124 type Service = AuthMiddleware<S>;
125
126 fn layer(&self, inner: S) -> Self::Service {
127 AuthMiddleware {
128 inner,
129 mode: self.mode.clone(),
130 }
131 }
132}
133
134#[derive(Clone)]
136pub struct AuthMiddleware<S> {
137 inner: S,
138 mode: AuthMode,
139}
140
141impl<S> Service<Request> for AuthMiddleware<S>
142where
143 S: Service<Request, Response = Response> + Send + 'static,
144 S::Future: Send + 'static,
145{
146 type Response = S::Response;
147 type Error = S::Error;
148 type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;
149
150 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
151 self.inner.poll_ready(cx)
152 }
153
154 fn call(&mut self, mut req: Request) -> Self::Future {
155 match self.mode.authenticate(&req) {
156 Ok(identity) => {
157 req.extensions_mut().insert(identity);
158 self.inner.call(req).boxed()
159 }
160 Err(e) => async move { Ok(e.into_response()) }.boxed(),
161 }
162 }
163}
164
165#[cfg(test)]
166mod tests {
167 use axum::body::Body;
168 use axum::http::StatusCode;
169 use tower::{ServiceBuilder, ServiceExt};
170
171 use super::*;
172
173 fn reverse_proxy(header: &str, allow_missing: bool) -> AuthMode {
174 AuthMode::ReverseProxy {
175 header: header.to_string(),
176 allow_missing,
177 }
178 }
179
180 #[test]
181 fn anonymous_always_anonymous() {
182 let req = Request::get("/").body(Body::empty()).unwrap();
183 assert_eq!(
184 AuthMode::Anonymous.authenticate(&req).unwrap(),
185 ForwardedIdentity::anonymous()
186 );
187 }
188
189 #[test]
190 fn reverse_proxy_extracts_forwarded_user() {
191 let mode = reverse_proxy(DEFAULT_FORWARDED_USER_HEADER, false);
192 let req = Request::get("/")
193 .header(DEFAULT_FORWARDED_USER_HEADER, "alice")
194 .body(Body::empty())
195 .unwrap();
196 assert_eq!(
197 mode.authenticate(&req).unwrap(),
198 ForwardedIdentity::user("alice")
199 );
200 }
201
202 #[test]
203 fn reverse_proxy_honors_custom_header_and_trims() {
204 let mode = reverse_proxy("x-user", false);
205 let req = Request::get("/")
206 .header("x-user", " bob ")
207 .body(Body::empty())
208 .unwrap();
209 assert_eq!(
210 mode.authenticate(&req).unwrap(),
211 ForwardedIdentity::user("bob")
212 );
213 }
214
215 #[test]
216 fn reverse_proxy_rejects_missing_header_by_default() {
217 let mode = reverse_proxy(DEFAULT_FORWARDED_USER_HEADER, false);
218 let req = Request::get("/").body(Body::empty()).unwrap();
219 assert!(mode.authenticate(&req).is_err());
220 }
221
222 #[test]
223 fn reverse_proxy_rejects_empty_header() {
224 let mode = reverse_proxy(DEFAULT_FORWARDED_USER_HEADER, false);
225 let req = Request::get("/")
226 .header(DEFAULT_FORWARDED_USER_HEADER, " ")
227 .body(Body::empty())
228 .unwrap();
229 assert!(mode.authenticate(&req).is_err());
230 }
231
232 #[test]
233 fn reverse_proxy_falls_back_to_anonymous_when_configured() {
234 let mode = reverse_proxy(DEFAULT_FORWARDED_USER_HEADER, true);
235 let req = Request::get("/").body(Body::empty()).unwrap();
236 assert_eq!(
237 mode.authenticate(&req).unwrap(),
238 ForwardedIdentity::anonymous()
239 );
240 }
241
242 #[tokio::test]
243 async fn layer_inserts_identity_extension() {
244 async fn echo(req: Request) -> Result<Response, std::convert::Infallible> {
245 assert!(req.extensions().get::<ForwardedIdentity>().is_some());
248 Ok(Response::new(Body::empty()))
249 }
250 let mut service = ServiceBuilder::new()
251 .layer(AuthLayer::new(AuthMode::Anonymous))
252 .service_fn(echo);
253 let req = Request::get("/").body(Body::empty()).unwrap();
254 let resp = service.ready().await.unwrap().call(req).await.unwrap();
255 assert_eq!(resp.status(), StatusCode::OK);
256 }
257
258 #[tokio::test]
259 async fn layer_missing_identity_maps_to_401() {
260 async fn ok(_req: Request) -> Result<Response, std::convert::Infallible> {
261 Ok(Response::new(Body::empty()))
262 }
263 let mut service = ServiceBuilder::new()
264 .layer(AuthLayer::new(reverse_proxy(
265 DEFAULT_FORWARDED_USER_HEADER,
266 false,
267 )))
268 .service_fn(ok);
269 let req = Request::get("/").body(Body::empty()).unwrap();
270 let resp = service.ready().await.unwrap().call(req).await.unwrap();
271 assert_eq!(resp.status(), StatusCode::UNAUTHORIZED);
272 }
273}