CryptoFlow: Building a secure and scalable system with Axum and SvelteKit - Part 2

Introduction

Part 1 saw us implement the login process and many other utility functions and setups with which the implementation was made seamless. In this part, we will build on that to allow comprehensive user management as promised. We'll write some middleware to help check whether or not a request is authenticated as well as enable user registration, email verification via token (a 6-digit cryptographically random token will be used), and user logout.

Source code

The source code for this series is hosted on GitHub via:

Sirneij / cryptoflow

A Q&A web application to demostrate how to build a secured and scalable client-server application with axum and sveltekit

CryptoFlow

CryptoFlow is a full-stack web application built with Axum and SvelteKit. It's a Q&A system tailored towards the world of cryptocurrency!

Its building process is explained in this series of articles.

View on GitHub

Implementation

Step 1: User session extraction via middleware

In almost any reasonable system, some services shouldn't be available to anonymous users. These kinds of services are protected by some sort of authentication and authorization mechanism. In our system, we already set up a mechanism to authenticate users. However, we haven't figured out a way to intercept every request coming into our system and check whether or not they should be allowed in. That's what we'll do here! We want to extract certain data from a request and make the decision to reject or accept such a request depending on the data extracted from it. To achieve this, we will roll out a middleware:

// backend/src/utils/middleware.rs
use crate::startup::AppState;
use crate::utils::get_user_id_from_session;
use axum::{extract::Request, middleware::Next};
use axum::{
    extract::State,
    response::{IntoResponse, Response},
};
use axum_extra::extract::PrivateCookieJar;

#[tracing::instrument(
    name = "validate_authentication_session",
    skip(cookies, state, req, next)
)]
pub async fn validate_authentication_session(
    cookies: PrivateCookieJar,
    State(state): State<AppState>,
    req: Request,
    next: Next,
) -> Result<impl IntoResponse, Response> {
    // Use the utility function to get the user ID from the session
    match get_user_id_from_session(&cookies, &state.redis_store, false).await {
        Ok(_user_id) => Ok(next.run(req).await),
        Err(error) => Err(error.into_response()),
    }
}

One major selling point of axum is its ease of writing middleware. Just look at what we have up there! It's just a beauty!!! Though some things have been abstracted away (especially retrieving the cookie and processing it), it is still unmatched! That said, writing middleware in axum can take different shapes, we decided to stick to using the axum::middleware::from_fn (axum::middleware::from_fn_with_state precisely) because our middleware is just for this system, built with axum, to extract request cookies and make decisions. In axum, a function regarded as middleware should take the Request and Next types (in v0.6, these types were required to be bounded by <B> generics) alongside other arguments and return something that implements IntoResponse. If all requirements are met, returning next.run(req).await allows the request to proceed to the main service it wanted.

An integral part of the middleware above is the get_user_id_from_session and it has this definition:

// backend/src/utils/user.rs
use crate::utils::{CustomAppError, ErrorContext};
use axum_extra::extract::PrivateCookieJar;
use bb8_redis::bb8;
use uuid::Uuid;

#[tracing::instrument(
    name = "get_user_id_from_session",
    skip(cookies, redis_store, is_logout)
)]
pub async fn get_user_id_from_session(
    cookies: &PrivateCookieJar,
    redis_store: &bb8::Pool<bb8_redis::RedisConnectionManager>,
    is_logout: bool,
) -> Result<(Uuid, String), CustomAppError> {
    let session_id = cookies
        .get("sessionid")
        .map(|cookie| cookie.value().to_owned())
        .ok_or_else(|| {
            CustomAppError::from((
                "Session ID not found because you are not authenticated".to_string(),
                ErrorContext::UnauthorizedAccess,
            ))
        })?;

    let mut redis_con = redis_store.get().await.map_err(|_| {
        CustomAppError::from((
            "Failed to get redis connection".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?;

    tracing::debug!("Session ID: {}", session_id);

    let user_id: String = bb8_redis::redis::cmd("GET")
        .arg(&session_id)
        .query_async(&mut *redis_con)
        .await
        .map_err(|_| {
            CustomAppError::from((
                "You are not authorized since you don't seem to have been authenticated"
                    .to_string(),
                ErrorContext::UnauthorizedAccess,
            ))
        })?;

    let user_uuid = Uuid::parse_str(&user_id).map_err(|_| {
        CustomAppError::from((
            "Invalid user ID format".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?;

    if is_logout {
        bb8_redis::redis::cmd("DEL")
            .arg(&session_id)
            .query_async::<_, i64>(&mut *redis_con)
            .await
            .map_err(|_| {
                CustomAppError::from((
                    "Failed to delete session ID from redis".to_string(),
                    ErrorContext::InternalServerError,
                ))
            })?;
    }

    Ok((user_uuid, session_id))
}

The function first tries retrieving the session_id stored in the cookies (as done in the login_user in part 1). If successful, it then goes to do the same from the redis. This is the second layer of security. In case the session_id isn't found in either, appropriate errors are returned. The function takes a flag, is_logout, which determines whether or not the session_id will be deleted from the redis (which is the case for the logout operation).

With that concluded, let's use it. But before then, let's write the logout handler.

Step 2: Logging users out

We'll create a new file, backend/src/routes/users/logout.rs, and fill it with this:

use crate::startup::AppState;
use crate::utils::CustomAppError;
use crate::utils::SuccessResponse;
use axum::{extract::State, http::StatusCode, response::IntoResponse};
use axum_extra::extract::cookie::{Cookie, PrivateCookieJar};

#[axum::debug_handler]
#[tracing::instrument(name = "logout_user", skip(cookies, state))]
pub async fn logout_user(
    cookies: PrivateCookieJar,
    State(state): State<AppState>,
) -> Result<(PrivateCookieJar, impl IntoResponse), CustomAppError> {
    // Get user_id and session_id from cookie and delete it
    let (_, _) = crate::utils::get_user_id_from_session(&cookies, &state.redis_store, true).await?;

    Ok((
        cookies.remove(Cookie::from("sessionid")),
        SuccessResponse {
            message: "The unauthentication process was successful.".to_string(),
            status_code: StatusCode::OK.as_u16(),
        }
        .into_response(),
    ))
}

It is very simple. We utilized the get_user_id_from_session utility function, passing is_logout=true to delete the session from redis in case everything goes as planned. Remember that to propagate any action on PrivateCookieJar, it must be returned with your response.

Next, let's fix it up with our routes:

// src/routes/users/mod.rs
...
mod logout

pub fn users_routes(state: crate::startup::AppState) -> Router<crate::startup::AppState> {
    Router::new()
        .route("/logout", post(logout::logout_user))
        .route_layer(axum::middleware::from_fn_with_state(
            state.clone(),
            validate_authentication_session,
        ))
        ...
}

users_routes now needs to take AppState as an argument (you need to pass this in backend/src/startup.rs accordingly). This is because from_fn_with_state needs an instance of it to work. Our middleware was also added up. This ensures that only authenticated users can access /api/users/logout route. As demonstrated, all routes before the application of a middleware have that middleware applied to them while those after are free from the middleware's requirements.

You can now test the login/logout process and everything should be great.

Step 3: Registering users

Before users can log in (not to mention logout), they need a (verified) account. It's time to start the process of getting such an account:

// backend/src/routes/users/register.rs
use crate::models::NewUser;
use crate::startup::AppState;
use crate::utils::SuccessResponse;
use crate::utils::{CustomAppError, CustomAppJson, ErrorContext};
use argon2::password_hash::rand_core::{OsRng, RngCore};
use axum::{extract::State, http::StatusCode, response::IntoResponse};
use sha2::{Digest, Sha256};

#[axum::debug_handler]
#[tracing::instrument(name = "register_user", skip(state, new_user),fields(user_email = new_user.email, user_first_name = new_user.first_name, user_last_name = new_user.last_name))]
pub async fn register_user(
    State(state): State<AppState>,
    CustomAppJson(new_user): CustomAppJson<NewUser>,
) -> Result<impl IntoResponse, CustomAppError> {
    let hashed_password = crate::utils::hash_password(&new_user.password.as_bytes()).await;

    let user = state
        .db_store
        .create_user(
            &new_user.first_name,
            &new_user.last_name,
            &new_user.email,
            &hashed_password,
        )
        .await?;

    // Generate a truly random activation code for the user using argon2::password_hash::rand_core::OsRng
    let activation_code = (OsRng.next_u32() % 900000 + 100000).to_string();
    // Hash the activation code
    let mut hasher = Sha256::new();
    hasher.update(activation_code.as_bytes());
    let hashed_activation_code = format!("{:x}", hasher.finalize());

    // Save activation code in redis
    let mut redis_con = state.redis_store.get().await.map_err(|_| {
        CustomAppError::from((
            "Failed to get redis connection".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?;

    let settings = crate::settings::get_settings().map_err(|_| {
        CustomAppError::from((
            "Failed to read settings".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?;
    let activation_code_expiration_in_seconds = settings.secret.token_expiration * 60;

    bb8_redis::redis::cmd("SET")
        .arg(user.id.to_string())
        .arg(hashed_activation_code)
        .arg("EX")
        .arg(activation_code_expiration_in_seconds)
        .query_async::<_, String>(&mut *redis_con)
        .await
        .map_err(|_| {
            CustomAppError::from((
                "Failed to save activation code".to_string(),
                ErrorContext::InternalServerError,
            ))
        })?;

    // Send activation code to user's email
    crate::utils::send_multipart_email(
        "Welcome to CryptoFlow with Rust (axum) and SvelteKit".to_string(),
        user,
        state.clone(),
        "user_welcome.html",
        activation_code,
    )
    .await
    .map_err(|_| {
        CustomAppError::from((
            "Failed to send activation email".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?;

    Ok(SuccessResponse {
        message: "Registration complete! Check your email for a verification code to activate your account.".to_string(),
        status_code: StatusCode::CREATED.as_u16(),
    }.into_response())
}

As usual, error handling took a bunch of space but it's okay! We just started by hashing the provided password and went straight to create the user in our database using the create_user method on db_store (was written in part 1). This creation sets the is_active to false pending when the user's email is confirmed. Then, a 6-digit cryptographically random string is generated and its sha256 hash (we don't want to save the plain token) is saved temporarily (for token_expiration * 60 seconds, 15 * 60 by default) in redis (the reason for sending a token instead of sending a verification link was explained here). The token is then sent to the email used in registration. Let's take a look at the email-sending utility:

// backend/src/utils/email.rs
use lettre::AsyncTransport;

#[tracing::instrument(
    name = "Generic e-mail sending function.",
    skip(
        subject,
        html_content,
        text_content
    ),
    fields(
        recipient_email = %user.email,
        recipient_first_name = %user.first_name,
        recipient_last_name = %user.last_name
    )
)]
pub async fn send_email(
    user: crate::models::UserVisible,
    subject: impl Into<String>,
    html_content: impl Into<String>,
    text_content: impl Into<String>,
) -> Result<(), String> {
    let settings = crate::settings::get_settings().expect("Failed to read settings.");

    let email = lettre::Message::builder()
        .from(
            format!(
                "{} <{}>",
                "CryptoFlow with axum and SvelteKit",
                settings.email.host_user.clone()
            )
            .parse()
            .map_err(|e| {
                tracing::error!("Could not parse 'from' email address: {:#?}", e);
                format!("Could not parse 'from' email address: {:#?}", e)
            })?,
        )
        .to(format!(
            "{} <{}>",
            [user.first_name, user.last_name].join(" "),
            user.email
        )
        .parse()
        .map_err(|e| {
            tracing::error!("Could not parse 'to' email address: {:#?}", e);
            format!("Could not parse 'to' email address: {:#?}", e)
        })?)
        .subject(subject)
        .multipart(
            lettre::message::MultiPart::alternative()
                .singlepart(
                    lettre::message::SinglePart::builder()
                        .header(lettre::message::header::ContentType::TEXT_PLAIN)
                        .body(text_content.into()),
                )
                .singlepart(
                    lettre::message::SinglePart::builder()
                        .header(lettre::message::header::ContentType::TEXT_HTML)
                        .body(html_content.into()),
                ),
        )
        .unwrap();

    let creds = lettre::transport::smtp::authentication::Credentials::new(
        settings.email.host_user,
        settings.email.host_user_password,
    );

    // Open a remote connection to gmail
    let mailer: lettre::AsyncSmtpTransport<lettre::Tokio1Executor> =
        lettre::AsyncSmtpTransport::<lettre::Tokio1Executor>::relay(&settings.email.host)
            .unwrap()
            .credentials(creds)
            .build();

    // Send the email
    match mailer.send(email).await {
        Ok(_) => {
            tracing::info!("Email sent successfully.");
            Ok(())
        }
        Err(e) => {
            tracing::error!("Could not send email: {:#?}", e);
            Err(format!("Could not send email: {:#?}", e))
        }
    }
}

#[tracing::instrument(
    name = "Generic multipart e-mail sending function.",
    skip(
        user,
        state,
        template_name
    ),
    fields(
        recipient_user_id = %user.id,
        recipient_email = %user.email,
        recipient_first_name = %user.first_name,
        recipient_last_name = %user.last_name
    )
)]
pub async fn send_multipart_email(
    subject: String,
    user: crate::models::UserVisible,
    state: crate::startup::AppState,
    template_name: &str,
    issued_token: String,
) -> Result<(), String> {
    let settings = crate::settings::get_settings().expect("Unable to load settings.");
    let title = subject.clone();

    let now = chrono::Local::now();
    let expiration_time = now + chrono::Duration::minutes(settings.secret.token_expiration);
    let exact_time = expiration_time.format("%A %B %d, %Y at %r").to_string();

    let template = state.env.get_template(template_name).unwrap();
    let ctx = minijinja::context! {
        title => &title,
        user_id => &user.id,
        domain => &settings.frontend_url,
        token => &issued_token,
        expiration_time => &settings.secret.token_expiration,
        exact_time => &exact_time,
    };
    let html_text = template.render(ctx).unwrap();

    let text = format!(
        r#"
        Thanks for signing up for a CryptoFlow with Rust (axum) and SvelteKit. We're excited to have you on board!

        For future reference, your user ID number is {}.

        Please visit {}/auth/activate/{} and input the token below to activate your account:

        {}


        Please note that this is a one-time use token and it will expire in {} minutes ({}).


        Thanks,

        CryptoFlow with Rust (axum) and SvelteKit Team
        "#,
        user.id,
        settings.frontend_url,
        user.id,
        issued_token,
        settings.secret.token_expiration,
        exact_time
    );

    tokio::spawn(send_email(user, subject, html_text, text));
    Ok(())
}

Briefly (a detailed explanation can be found here), it uses the lettre crate to build and send emails. The emails support HTML and have plaintext fallback in case the email server doesn't support HTML. For templating, we used minijinja which we will set up next!

// backend/src/startup.rs
...
#[derive(Clone)]
pub struct AppState {
    ...
    pub env: minijinja::Environment<'static>,
    ...
}
...

async fn run(
    listener: tokio::net::TcpListener,
    store: crate::store::Store,
    settings: crate::settings::Settings,
) {
    ...
    let mut env = minijinja::Environment::new();
    env.set_loader(minijinja::path_loader("templates"));
    let app_state = AppState {
        ...
        env,
        ...
    };
    ...
}
...

Kindly create the templates folder at the root of the backend folder. In it, put this HTML file:

<!--backend/templates/user_welcome.html-->
<!DOCTYPE html>
<html>
  <head>
    <meta name="viewport" content="width=device-width" />
    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8" />
    <title>{{ title }}</title>
  </head>
  <body>
    <table style="background: #ffffff; border-radius: 1rem; padding: 30px 0px">
      <tbody>
        <tr>
          <td style="padding: 0px 30px">
            <h3 style="margin-bottom: 0px; color: #000000">Hello,</h3>
            <p>
              Thanks for signing up for a CryptoFlow with Rust (axum) and
              SvelteKit. We're excited to have you on board!
            </p>
          </td>
        </tr>
        <tr>
          <td style="padding: 0px 30px">
            <p>For future reference, your user ID is #{{ user_id }}.</p>
            <p>
              Please visit
              <a href="{{ domain }}/auth/activate/{{ user_id }}">
                {{ domain }}/auth/activate/{{ user_id }}
              </a>
              and input the OTP below to activate your account:
            </p>
          </td>
        </tr>

        <tr>
          <td style="padding: 10px 30px; text-align: center">
            <strong style="display: block; color: #00a856">
              One Time Password (OTP)
            </strong>
            <table style="margin: 10px 0px" width="100%">
              <tbody>
                <tr>
                  <td
                    style="
                      padding: 25px;
                      background: #faf9f5;
                      border-radius: 1rem;
                    "
                  >
                    <strong
                      style="
                        letter-spacing: 8px;
                        font-size: 24px;
                        color: #000000;
                      "
                    >
                      {{ token }}
                    </strong>
                  </td>
                </tr>
              </tbody>
            </table>
            <small style="display: block; color: #6c757d; line-height: 19px">
              <strong>
                Please note that this is a one-time use token and it will expire
                in {{ expiration_time }} minutes ({{ exact_time }}).
              </strong>
            </small>
          </td>
        </tr>

        <tr>
          <td style="padding: 0px 30px">
            <hr style="margin: 0" />
          </td>
        </tr>
        <tr>
          <td style="padding: 30px 30px">
            <table>
              <tbody>
                <tr>
                  <td>
                    <strong>
                      Kind Regards,<br />
                      CryptoFlow with Rust (axum) and SvelteKit Team
                    </strong>
                  </td>
                  <td></td>
                </tr>
              </tbody>
            </table>
          </td>
        </tr>
      </tbody>
    </table>
  </body>
</html>

It's a simple but nice-looking email template built in this series.

To wrap up, let's write the user activation or token verification handler.

Step 4: User activation and token verification

Currently, the registration process isn't complete yet as the registered users cannot log in to our system. Their email addresses must be verified before such can be allowed. Here is the function that handles the verification:

// backend/src/routes/users/activate_account.rs
use crate::{
    models::ActivateUser,
    startup::AppState,
    utils::{CustomAppError, CustomAppJson, ErrorContext, SuccessResponse},
};
use axum::{extract::State, http::StatusCode, response::IntoResponse};
use sha2::{Digest, Sha256};

#[axum::debug_handler]
#[tracing::instrument(name = "activate_user_account", skip(state, acc_user))]
pub async fn activate_user_account(
    State(state): State<AppState>,
    CustomAppJson(acc_user): CustomAppJson<ActivateUser>,
) -> Result<impl IntoResponse, CustomAppError> {
    let mut redis_con = state.redis_store.get().await.map_err(|_| {
        CustomAppError::from((
            "Failed to get redis connection".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?;

    let mut hasher = Sha256::new();
    hasher.update(acc_user.token.as_bytes());
    let hashed_token = format!("{:x}", hasher.finalize());

    let hashed_activation_code: String = bb8_redis::redis::cmd("GET")
        .arg(&acc_user.id.to_string())
        .query_async(&mut *redis_con)
        .await
        .map_err(|_| {
            CustomAppError::from((
                "This activation has been used or expired".to_string(),
                ErrorContext::BadRequest,
            ))
        })?;

    if hashed_activation_code == hashed_token {
        state.db_store.activate_user(&acc_user.id).await?;

        // Delete activation code from redis
        bb8_redis::redis::cmd("DEL")
            .arg(&acc_user.id.to_string())
            .query_async::<_, i64>(&mut *redis_con)
            .await
            .map_err(|_| {
                CustomAppError::from((
                    "Failed to delete activation code from Redis".to_string(),
                    ErrorContext::InternalServerError,
                ))
            })?;

        Ok(SuccessResponse {
            message: "The activation process was successful.".to_string(),
            status_code: StatusCode::OK.as_u16(),
        }
        .into_response())
    } else {
        return Err(CustomAppError::from((
            "Activation code not found or expired".to_string(),
            ErrorContext::BadRequest,
        )));
    }
}

It's simple. We require that the user provides the token sent. We then find its hash and compare it to what we have in redis for that user. If they are the same, the user gets activated. Otherwise, we send an error. Though not covered in this series, a nice feature to have is allowing users to regenerate tokens so that they won't be locked out of our system forever (we want users!!!).

NOTE: I also implemented a handler that retrieves the currently logged-in user. It's simple and can be seen here.

The entire user_routes should now look like this:

// backend/src/routes/users/mod.rs

use crate::utils::validate_authentication_session;
use axum::{
    routing::{get, post},
    Router,
};

mod activate_account;
mod current_user;
mod login;
mod logout;
mod register;

pub fn users_routes(state: crate::startup::AppState) -> Router<crate::startup::AppState> {
    Router::new()
        .route("/logout", post(logout::logout_user))
        .route("/current", get(current_user::get_current_user))
        .route_layer(axum::middleware::from_fn_with_state(
            state.clone(),
            validate_authentication_session,
        ))
        .route("/login", post(login::login_user))
        .route("/register", post(register::register_user))
        .route("/activate", post(activate_account::activate_user_account))
}

With that, we are done with user management stuff. It's not feature-complete though (I gave suggestions). Let's move on to the Q&A service in the next few articles. See ya!!!

Outro

Enjoyed this article? Consider contacting me for a job, something worthwhile or buying a coffee ☕. You can also connect with/follow me on LinkedIn and Twitter. It isn't bad if you help share this article for wider coverage. I would appreciate it...