121 lines
3.2 KiB
Rust
121 lines
3.2 KiB
Rust
use std::sync::Arc;
|
|
|
|
use axum::{
|
|
http::{Request, StatusCode},
|
|
response::IntoResponse,
|
|
};
|
|
use tower::{Layer, Service};
|
|
|
|
use crate::{auth::AuthenticatedUser, error::AppError, models::ApiCapability};
|
|
|
|
#[derive(Clone, Copy)]
|
|
pub enum CapabilityStrategy {
|
|
All,
|
|
Any,
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct RequireCapabilitiesLayer {
|
|
required: Arc<Vec<ApiCapability>>,
|
|
strategy: CapabilityStrategy,
|
|
}
|
|
|
|
impl RequireCapabilitiesLayer {
|
|
pub fn all<I>(caps: I) -> Self
|
|
where
|
|
I: IntoIterator<Item = ApiCapability>,
|
|
{
|
|
Self {
|
|
required: Arc::new(caps.into_iter().collect()),
|
|
strategy: CapabilityStrategy::All,
|
|
}
|
|
}
|
|
|
|
pub fn any<I>(caps: I) -> Self
|
|
where
|
|
I: IntoIterator<Item = ApiCapability>,
|
|
{
|
|
Self {
|
|
required: Arc::new(caps.into_iter().collect()),
|
|
strategy: CapabilityStrategy::Any,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl<S> Layer<S> for RequireCapabilitiesLayer {
|
|
type Service = RequireCapabilities<S>;
|
|
|
|
fn layer(&self, inner: S) -> Self::Service {
|
|
RequireCapabilities {
|
|
inner,
|
|
required: Arc::clone(&self.required),
|
|
strategy: self.strategy,
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub struct RequireCapabilities<S> {
|
|
inner: S,
|
|
required: Arc<Vec<ApiCapability>>,
|
|
strategy: CapabilityStrategy,
|
|
}
|
|
|
|
impl<S, B> Service<Request<B>> for RequireCapabilities<S>
|
|
where
|
|
S: Service<Request<B>, Response = axum::response::Response> + Send,
|
|
S::Future: Send + 'static,
|
|
B: Send + 'static,
|
|
{
|
|
type Response = S::Response;
|
|
type Error = S::Error;
|
|
type Future = std::pin::Pin<
|
|
Box<dyn std::future::Future<Output = Result<Self::Response, Self::Error>> + Send>,
|
|
>;
|
|
|
|
fn poll_ready(
|
|
&mut self,
|
|
cx: &mut std::task::Context<'_>,
|
|
) -> std::task::Poll<Result<(), Self::Error>> {
|
|
self.inner.poll_ready(cx)
|
|
}
|
|
|
|
fn call(&mut self, req: Request<B>) -> Self::Future {
|
|
if self.required.is_empty() {
|
|
let fut = self.inner.call(req);
|
|
return Box::pin(async move { fut.await });
|
|
}
|
|
|
|
let (parts, body) = req.into_parts();
|
|
let user = match parts.extensions.get::<AuthenticatedUser>() {
|
|
Some(user) => user,
|
|
None => {
|
|
let response = AppError::unauthorized().into_response();
|
|
return Box::pin(async move { Ok(response) });
|
|
}
|
|
};
|
|
|
|
let allowed = match self.strategy {
|
|
CapabilityStrategy::All => self
|
|
.required
|
|
.iter()
|
|
.all(|cap| user.capabilities.contains(cap)),
|
|
CapabilityStrategy::Any => self
|
|
.required
|
|
.iter()
|
|
.any(|cap| user.capabilities.contains(cap)),
|
|
};
|
|
|
|
if !allowed {
|
|
let response = AppError::new(StatusCode::FORBIDDEN, "missing required capability")
|
|
.with_code("missing_capability")
|
|
.into_response();
|
|
return Box::pin(async move { Ok(response) });
|
|
}
|
|
|
|
let req = Request::from_parts(parts, body);
|
|
let fut = self.inner.call(req);
|
|
Box::pin(async move { fut.await })
|
|
}
|
|
}
|