#Axum duplicate tower-governer key extractor and axum extractor issue

14 messages · Page 1 of 1 (latest)

zinc crystal
#

I have this axum extractor for extracting an authenticated user from a bearer token:

impl FromRequestParts<AppState> for AuthenticatedUser {
    type Rejection = Error;

    async fn from_request_parts(
        parts: &mut Parts,
        state: &AppState,
    ) -> Result<Self, Self::Rejection> {
        let TypedHeader(Authorization(bearer)) = parts
            .extract::<TypedHeader<Authorization<Bearer>>>()
            .await
            .map_err(|_| AuthError::MissingToken)?;

        let token = bearer.token();
        let bearer_token = format!("Bearer {}", token);

        match Claims::from_bearer_token(&bearer_token, &state.jwks_cache).await {
            Ok(claims) => {
                // Extract user information from claims
                let user_id = claims
                    .user_id()
                    .parse()
                    .map_err(|_| AuthError::InvalidUserId)?;
                // ...

                Ok(AuthenticatedUser {
                    user_id,
                    // ...
                })
            }
            Err(err) => Err(AuthError::VerificationFailed(err.to_string()).into()),
        }
    }
}

I now want to add tower-governer to implement rate limiting, where the rate limit key is per authenticated user. For this, I need to implement KeyExtractor for a type. Ideally, I want the key extractor to be the users Uuid, however I don't want each request to perform the bearer token extraction and validation twice on every request (once for KeyExtractor rate limiting, again for the handler to extract the AuthenticatedUser). Any suggestions on how I could make this work be done just once, and still have access to the authenticated user from my extractor?

For now, I've just done the key extraction on the bearer token directly, since thats quite cheap.

#

Here's my current KeyExtractor implementation:

impl KeyExtractor for AuthenticatedUser {
    type Key = String;

    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, GovernorError> {
        let bearer = Bearer::decode(req.headers().get(AUTHORIZATION).ok_or_else(|| {
            let err = Error::from(AuthError::MissingToken);
            GovernorError::Other {
                code: err.code,
                msg: Some(err.msg.into_owned()),
                headers: Some(err.headers),
            }
        })?)
        .ok_or_else(|| {
            let err = Error::from(AuthError::InvalidToken);
            GovernorError::Other {
                code: err.code,
                msg: Some(err.msg.into_owned()),
                headers: Some(err.headers),
            }
        })?;

        Ok(bearer.token().to_string())
    }
}

But still, I'd prefer the key to be the users Uuid

tall crypt
#

iiuc, you could do something like axum-extra's Cached extractor (not literally; but your own version)

the idea is to check the request extensions first, and otherwise perform your expensive operation and then put the result into extensions

#

that way, nomatter the order, extraction should only "actually" be performed once

#

and every other access should get the cached version from the extensions

zinc crystal
#

Yea thats what I'm looking into now. I am trying req.extensions().get::<AuthenticatedUser>().ok_or(...), but it seems like the authenticated user doesn't exist when the key extractor runs. And I can't insert an extension into the request since I only have an immutable reference

#

For example, I tried this:

#[derive(Clone, Copy)]
pub struct AuthenticatedUserKey;

impl KeyExtractor for AuthenticatedUserKey {
    type Key = Uuid;

    fn extract<T>(&self, req: &Request<T>) -> Result<Self::Key, GovernorError> {
        let user = req.extensions().get::<AuthenticatedUser>().ok_or_else(|| {
            let err = Error::from(AuthError::MissingToken);
            GovernorError::Other {
                code: err.code,
                msg: Some(err.msg.into_owned()),
                headers: Some(err.headers),
            }
        })?;
        Ok(user.user_id)
    }
}

But AuthenticatedUser doesn't exist, despite my route handler having it as an extractor

tall crypt
#

you wouldn't be able to run the validation in there anyway since it's async

#

what you could do is move your auth logic to a middleware, then have both your key extractor and your AuthenticatedUser extractor just access extensions

#

the middleware can forward any requests as-is if the header is missing, and otherwise perform validation

#

the extractor can return unauthorized on a missing extension

zinc crystal
#

oh yea that might be the way to go! And the auth middleware can still succeed even if there's no bearer token, but AuthenticatedUser extractor will fail, so I'll still get the same behvaiour which is nice

tall crypt
#

yep