use std::{future::Future, marker::PhantomData, pin::Pin, rc::Rc, sync::Arc};
use actix_web::{
dev::{Service, ServiceRequest, ServiceResponse, Transform},
Error,
};
use futures_util::{
future::{self, FutureExt as _, LocalBoxFuture, TryFutureExt as _},
ready,
task::{Context, Poll},
};
use crate::extractors::{basic, bearer, AuthExtractor};
#[derive(Debug, Clone)]
pub struct HttpAuthentication<T, F>
where
T: AuthExtractor,
{
process_fn: Arc<F>,
_extractor: PhantomData<T>,
}
impl<T, F, O> HttpAuthentication<T, F>
where
T: AuthExtractor,
F: Fn(ServiceRequest, T) -> O,
O: Future<Output = Result<ServiceRequest, Error>>,
{
pub fn with_fn(process_fn: F) -> HttpAuthentication<T, F> {
HttpAuthentication {
process_fn: Arc::new(process_fn),
_extractor: PhantomData,
}
}
}
impl<F, O> HttpAuthentication<basic::BasicAuth, F>
where
F: Fn(ServiceRequest, basic::BasicAuth) -> O,
O: Future<Output = Result<ServiceRequest, Error>>,
{
pub fn basic(process_fn: F) -> Self {
Self::with_fn(process_fn)
}
}
impl<F, O> HttpAuthentication<bearer::BearerAuth, F>
where
F: Fn(ServiceRequest, bearer::BearerAuth) -> O,
O: Future<Output = Result<ServiceRequest, Error>>,
{
pub fn bearer(process_fn: F) -> Self {
Self::with_fn(process_fn)
}
}
impl<S, B, T, F, O> Transform<S, ServiceRequest> for HttpAuthentication<T, F>
where
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
S::Future: 'static,
F: Fn(ServiceRequest, T) -> O + 'static,
O: Future<Output = Result<ServiceRequest, Error>> + 'static,
T: AuthExtractor + 'static,
{
type Response = ServiceResponse<B>;
type Error = Error;
type Transform = AuthenticationMiddleware<S, F, T>;
type InitError = ();
type Future = future::Ready<Result<Self::Transform, Self::InitError>>;
fn new_transform(&self, service: S) -> Self::Future {
future::ok(AuthenticationMiddleware {
service: Rc::new(service),
process_fn: self.process_fn.clone(),
_extractor: PhantomData,
})
}
}
#[doc(hidden)]
pub struct AuthenticationMiddleware<S, F, T>
where
T: AuthExtractor,
{
service: Rc<S>,
process_fn: Arc<F>,
_extractor: PhantomData<T>,
}
impl<S, B, F, T, O> Service<ServiceRequest> for AuthenticationMiddleware<S, F, T>
where
S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error> + 'static,
S::Future: 'static,
F: Fn(ServiceRequest, T) -> O + 'static,
O: Future<Output = Result<ServiceRequest, Error>> + 'static,
T: AuthExtractor + 'static,
{
type Response = ServiceResponse<B>;
type Error = S::Error;
type Future = LocalBoxFuture<'static, Result<ServiceResponse<B>, Error>>;
actix_service::forward_ready!(service);
fn call(&self, req: ServiceRequest) -> Self::Future {
let process_fn = Arc::clone(&self.process_fn);
let service = Rc::clone(&self.service);
async move {
let (req, credentials) = match Extract::<T>::new(req).await {
Ok(req) => req,
Err((err, req)) => {
return Ok(req.error_response(err));
}
};
let req = process_fn(req, credentials).await?;
service.call(req).await
}
.boxed_local()
}
}
struct Extract<T> {
req: Option<ServiceRequest>,
f: Option<LocalBoxFuture<'static, Result<T, Error>>>,
_extractor: PhantomData<fn() -> T>,
}
impl<T> Extract<T> {
pub fn new(req: ServiceRequest) -> Self {
Extract {
req: Some(req),
f: None,
_extractor: PhantomData,
}
}
}
impl<T> Future for Extract<T>
where
T: AuthExtractor,
T::Future: 'static,
T::Error: 'static,
{
type Output = Result<(ServiceRequest, T), (Error, ServiceRequest)>;
fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
if self.f.is_none() {
let req = self.req.as_ref().expect("Extract future was polled twice!");
let f = T::from_service_request(req).map_err(Into::into);
self.f = Some(f.boxed_local());
}
let f = self
.f
.as_mut()
.expect("Extraction future should be initialized at this point");
let credentials = ready!(f.as_mut().poll(ctx)).map_err(|err| {
(
err,
self.req.take().expect("Extract future was polled twice!"),
)
})?;
let req = self.req.take().expect("Extract future was polled twice!");
Poll::Ready(Ok((req, credentials)))
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::extractors::bearer::BearerAuth;
use actix_service::{into_service, Service};
use actix_web::error;
use actix_web::test::TestRequest;
#[actix_rt::test]
async fn test_middleware_panic() {
let middleware = AuthenticationMiddleware {
service: Rc::new(into_service(|_: ServiceRequest| async move {
actix_rt::time::sleep(std::time::Duration::from_secs(1)).await;
Err::<ServiceResponse, _>(error::ErrorBadRequest("error"))
})),
process_fn: Arc::new(|req, _: BearerAuth| async { Ok(req) }),
_extractor: PhantomData,
};
let req = TestRequest::get()
.append_header(("Authorization", "Bearer 1"))
.to_srv_request();
let f = middleware.call(req).await;
let _res = futures_util::future::lazy(|cx| middleware.poll_ready(cx)).await;
assert!(f.is_err());
}
#[actix_rt::test]
async fn test_middleware_panic_several_orders() {
let middleware = AuthenticationMiddleware {
service: Rc::new(into_service(|_: ServiceRequest| async move {
actix_rt::time::sleep(std::time::Duration::from_secs(1)).await;
Err::<ServiceResponse, _>(error::ErrorBadRequest("error"))
})),
process_fn: Arc::new(|req, _: BearerAuth| async { Ok(req) }),
_extractor: PhantomData,
};
let req = TestRequest::get()
.append_header(("Authorization", "Bearer 1"))
.to_srv_request();
let f1 = middleware.call(req).await;
let req = TestRequest::get()
.append_header(("Authorization", "Bearer 1"))
.to_srv_request();
let f2 = middleware.call(req).await;
let req = TestRequest::get()
.append_header(("Authorization", "Bearer 1"))
.to_srv_request();
let f3 = middleware.call(req).await;
let _res = futures_util::future::lazy(|cx| middleware.poll_ready(cx)).await;
assert!(f1.is_err());
assert!(f2.is_err());
assert!(f3.is_err());
}
}