use std::{
future::Future,
pin::Pin,
sync::Arc,
task::{Context, Poll},
};
use axum::{
body::Body,
http::{header, Request, Response, StatusCode},
};
use base64::{engine::general_purpose::STANDARD, Engine};
use tower::{Layer, Service};
use rok_auth::Claims;
pub type VerifyFn = Arc<
dyn Fn(String, String) -> Pin<Box<dyn Future<Output = Option<Claims>> + Send>> + Send + Sync,
>;
#[derive(Clone)]
pub struct BasicAuthLayer {
verifier: VerifyFn,
realm: String,
}
impl BasicAuthLayer {
pub fn new<F, Fut>(realm: impl Into<String>, verifier: F) -> Self
where
F: Fn(String, String) -> Fut + Send + Sync + 'static,
Fut: Future<Output = Option<Claims>> + Send + 'static,
{
Self {
realm: realm.into(),
verifier: Arc::new(move |u, p| Box::pin(verifier(u, p))),
}
}
}
impl<Svc> Layer<Svc> for BasicAuthLayer {
type Service = BasicAuthService<Svc>;
fn layer(&self, inner: Svc) -> Self::Service {
BasicAuthService {
inner,
verifier: Arc::clone(&self.verifier),
realm: self.realm.clone(),
}
}
}
#[derive(Clone)]
pub struct BasicAuthService<Svc> {
inner: Svc,
verifier: VerifyFn,
realm: String,
}
type BoxFut<T, E> = Pin<Box<dyn Future<Output = Result<T, E>> + Send>>;
impl<Svc, ReqBody> Service<Request<ReqBody>> for BasicAuthService<Svc>
where
Svc: Service<Request<ReqBody>, Response = Response<Body>> + Clone + Send + 'static,
Svc::Error: Send,
Svc::Future: Send,
ReqBody: Send + 'static,
{
type Response = Response<Body>;
type Error = Svc::Error;
type Future = BoxFut<Self::Response, Self::Error>;
fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}
fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
let verifier = Arc::clone(&self.verifier);
let realm = self.realm.clone();
let clone = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, clone);
Box::pin(async move {
let Some((username, password)) = parse_basic_header(req.headers()) else {
return Ok(deny(&realm));
};
let Some(claims) = verifier(username, password).await else {
return Ok(deny(&realm));
};
req.extensions_mut().insert(claims);
inner.call(req).await
})
}
}
pub fn parse_basic_header(headers: &axum::http::HeaderMap) -> Option<(String, String)> {
let value = headers.get(header::AUTHORIZATION)?.to_str().ok()?;
let encoded = value.strip_prefix("Basic ")?;
let decoded = STANDARD.decode(encoded.trim()).ok()?;
let s = String::from_utf8(decoded).ok()?;
let mut parts = s.splitn(2, ':');
let username = parts.next()?.to_string();
let password = parts.next()?.to_string();
Some((username, password))
}
fn deny(realm: &str) -> Response<Body> {
Response::builder()
.status(StatusCode::UNAUTHORIZED)
.header(header::WWW_AUTHENTICATE, format!("Basic realm=\"{realm}\""))
.header(header::CONTENT_TYPE, "application/json")
.body(Body::from(
r#"{"error":"Unauthorized","message":"Basic authentication required"}"#,
))
.unwrap()
}