axum-jwt 0.1.2

//! Middleware types and traits.
//!
//! If you just need to access token data in a handler, use
//! the [`Token`] and [`Claims`](crate::Claims) extractors directly.
//!
//! The [`layer`] function creates a configurable axum [middleware] layer.
//! When a request is made to a handler wrapped in this layer, the
//! JSON Web Token is validated. If validation succeeds, the handler is called.
//! If validation fails, a `401 Unauthorized` status code is returned, though
//! more fine-grained [configuration] is possible.
//!
//! [middleware]: https://docs.rs/axum/latest/axum/middleware/index.html
//! [configuration]: #configuration
//!
//! # Examples
//!
//! ```
//! use {
//!     axum::{Router, routing},
//!     axum_jwt::{Decoder, jsonwebtoken::DecodingKey},
//! };
//!
//! // This handler will be called only if the token is successfully validated.
//! async fn hello() -> String {
//!     "Hello, Anonimus!".to_owned()
//! }
//!
//! # async fn f() -> std::io::Result<()> {
//! let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
//!
//! let app = Router::new()
//!     .route("/", routing::get(hello))
//!     .layer(axum_jwt::layer(decoder));
//!
//! # use tokio::net::TcpListener;
//! let listener = TcpListener::bind("0.0.0.0:3000").await?;
//! axum::serve(listener, app).await?;
//! # Ok(())
//! # }
//! ```
//!
//! # Configuration
//!
//! The [`layer`] function accepts a [decoder](Decoder) that defines how to
//! decode and validate the token.
//!
//! Additionally, the layer itself can be
//! configured: set a [filter](JwtLayer::with_filter) to define the token's
//! data type and perform extra checks, store the token in
//! [extensions](JwtLayer::store_to_extension) so it can later be retrieved in
//! the handler via an extractor, or specify a custom
//! method of [extracting](JwtLayer::with_extract) the token from the request.

use {
    crate::{
        decode::Decoder,
        error::Error,
        extract::{Bearer, Extract, Token},
    },
    axum_core::{
        extract::Request,
        response::{IntoResponse, Response},
    },
    http::{Extensions, StatusCode},
    jsonwebtoken::TokenData,
    serde::de::{DeserializeOwned, IgnoredAny},
    std::{
        any,
        convert::Infallible,
        fmt,
        marker::PhantomData,
        mem,
        pin::Pin,
        task::{self, Context, Poll},
    },
    tower_layer::Layer,
    tower_service::Service,
};

/// Creates a [layer](JwtLayer) for middleware.
///
/// # Examples
///
/// ```
/// use {
///     axum::{Router, routing},
///     axum_jwt::{Decoder, jsonwebtoken::DecodingKey},
/// };
///
/// // This handler will be called only if the token is successfully validated.
/// async fn hello() -> String {
///     "Hello, Anonimus!".to_owned()
/// }
///
/// let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
///
/// let app = Router::new()
///     .route("/", routing::get(hello))
///     .layer(axum_jwt::layer(decoder));
/// # let _: Router = app;
/// ```
pub fn layer(decoder: Decoder) -> JwtLayer {
    JwtLayer {
        decoder,
        validate: Discard,
        store: |_, _| {},
        extract: PhantomData,
    }
}

/// Layer type for creating middleware.
///
/// To configure the layer and create the middleware service, call
/// the [`layer`] function.
pub struct JwtLayer<I = IgnoredAny, H = Discard, X = Bearer> {
    decoder: Decoder,
    validate: H,
    store: fn(Token<I>, &mut Extensions),
    extract: PhantomData<fn() -> X>,
}

impl<I, X> JwtLayer<I, Discard, X> {
    /// Sets a filter for additional validation.
    ///
    /// By default, the layer only validates the token header, ignoring all
    /// its claims. This method allows you to specify an arbitrary data type
    /// for the claims and perform additional token checks.
    ///
    /// The claims type must implement [`Deserialize`](serde::Deserialize).
    ///
    /// # Examples
    ///
    /// ```
    /// use {
    ///     axum::{Router, routing},
    ///     axum_jwt::{Decoder, Token, jsonwebtoken::DecodingKey},
    ///     serde::Deserialize,
    /// };
    ///
    /// #[derive(Deserialize)]
    /// struct User {
    ///     roles: Vec<String>,
    /// }
    ///
    /// // Checks that the user's token contains the admin role.
    /// fn check_access(t: &Token<User>) -> bool {
    ///     t.claims.roles.iter().any(|role| role == "admin")
    /// }
    ///
    /// // Called only if the role check is successful.
    /// async fn hello() -> String {
    ///     "Hello, Admin!".to_owned()
    /// }
    ///
    /// let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
    ///
    /// let app = Router::new()
    ///     .route("/", routing::get(hello))
    ///     .layer(axum_jwt::layer(decoder).with_filter(check_access));
    /// # let _: Router = app;
    /// ```
    ///
    /// If you also need to use the token inside the handler,
    /// see the [`store_to_extension`](JwtLayer::store_to_extension) method.
    ///
    /// # Callback return value
    ///
    /// The return value of the provided callback can be:
    ///
    /// * `bool`: where `true` means validation succeeded, and `false` means
    ///   it failed returning an HTTP status code `401 Unauthorized`.
    /// * `Result<(), E>`: where `Ok(())` means validation succeeded
    ///   and `Err(e)` means it failed. The error type must implement
    ///   [`IntoResponse`], which will be called on failure to return the
    ///   corresponding response.
    pub fn with_filter<H, N, O>(self, validate: H) -> JwtLayer<N, H, X>
    where
        H: FnMut(&Token<N>) -> O,
        N: DeserializeOwned,
        O: Output,
    {
        JwtLayer {
            decoder: self.decoder,
            validate,
            store: |_, _| {},
            extract: PhantomData,
        }
    }
}

impl<I, H, X> JwtLayer<I, H, X> {
    /// Configures the layer to store the token in the [extension].
    ///
    /// [extension]: https://docs.rs/axum/latest/axum/struct.Extension.html
    ///
    /// If you just need to access token data in a handler, use
    /// the [`Token`] and [`Claims`](crate::Claims) extractors directly.
    /// This function is useful only if you want to use the middleware
    /// but still access token data in some handlers.
    ///
    /// After calling this method, the middleware will store the parsed token
    /// in the extension, which can later be retrieved, for example,
    /// in a handler.
    ///
    /// The token is stored only after *successful validation*, including
    /// the configured [filter](JwtLayer::with_filter). This is usually what
    /// you want, but if you reuse the handler elsewhere, keep in mind that
    /// extracting it from [`Extension`] may fail.
    ///
    /// [`Extension`]: https://docs.rs/axum/latest/axum/struct.Extension.html
    ///
    /// # Examples
    ///
    /// ```
    /// use {
    ///     axum::{Extension, Router, routing},
    ///     axum_jwt::{Decoder, Token, jsonwebtoken::DecodingKey},
    ///     serde::Deserialize,
    /// };
    ///
    /// // To store a value in the extension, it must implement `Clone`.
    /// #[derive(Clone, Deserialize)]
    /// struct User {
    ///     sub: String,
    ///     roles: Vec<String>,
    /// }
    ///
    /// // Checks that the user's token contains the admin role.
    /// fn check_access(t: &Token<User>) -> bool {
    ///     t.claims.roles.iter().any(|role| role == "admin")
    /// }
    ///
    /// // Called only if the role check is successful.
    /// async fn hello(Extension(t): Extension<Token<User>>) -> String {
    ///     // We can also access the parsed token
    ///     format!("Hello, {}!", t.claims.sub)
    /// }
    ///
    /// let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
    ///
    /// let app = Router::new()
    ///     .route("/", routing::get(hello))
    ///     .layer(
    ///         axum_jwt::layer(decoder)
    ///             .with_filter(check_access)
    ///             .store_to_extension(),
    ///     );
    /// # let _: Router = app;
    /// ```
    ///
    /// <section class="warning">
    ///
    /// Note that the `store_to_extension` call comes after setting
    /// the `with_filter` filter.
    ///
    /// This is important because
    /// `store_to_extension` applies to the current token type.
    /// The `with_filter` call may change the type and therefore it always
    /// resets the `store_to_extension` configuration.
    ///
    /// ```
    /// # use {
    /// #     axum::{Router, routing},
    /// #     axum_jwt::{Decoder, Token, jsonwebtoken::DecodingKey},
    /// # };
    /// # fn check_access(t: &Token) -> bool { true }
    /// # async fn hello() {}
    /// # let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
    /// let app = Router::new()
    ///     .route("/", routing::get(hello))
    ///     .layer(
    ///         axum_jwt::layer(decoder)
    ///             // Incorrect order!
    ///             .store_to_extension()
    ///             .with_filter(check_access),
    ///     );
    /// # let _: Router = app;
    /// ```
    ///
    /// </section>
    ///
    /// # Read header only
    ///
    /// If you just need to retrieve the token without any additional payload,
    /// omit the `with_filter` call and use the `Token` type without
    /// a parameter:
    ///
    /// ```
    /// use {
    ///     axum::{Extension, Router, routing},
    ///     axum_jwt::{Decoder, Token, jsonwebtoken::DecodingKey},
    /// };
    ///
    /// async fn hello(Extension(t): Extension<Token>) -> String {
    ///     // Access the parsed token
    ///     format!("Decoded with {:?} algorithm", t.header.alg)
    /// }
    ///
    /// let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
    ///
    /// let app = Router::new()
    ///     .route("/", routing::get(hello))
    ///     .layer(axum_jwt::layer(decoder).store_to_extension());
    /// # let _: Router = app;
    /// ```
    pub fn store_to_extension(mut self) -> Self
    where
        I: Clone + Send + Sync + 'static,
    {
        self.store = |claims, extensions| {
            extensions.insert(claims);
        };

        self
    }
}

impl<I, H> JwtLayer<I, H> {
    /// Applies a token extractor to the layer.
    ///
    /// By default, the token is extracted from the `Authorization` header using
    /// the `Bearer` scheme. If you want to change this behavior, create a new type
    /// and implement [`Extract`] for it. Then, you can pass this type into the
    /// layer configuration:
    ///
    /// ```
    /// use {
    ///     axum::{Extension, Router, http::request::Parts, routing},
    ///     axum_jwt::{Decoder, Extract, Token, jsonwebtoken::DecodingKey},
    /// };
    ///
    /// struct Custom;
    ///
    /// impl Extract for Custom {
    ///     fn extract(parts: &mut Parts) -> Option<&str> {
    ///         parts.headers.get("X-Auth-Token")?.to_str().ok()
    ///     }
    /// }
    ///
    /// async fn hello(Extension(t): Extension<Token>) -> String {
    ///     format!("Decoded with {:?} algorithm", t.header.alg)
    /// }
    ///
    /// let decoder = Decoder::from_key(DecodingKey::from_secret(b"secret"));
    ///
    /// let app = Router::new()
    ///     .route("/", routing::get(hello))
    ///     .layer(axum_jwt::layer(decoder).with_extract(Custom));
    /// # let _: Router = app;
    /// ```
    pub fn with_extract<X>(self, extract: X) -> JwtLayer<I, H, X>
    where
        X: Extract,
    {
        _ = extract;
        JwtLayer {
            decoder: self.decoder,
            validate: self.validate,
            store: self.store,
            extract: PhantomData,
        }
    }
}

impl<I, H, X> Clone for JwtLayer<I, H, X>
where
    H: Clone,
{
    fn clone(&self) -> Self {
        Self {
            decoder: self.decoder.clone(),
            validate: self.validate.clone(),
            store: self.store,
            extract: PhantomData,
        }
    }
}

impl<I, H, X> fmt::Debug for JwtLayer<I, H, X> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("JwtLayer")
            .field("decoder", &self.decoder)
            .field("validate", &"..")
            .field("store", &"..")
            .field("extract", &any::type_name::<H>())
            .finish()
    }
}

impl<S, I, H, X> Layer<S> for JwtLayer<I, H, X>
where
    H: Clone,
{
    type Service = Jwt<S, I, H>;

    fn layer(&self, svc: S) -> Self::Service {
        Jwt {
            svc,
            decoder: self.decoder.clone(),
            validate: self.validate.clone(),
            store: self.store,
            extract: PhantomData,
        }
    }
}

/// Trait for additional token validation.
pub trait Validate<I> {
    type Output: Output;
    fn validate(&mut self, input: &Token<I>) -> Self::Output;
}

/// The output value of the [validation](Validate).
pub trait Output {
    fn output(self) -> Option<Response>;
}

impl<E> Output for Result<(), E>
where
    E: IntoResponse,
{
    fn output(self) -> Option<Response> {
        self.err().map(E::into_response)
    }
}

impl Output for bool {
    fn output(self) -> Option<Response> {
        if self {
            None
        } else {
            Some(StatusCode::UNAUTHORIZED.into_response())
        }
    }
}

/// Discards any token data and returns success.
#[derive(Clone)]
pub struct Discard;

impl<I> Validate<I> for Discard {
    type Output = bool;

    fn validate(&mut self, _: &Token<I>) -> Self::Output {
        true
    }
}

impl<F, I, O> Validate<I> for F
where
    F: FnMut(&Token<I>) -> O,
    I: DeserializeOwned,
    O: Output,
{
    type Output = O;

    fn validate(&mut self, input: &Token<I>) -> Self::Output {
        self(input)
    }
}

/// Axum [middleware] for token validation.
///
/// [middleware]: https://docs.rs/axum/latest/axum/middleware/index.html
///
/// To configure the layer and create the middleware service, call
/// the [`layer`] function.
pub struct Jwt<S, I, H = Discard, X = Bearer> {
    svc: S,
    decoder: Decoder,
    validate: H,
    store: fn(Token<I>, &mut Extensions),
    extract: PhantomData<fn() -> X>,
}

impl<S, I, H, X> Clone for Jwt<S, I, H, X>
where
    S: Clone,
    H: Clone,
{
    fn clone(&self) -> Self {
        Self {
            svc: self.svc.clone(),
            decoder: self.decoder.clone(),
            validate: self.validate.clone(),
            store: self.store,
            extract: PhantomData,
        }
    }
}

impl<S, I, H, X> fmt::Debug for Jwt<S, I, H, X>
where
    S: fmt::Debug,
{
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("Jwt")
            .field("svc", &self.svc)
            .field("decoder", &self.decoder)
            .field("validate", &"..")
            .field("store", &"..")
            .field("extract", &any::type_name::<X>())
            .finish()
    }
}

impl<S, I, H, X> Service<Request> for Jwt<S, I, H, X>
where
    S: Service<Request> + Clone,
    I: DeserializeOwned,
    H: Validate<I>,
    X: Extract,
    Result<S::Response, S::Error>: IntoResponse,
{
    type Response = Response;
    type Error = Infallible;
    type Future = JwtFuture<S>;

    fn poll_ready(&mut self, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        Poll::Ready(Ok(()))
    }

    fn call(&mut self, req: Request) -> Self::Future {
        let validate = |parts| -> Result<Token<I>, Error> {
            let token = X::extract(parts).ok_or(Error::Extract)?;
            let TokenData { header, claims }: TokenData<I> =
                self.decoder.decode(token).map_err(Error::Jwt)?;

            Ok(Token::new(header, claims))
        };

        let (mut parts, body) = req.into_parts();
        match validate(&mut parts) {
            Ok(token) => {
                if let Some(res) = self.validate.validate(&token).output() {
                    return JwtFuture::ready(res);
                }

                (self.store)(token, &mut parts.extensions);

                let req = Request::from_parts(parts, body);
                let clone = self.svc.clone();
                let svc = mem::replace(&mut self.svc, clone);
                JwtFuture::not_ready(svc, req)
            }
            Err(e) => JwtFuture::ready(e.into_response()),
        }
    }
}

pin_project_lite::pin_project! {
    /// Middleware future.
    pub struct JwtFuture<S>
    where
        S: Service<Request>,
    {
        #[pin]
        state: State<S, S::Future>,
    }
}

impl<S> JwtFuture<S>
where
    S: Service<Request>,
{
    fn not_ready(svc: S, req: Request) -> Self {
        Self {
            state: State::NotReady { svc, req },
        }
    }

    fn ready(res: Response) -> Self {
        Self {
            state: State::Ready { res },
        }
    }
}

impl<S> Future for JwtFuture<S>
where
    S: Service<Request>,
    Result<S::Response, S::Error>: IntoResponse,
{
    type Output = Result<Response, Infallible>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut state = self.project().state;
        let res = loop {
            match state.as_mut().project() {
                StateProj::NotReady { svc, req } => {
                    if let Err(e) = task::ready!(svc.poll_ready(cx)) {
                        state.set(State::Done);
                        break Err(e).into_response();
                    }

                    let req = mem::take(req);
                    let fut = svc.call(req);
                    state.set(State::Called { fut });
                }
                StateProj::Called { fut } => {
                    let res = task::ready!(fut.poll(cx));
                    state.set(State::Done);
                    break res.into_response();
                }
                StateProj::Ready { res } => {
                    let res = mem::take(res);
                    state.set(State::Done);
                    break res;
                }
                StateProj::Done => panic!("polled after completion"),
            }
        };

        Poll::Ready(Ok(res))
    }
}

pin_project_lite::pin_project! {
    #[project = StateProj]
    enum State<S, F> {
        NotReady { svc: S, req: Request },
        Called {
            #[pin]
            fut: F,
        },
        Ready { res: Response },
        Done,
    }
}