use std::{
convert::Infallible,
fmt,
future::Future,
net::{IpAddr, SocketAddr},
pin::Pin,
sync::Arc,
task::{Context, Poll},
time::Duration,
};
use ::axum::{
body::Body,
extract::{ConnectInfo, FromRef, FromRequestParts, OptionalFromRequestParts},
response::{IntoResponse, IntoResponseParts, Response, ResponseParts},
};
use http::{Extensions, HeaderMap, HeaderValue, Request, StatusCode, header, request::Parts};
use secrecy::{ExposeSecret, SecretString};
use tower_layer::Layer;
use tower_service::Service;
use uuid::Uuid;
use crate::{
AuthSession, Checks, ClientContext, Guard, GuardError, IssuedToken, NewSession, Outcome, Rejection, StepUpOutcome,
store::SessionStore, subject::Subject,
};
type Predicate = Arc<dyn Fn(&HeaderMap) -> bool + Send + Sync>;
impl<St: SessionStore, S: Subject> Guard<St, S> {
pub fn layer(&self) -> GuardLayer<St, S> {
GuardLayer {
shared: Arc::new(Shared {
guard: self.clone(),
checks: None,
recent_auth: None,
bearer_prefixes: Vec::new(),
predicates: Vec::new(),
}),
}
}
}
struct Shared<St, S> {
guard: Guard<St, S>,
checks: Option<Checks>,
recent_auth: Option<Duration>,
bearer_prefixes: Vec<String>,
predicates: Vec<Predicate>,
}
impl<St, S> Clone for Shared<St, S> {
fn clone(&self) -> Self {
Self {
guard: self.guard.clone(),
checks: self.checks,
recent_auth: self.recent_auth,
bearer_prefixes: self.bearer_prefixes.clone(),
predicates: self.predicates.clone(),
}
}
}
impl<St, S> Shared<St, S> {
fn bypassed(&self, headers: &HeaderMap) -> bool {
bearer_token(headers).is_some_and(|token| self.bearer_prefixes.iter().any(|p| token.starts_with(p.as_str())))
|| self.predicates.iter().any(|p| p(headers))
}
}
fn bearer_token(headers: &HeaderMap) -> Option<&str> {
let value = headers.get(header::AUTHORIZATION)?.to_str().ok()?.trim();
let (scheme, token) = value.split_once(' ')?;
scheme.eq_ignore_ascii_case("bearer").then(|| token.trim_start())
}
pub struct GuardLayer<St, S> {
shared: Arc<Shared<St, S>>,
}
impl<St, S> Clone for GuardLayer<St, S> {
fn clone(&self) -> Self {
Self {
shared: Arc::clone(&self.shared),
}
}
}
impl<St, S> fmt::Debug for GuardLayer<St, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GuardLayer")
.field("checks", &self.shared.checks)
.field("recent_auth", &self.shared.recent_auth)
.field("bearer_prefixes", &self.shared.bearer_prefixes)
.field("predicates", &self.shared.predicates.len())
.finish_non_exhaustive()
}
}
impl<St, S> GuardLayer<St, S> {
fn update(mut self, f: impl FnOnce(&mut Shared<St, S>)) -> Self {
f(Arc::make_mut(&mut self.shared));
self
}
#[must_use]
pub fn with(self, checks: Checks) -> Self {
self.update(|s| s.checks = Some(checks))
}
#[must_use]
pub fn require_recent_auth(self, max_age: Duration) -> Self {
self.update(|s| s.recent_auth = Some(max_age))
}
#[must_use]
pub fn bypass_bearer_prefixes<I, P>(self, prefixes: I) -> Self
where
I: IntoIterator<Item = P>,
P: Into<String>,
{
let prefixes: Vec<String> = prefixes.into_iter().map(Into::into).filter(|p| !p.is_empty()).collect();
self.update(|s| s.bearer_prefixes.extend(prefixes))
}
#[must_use]
pub fn bypass_when<F>(self, predicate: F) -> Self
where
F: Fn(&HeaderMap) -> bool + Send + Sync + 'static,
{
self.update(|s| s.predicates.push(Arc::new(predicate)))
}
}
impl<I, St, S> Layer<I> for GuardLayer<St, S> {
type Service = GuardService<I, St, S>;
fn layer(&self, inner: I) -> Self::Service {
GuardService {
inner,
shared: Arc::clone(&self.shared),
}
}
}
pub struct GuardService<I, St, S> {
inner: I,
shared: Arc<Shared<St, S>>,
}
impl<I: Clone, St, S> Clone for GuardService<I, St, S> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
shared: Arc::clone(&self.shared),
}
}
}
impl<I: fmt::Debug, St, S> fmt::Debug for GuardService<I, St, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("GuardService")
.field("inner", &self.inner)
.finish_non_exhaustive()
}
}
type BoxFuture<T> = Pin<Box<dyn Future<Output = T> + Send>>;
impl<I, B, St, S> Service<Request<B>> for GuardService<I, St, S>
where
I: Service<Request<B>, Response = Response> + Clone + Send + 'static,
I::Future: Send + 'static,
B: Send + 'static,
St: SessionStore,
S: Subject,
{
type Response = Response;
type Error = I::Error;
type Future = BoxFuture<Result<Response, I::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<B>) -> Self::Future {
let clone = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, clone);
let shared = Arc::clone(&self.shared);
Box::pin(async move {
if shared.bypassed(req.headers()) {
return inner.call(req).await;
}
let guard = &shared.guard;
let Some(peer) = peer_ip(req.extensions()) else {
return Ok(missing_connect_info());
};
let ctx = guard.client_context(peer, req.headers());
let Some(token) = guard.cookie_token(req.headers()) else {
return Ok(rejection_response(Rejection::NotAuthenticated, guard));
};
let outcome = match &shared.checks {
Some(checks) => guard.authenticate_with(token, &ctx, checks).await,
None => guard.authenticate(token, &ctx).await,
};
let auth = match outcome {
Ok(Outcome::Authenticated(auth)) => auth,
Ok(Outcome::Rejected(rejection)) => return Ok(rejection_response(rejection, guard)),
Err(error) => return Ok(guard_error(&error)),
};
let rotated = auth.rotated.as_ref().map(crate::IssuedToken::set_cookie);
if let Some(max_age) = shared.recent_auth
&& guard.requires_reauth(&auth.session, max_age)
{
let mut response = rejection_response(Rejection::StepUpRequired, guard);
append_set_cookie(response.headers_mut(), rotated);
return Ok(response);
}
req.extensions_mut().insert(auth.session.clone());
let mut response = inner.call(req).await?;
let headers = response.headers_mut();
append_set_cookie(headers, rotated);
headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
Ok(response)
})
}
}
fn peer_ip(extensions: &Extensions) -> Option<IpAddr> {
extensions
.get::<ConnectInfo<SocketAddr>>()
.map(|ConnectInfo(addr)| addr.ip())
}
fn append_set_cookie(headers: &mut HeaderMap, set_cookie: Option<&str>) {
let Some(set_cookie) = set_cookie else { return };
if let Ok(value) = HeaderValue::from_str(set_cookie) {
headers.append(header::SET_COOKIE, value);
} else {
tracing::error!("session cookie is not a valid header value; not sent");
}
}
fn json_response(status: StatusCode, error: &str) -> Response {
let mut response = Response::new(Body::from(format!(r#"{{"error":"{error}"}}"#)));
*response.status_mut() = status;
let headers = response.headers_mut();
headers.insert(header::CONTENT_TYPE, HeaderValue::from_static("application/json"));
headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
response
}
fn internal_error() -> Response {
json_response(StatusCode::INTERNAL_SERVER_ERROR, "internal")
}
fn missing_connect_info() -> Response {
log_missing_connect_info();
internal_error()
}
fn log_missing_connect_info() {
tracing::error!(
"session-guard: no ConnectInfo<SocketAddr> in request extensions; serve the app with \
`into_make_service_with_connect_info::<SocketAddr>()`. Failing the request closed."
);
}
fn guard_error(error: &GuardError) -> Response {
tracing::error!(%error, "session-guard: authentication failed");
internal_error()
}
pub fn rejection_response<St, S>(rejection: Rejection, guard: &Guard<St, S>) -> Response
where
St: SessionStore,
S: Subject,
{
rejection_with(rejection, &guard.clear_cookie())
}
fn rejection_with(rejection: Rejection, clear_cookie: &str) -> Response {
let mut response = json_response(StatusCode::UNAUTHORIZED, rejection.as_str());
let headers = response.headers_mut();
if rejection.clears_cookie() {
append_set_cookie(headers, Some(clear_cookie));
}
if rejection.clears_site_data() {
headers.insert("clear-site-data", HeaderValue::from_static("\"cookies\""));
}
response
}
pub fn client_context<St, S>(guard: &Guard<St, S>, parts: &Parts) -> Result<ClientContext, MissingConnectInfo>
where
St: SessionStore,
S: Subject,
{
let Some(peer) = peer_ip(&parts.extensions) else {
log_missing_connect_info();
return Err(MissingConnectInfo);
};
Ok(guard.client_context(peer, &parts.headers))
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MissingConnectInfo;
impl fmt::Display for MissingConnectInfo {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.write_str("no ConnectInfo<SocketAddr> in request extensions")
}
}
impl std::error::Error for MissingConnectInfo {}
impl IntoResponse for MissingConnectInfo {
fn into_response(self) -> Response {
internal_error()
}
}
impl<A, S> FromRequestParts<A> for AuthSession<S>
where
A: Send + Sync,
S: Clone + Send + Sync + 'static,
{
type Rejection = Response;
fn from_request_parts(parts: &mut Parts, _state: &A) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
std::future::ready(
parts
.extensions
.get::<AuthSession<S>>()
.cloned()
.ok_or_else(|| json_response(StatusCode::UNAUTHORIZED, Rejection::NotAuthenticated.as_str())),
)
}
}
impl<A, S> OptionalFromRequestParts<A> for AuthSession<S>
where
A: Send + Sync,
S: Clone + Send + Sync + 'static,
{
type Rejection = std::convert::Infallible;
fn from_request_parts(
parts: &mut Parts,
_state: &A,
) -> impl Future<Output = Result<Option<Self>, Self::Rejection>> + Send {
std::future::ready(Ok(parts.extensions.get::<AuthSession<S>>().cloned()))
}
}
fn cookie_parts(mut parts: ResponseParts, set_cookie: &str) -> ResponseParts {
let headers = parts.headers_mut();
append_set_cookie(headers, Some(set_cookie));
headers.insert(header::CACHE_CONTROL, HeaderValue::from_static("no-store"));
parts
}
impl IntoResponseParts for IssuedToken {
type Error = Infallible;
fn into_response_parts(self, parts: ResponseParts) -> Result<ResponseParts, Self::Error> {
Ok(cookie_parts(parts, self.set_cookie()))
}
}
impl IntoResponseParts for NewSession {
type Error = Infallible;
fn into_response_parts(self, parts: ResponseParts) -> Result<ResponseParts, Self::Error> {
self.token.into_response_parts(parts)
}
}
impl IntoResponse for GuardError {
fn into_response(self) -> Response {
guard_error(&self)
}
}
pub struct Sessions<St, S = Uuid> {
guard: Guard<St, S>,
client: ClientContext,
token: Option<SecretString>,
}
impl<St, S> fmt::Debug for Sessions<St, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Sessions")
.field("client", &self.client)
.field("token", &self.token.as_ref().map(|_| "<redacted>"))
.finish_non_exhaustive()
}
}
impl<St: SessionStore, S: Subject> Sessions<St, S> {
pub async fn login(&self, subject: &S) -> Result<NewSession, GuardError> {
self.guard.login(subject, &self.client, self.token()).await
}
pub async fn logout(&self) -> Result<LoggedOut, GuardError> {
let revoked = match self.token() {
Some(token) => self.guard.logout(token, &self.client).await?,
None => false,
};
Ok(LoggedOut {
revoked,
clear_cookie: self.guard.clear_cookie(),
})
}
pub async fn complete_step_up(&self, verified_subject: &S) -> Result<IssuedToken, StepUpError> {
let failure = match self.token() {
None => StepUpFailure::Rejected(Rejection::NotAuthenticated),
Some(token) => match self.guard.complete_step_up(token, &self.client, verified_subject).await {
Ok(StepUpOutcome::Completed(token)) => return Ok(token),
Ok(StepUpOutcome::SubjectMismatch) => StepUpFailure::SubjectMismatch,
Ok(StepUpOutcome::Rejected(rejection)) => StepUpFailure::Rejected(rejection),
Err(error) => StepUpFailure::Guard(error),
},
};
Err(StepUpError {
failure,
clear_cookie: self.guard.clear_cookie(),
})
}
pub fn guard(&self) -> &Guard<St, S> {
&self.guard
}
pub fn client(&self) -> &ClientContext {
&self.client
}
pub fn token(&self) -> Option<&str> {
self.token.as_ref().map(ExposeSecret::expose_secret)
}
}
impl<T, St, S> FromRequestParts<T> for Sessions<St, S>
where
T: Send + Sync,
Guard<St, S>: FromRef<T>,
St: SessionStore,
S: Subject,
{
type Rejection = MissingConnectInfo;
fn from_request_parts(parts: &mut Parts, state: &T) -> impl Future<Output = Result<Self, Self::Rejection>> + Send {
let guard = Guard::from_ref(state);
let sessions = client_context(&guard, parts).map(|client| {
let token = guard
.cookie_token(&parts.headers)
.map(|token| SecretString::from(token.to_owned()));
Self { guard, client, token }
});
std::future::ready(sessions)
}
}
#[derive(Clone, Debug)]
pub struct LoggedOut {
revoked: bool,
clear_cookie: String,
}
impl LoggedOut {
pub fn revoked(&self) -> bool {
self.revoked
}
}
impl IntoResponseParts for LoggedOut {
type Error = Infallible;
fn into_response_parts(self, parts: ResponseParts) -> Result<ResponseParts, Self::Error> {
Ok(cookie_parts(parts, &self.clear_cookie))
}
}
#[derive(Debug)]
pub struct StepUpError {
failure: StepUpFailure,
clear_cookie: String,
}
#[derive(Debug)]
#[non_exhaustive]
pub enum StepUpFailure {
SubjectMismatch,
Rejected(Rejection),
Guard(GuardError),
}
impl StepUpError {
pub fn failure(&self) -> &StepUpFailure {
&self.failure
}
pub fn into_failure(self) -> StepUpFailure {
self.failure
}
}
impl fmt::Display for StepUpError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match &self.failure {
StepUpFailure::SubjectMismatch => f.write_str("verified subject does not own the session"),
StepUpFailure::Rejected(rejection) => write!(f, "step-up rejected: {rejection}"),
StepUpFailure::Guard(_) => f.write_str("step-up failed"),
}
}
}
impl std::error::Error for StepUpError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match &self.failure {
StepUpFailure::Guard(error) => Some(error),
_ => None,
}
}
}
impl IntoResponse for StepUpError {
fn into_response(self) -> Response {
match self.failure {
StepUpFailure::SubjectMismatch => json_response(StatusCode::FORBIDDEN, "subject_mismatch"),
StepUpFailure::Rejected(rejection) => rejection_with(rejection, &self.clear_cookie),
StepUpFailure::Guard(error) => guard_error(&error),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn headers(auth: &str) -> HeaderMap {
let mut h = HeaderMap::new();
h.insert(header::AUTHORIZATION, HeaderValue::from_str(auth).unwrap());
h
}
#[test]
fn bearer_parsing() {
assert_eq!(bearer_token(&headers("Bearer abc")), Some("abc"));
assert_eq!(bearer_token(&headers("BEARER abc")), Some("abc"));
assert_eq!(bearer_token(&headers("Basic abc")), None);
assert_eq!(bearer_token(&headers("Bearer")), None);
assert_eq!(bearer_token(&HeaderMap::new()), None);
}
#[test]
fn json_body_is_exact() {
let response = json_response(StatusCode::UNAUTHORIZED, "session_expired");
assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
}
}