Implementing JWT Authentication in Rust

Hey there! Following on from our ShuttleBytes talk which we held on Tuesday, we're going to talk about how you can implement authentication using JSON Web Tokens (JWTs) in Rust.

What is a JWT?

A JSON Web Token (JWT) is a compact, URL-safe way to transfer data ("claims") between two parties over the Web. The data is typically encoded using a JSON Web Signature or as part of a JSON Web Encryption (JWE) structure, and can be encrypted (and signed!).

JWTs are a popular option when deciding on an auth strategy. The client stores all the information via the JWT, allowing for a stateless API. This can make user authentication much easier in some cases. While unencrypted and unsigned JWTs can be manipulated, it is simple for the server to be able to disregard manipulated JWTs as long as the secret key used to create them remains secret.

Getting started

To get started, let's initialise a project using shuttle init, making sure to pick Axum as the framework.

Then we'll add our dependencies:

cargo add axum-extra@0.9.2 -F typed-header
cargo add chrono@0.4.34 -F serde,clock
cargo add jsonwebtoken@9.2.0
cargo add once_cell@1.19.0
cargo add serde@1.0.196 -F derive
cargo add serde-json@1.0.113

Writing our web service

Setting up keys

To start with, we'll need to declare a struct that holds decoding and encoding key, with a method that can take a &[u8] (u8 slice) to generate the struct:

use jsonwebtoken::{DecodingKey, EncodingKey};

struct Keys {
    encoding: EncodingKey,
    decoding: DecodingKey,
}

impl Keys {
    fn new(secret: &[u8]) -> Self {
        Self {
            encoding: EncodingKey::from_secret(secret),
            decoding: DecodingKey::from_secret(secret),
        }
    }
}

This struct needs to be generated from a secret key as it's what we will use to generate the JWTs from. For this example, we'll be randomly generating our bytes from a String and then turning it into bytes. It will then be stored in a once_cell::LazyCell that can be accessed globally in our application:

use once_cell::sync::Lazy;

static KEYS: Lazy<Keys> = Lazy::new(|| {
    let secret = Alphanumeric.sample_string(&mut rand::thread_rng(), 60);
    Keys::new(secret.as_bytes())
});

Note that there's many different sources you can use to generate the byte array from for this - generating a random string and turning it into bytes is just one of them. For production usage, you may also want to use a cryptographically safe algorithm.

Writing our JWT Claim

The next step is to implement our claim. A claim (in JWT context) is the data transmitted by a JWT and gets encoded or decoded by the server. We can write our own Claim implementation by creating a struct that holds a username and expiry date, then implementing the FromRequestParts trait (from Axum) for the struct. This allows us to use it as an Axum extractor and saves us from having to implement any middleware!

However before we write the actual implementation itself, FromRequestParts requires that we have a custom error type. We can write one that represents JWT failures and implement IntoResponse for it - which will then allow us to use it in the implementation.

use axum::response::{ IntoResponse, Response };
use axum::http::StatusCode;
use serde_json::json;

pub enum AuthError {
    InvalidToken,
    WrongCredentials,
    TokenCreation,
    MissingCredentials,
}

impl IntoResponse for AuthError {
    fn into_response(self) -> Response {
        let (status, error_message) = match self {
            AuthError::WrongCredentials => (StatusCode::UNAUTHORIZED, "Wrong credentials"),
            AuthError::MissingCredentials => (StatusCode::BAD_REQUEST, "Missing credentials"),
            AuthError::TokenCreation => (StatusCode::INTERNAL_SERVER_ERROR, "Token creation error"),
            AuthError::InvalidToken => (StatusCode::BAD_REQUEST, "Invalid token"),
        };
        let body = Json(json!({
            "error": error_message,
        }));
        (status, body).into_response()
    }
}

Implementing IntoResponse for AuthError allows it to be used as the Rejection type in the FromRequestParts trait. Note that to be able to return AuthError in the FromPartsRequest trait, we use map_err to turn the error type into AuthError so that it can be propagated. We also use de-structuring here to extract the bearer struct from the TypedHeader<Authorization<Bearer>> type as it's much easier to access.

use serde::{ Serialize, Deserialize };
use axum::{ http::{ request::Parts }, extract::FromRequestParts, RequestPartsExt };

#[derive(Debug, Serialize, Deserialize)]
pub struct Claims {
    username: String,
    exp: usize,
}

#[async_trait]
impl<S> FromRequestParts<S> for Claims where S: Send + Sync {
    type Rejection = AuthError;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        // Extract the token from the authorization header
        let TypedHeader(Authorization(bearer)) = parts
            .extract::<TypedHeader<Authorization<Bearer>>>().await
            .map_err(|_| AuthError::InvalidToken)?;
        // Decode the user data
        let token_data = decode::<Claims>(
            bearer.token(),
            &KEYS.decoding,
            &Validation::default()
        ).map_err(|_| AuthError::InvalidToken)?;

        Ok(token_data.claims)
    }
}

Now that we're done setting up our boilerplate for how the JWT should work, we can move onto creating our routes!

Creating routes

The next step is to write our endpoint when authorizing a user - when returning the token, we'll create a new AuthBody with the token in it and the type of token it is. We'll be using this later on:

#[derive(Debug, Serialize)]
struct AuthBody {
    access_token: String,
    token_type: String,
}

impl AuthBody {
    fn new(access_token: String) -> Self {
        Self {
            access_token,
            token_type: "Bearer".to_string(),
        }
    }
}

Now that we've created our AuthBody, we can create an endpoint that will take a client ID and secret and verify it. Then it will create a claim, encode it and return it as JSON.

use axum::Json;
use chrono::Utc;

#[derive(Debug, Deserialize)]
struct AuthPayload {
    client_id: String,
    client_secret: String,
}

async fn authorize(Json(payload): Json<AuthPayload>) -> Result<Json<AuthBody>, AuthError> {
    // Check if the user sent the credentials
    if payload.client_id.is_empty() || payload.client_secret.is_empty() {
        return Err(AuthError::MissingCredentials);
    }
    // Here, basic verification is used but normally you would use a database
    if &payload.client_id != "foo" || &payload.client_secret != "bar" {
        return Err(AuthError::WrongCredentials);
    }

    // create the timestamp for the expiry time - here the expiry time is 1 day
    // in production you may not want to have such a long JWT life
    let exp = (Utc::now().naive_utc() + chrono::naive::Days::new(1)).timestamp() as usize;
    let claims = Claims {
        username: payload.client_id,
        exp,
    };
    // Create the authorization token
    let token = encode(&Header::default(), &claims, &KEYS.encoding).map_err(
        |_| AuthError::TokenCreation
    )?;

    // Send the authorized token
    Ok(Json(AuthBody::new(token)))
}

Now that we have a route to generate the JWT, we can create a route to try our token out!

async fn protected(claims: Claims) -> String {
    // Send the protected data to the user
    format!("Welcome to the protected area, {}!", claims.username)
}

Now that all of the routes have been written we can hook this all back up to our main function:

use axum::{Router, routing::{get, post}};

#[shuttle_runtime::main]
async fn main() -> shuttle_axum::ShuttleAxum {
    let router = Router::new()
        .route("/", get(hello_world))
        .route("/protected", get(protected))
        .route("/login", post(authorize));

    Ok(router.into())
}

Testing

To test that the API works, we can use shuttle run to serve the API locally.

To test our /login endpoint, we'll send a cURL request to it with the data:

curl localhost:8000/login -H 'Content-Type: application/json/' \
 -d '{"client_id":"foo","client_secret":"bar"}'

When we run this, it should output some JSON containing the JWT and type of token - which should be the Bearer type. This should be a Bearer token because in FromRequestParts we extract from the Authorization: Bearer ... header.

To now test the protected route, we need to use the JWT token that was sent to us in the Authorization header:

curl localhost:8000/protected -H 'Authorization: Bearer <your-jwt-here>'

You should receive a text response that looks like this:

Welcome to the protected area, foo!

Deploying

Now that we've written our whole project, we can deploy! Type in shuttle deploy and press enter (add --ad flag if on a dirty Git branch to allow dirty deploys). When finished, our terminal should print out all of the data about our deployment and project and a link to the live project!

Extending this project

Interested in extending this project? Here's a couple ideas:

Use Postgres to store user logins.
Try using encryption and signing to strengthen your JWTs, as well as storing them in a cookie.
Try adding integration tests so you don't need manual testing!