CryptoFlow: Building a secure and scalable system with Axum and SvelteKit - Part 1
Published: 2024-01-05 00:02:06

## Introduction

In [part 0][0], we laid some solid background in building our proposed system. The system's structure, database schema, and other details were laid bare. This article builds on that foundation.

> **NOTE: The program isn't feature-complete yet! Contributions are welcome.**

## Source code

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

{% github Sirneij/cryptoflow %}

## Implementation

### Step 1: Cookies and user management

As stated in [part 0][0], our system's authentication is cookie(session)-based. To allow this, we need to set it up. Right from the onset, we installed [axum-extra][1] (with `cookie-private` and `cookie` features enabled) and [tower-http][2] (`cors` feature enabled). These will allow us to achieve our aims. Let's kick start by configuring our app to allow cookies and enable an origin (our frontend app) to directly access the application. To do this, we'll add a `key` attribute to the app's state and make some configurations:

```rust
// backend/src/startup.rs
...
use axum::extract::FromRef;
...

#[derive(Clone)]
pub struct AppState {
    ...
    key: axum_extra::extract::cookie::Key,
}

impl FromRef<AppState> for axum_extra::extract::cookie::Key {
    fn from_ref(state: &AppState) -> Self {
        state.key.clone()
    }
}
...
async fn run(
    listener: tokio::net::TcpListener,
    store: crate::store::Store,
    settings: crate::settings::Settings,
) {
    let cors = tower_http::cors::CorsLayer::new()
        .allow_credentials(true)
        .allow_methods(vec![
            axum::http::Method::OPTIONS,
            axum::http::Method::GET,
            axum::http::Method::POST,
            axum::http::Method::PUT,
            axum::http::Method::DELETE,
        ])
        .allow_headers(vec![
            axum::http::header::ORIGIN,
            axum::http::header::AUTHORIZATION,
            axum::http::header::ACCEPT,
        ])
        .allow_origin(
            settings
                .frontend_url
                .parse::<axum::http::HeaderValue>()
                .unwrap(),
        );

    let app_state = AppState {
        ...
        key: axum_extra::extract::cookie::Key::from(
            std::env::var("COOKIE_SECRET")
                .expect("Failed to get COOKIE_SECRET.")
                .as_bytes(),
        ),
    };
    // build our application with a route
    let app = axum::Router::new()
        ...
        .layer(cors);
    ...
}
...
```

We allowed the basic HTTP methods, certain headers (particularly the authorization), and our frontend application via the [tower-http][2] `CorsLayer`. We want our cookies to be very private so we'll be using [axum_extra::extract::cookie::PrivateCookieJar][3] which requires a `key` for data encryption or to sign the cookies. The key can be generated via the `axum_extra::extract::cookie::Key::generate()` but I chose to generate it from a `512-bit` (64 bytes) cryptographically random string which I saved in `.env`. An example of this is:

```text
// .env
COOKIE_SECRET=3bbefd8d24c89aefd3ad0b8b95afd2ea996e47b89d93d4090b481a091b4e73e5543305f2e831d0b47737d9807a1b5b5773dba3bbb63623bd42de84389fbfa3d1
```

To inform `SignedCookieJar` how to access the key from our `AppState`, we made this `impl`:

```rust
...
impl FromRef<AppState> for axum_extra::extract::cookie::Key {
    fn from_ref(state: &AppState) -> Self {
        state.key.clone()
    }
}
...
```

Now to user data management. For modularity, we'll have a `users` submodule in the `routes` module. The submodule will house all user-related handlers and will expose the `Router` instance specific to user management alone. To achieve this, create two files, `src/routes/users/mod.rs` and `src/routes/users/login.rs`. The former is a special filename used for module organization (I assume you already know this) while the latter houses the handler for user login. Before we write their codes, we need to first write some utilities that will make the codes compile and functional. In the `src/utils` folder, create a `password.rs` file and make it look like this:

```rust
use argon2::{
    password_hash::{rand_core::OsRng, PasswordHash, PasswordHasher, PasswordVerifier, SaltString},
    Argon2,
};

#[tracing::instrument(name = "Hashing user password", skip(password))]
pub async fn hash_password(password: &[u8]) -> String {
    let salt = SaltString::generate(&mut OsRng);
    Argon2::default()
        .hash_password(password, &salt)
        .expect("Unable to hash password.")
        .to_string()
}

#[tracing::instrument(name = "Verifying user password", skip(password, hash))]
pub fn verify_password(hash: &str, password: &[u8]) -> Result<(), argon2::password_hash::Error> {
    let parsed_hash = PasswordHash::new(hash)?;
    Argon2::default().verify_password(password, &parsed_hash)
}
```

There are two functions there:

- `hash_password`: This hashes a plain text using the [argon2][4] hashing algorithm. We used the default settings of [argon2 crate][5] which uses `v19` of `Argon2id` &mdash; memory cost is `19 * 1024` (19 MiB), number of iterations is `2`, and the degree of parallelism is `1`. This is as recommended by [OWASP Password Storage Cheat Sheet][6]. The stringified result of this is what will be saved in the `password` attribute of the `users` relation.
- `verify_password`: A returning user will normally provide the plaintext password used for registration alongside the user's email. This password will need to be "compared" with the saved password to ensure that it's correct. To ascertain this correctness, this function was written.

Next in the list of utilities is error handling. This part was heavily influenced by this [example][7] provided by the axum team.

> **NOTE: The code is just part of the implementation. Check out [this file][8] for the complete code.**

```rust
// src/utils/errors.rs
use crate::utils::CustomAppJson;
use argon2::password_hash::Error as ArgonError;
use axum::{
    extract::rejection::JsonRejection,
    http::StatusCode,
    response::{IntoResponse, Response},
};

use serde::Serialize;

pub enum ErrorContext {
    UnauthorizedAccess,
    InternalServerError,
    BadRequest,
    NotFound,
}

pub enum CustomAppError {
    JsonRejection(JsonRejection),
    DatabaseQueryError(sqlx::Error),
    PasswordHashError(ArgonError),
    RedisError(bb8_redis::redis::RedisError),
    UUIDError(uuid::Error),
    Unauthorized(String),
    InternalError(String),
    BadRequest(String),
    NotFound(String),
    ReqwestError(reqwest::Error),
}

impl IntoResponse for CustomAppError {
    fn into_response(self) -> Response {
        // How we want error responses to be serialized

        #[derive(Serialize)]
        struct ErrorResponse {
            message: String,
            status_code: u16,
        }

        let (status, message) = match self {
            CustomAppError::JsonRejection(rejection) => {
                // This error is caused by bad user input so don't log it
                tracing::error!("Bad user input: {:?}", rejection);
                (rejection.status(), rejection.body_text())
            }
            CustomAppError::DatabaseQueryError(error) => {
                match &error {
                    sqlx::Error::RowNotFound => {
                        tracing::error!("Resource not found: {}", error);
                        (
                            StatusCode::NOT_FOUND,
                            "Resource not found or you are not allowed to perform this operation"
                                .to_string(),
                        )
                    }
                    ...
                }
            }
            CustomAppError::PasswordHashError(error) => match error {
                ArgonError::Password => {
                    tracing::info!("Password mismatch error");
                    (
                        StatusCode::BAD_REQUEST,
                        "Email and Password combination does not match.".to_string(),
                    )
                }
                ...
            },
            ...
        };

        (
            status,
            CustomAppJson(ErrorResponse {
                message,
                status_code: status.as_u16(),
            }),
        )
            .into_response()
    }
}

impl From<JsonRejection> for CustomAppError {
    fn from(rejection: JsonRejection) -> Self {
        Self::JsonRejection(rejection)
    }
}

...

impl From<(String, ErrorContext)> for CustomAppError {
    fn from((message, context): (String, ErrorContext)) -> Self {
        match context {
            ErrorContext::UnauthorizedAccess => CustomAppError::Unauthorized(message),
            ErrorContext::InternalServerError => CustomAppError::InternalError(message),
            ErrorContext::BadRequest => CustomAppError::BadRequest(message),
            ErrorContext::NotFound => CustomAppError::NotFound(message),
        }
    }
}
```

It's some straightforward code that allows many of the expected errors (from `SQLx`, `bb8_redis`, `argon2`, `uuid` and others) to be gracefully handled.

For data extraction from requests' bodies, I also made a simple `JSON` extractor (for now) in `src/utils/responses.rs`:

```rust
use axum::{
    extract::FromRequest,
    http::StatusCode,
    response::{IntoResponse, Response},
};

use crate::utils::CustomAppError;

use serde::Serialize;

#[derive(FromRequest)]
#[from_request(via(axum::Json), rejection(CustomAppError))]
pub struct CustomAppJson<T>(pub T);

impl<T> IntoResponse for CustomAppJson<T>
where
    axum::Json<T>: IntoResponse,
{
    fn into_response(self) -> Response {
        axum::Json(self.0).into_response()
    }
}
...
```

It does what `axum::Json` would do with an extra.

The next utility functions will be creating and retrieving (for now) users from the database. To write them, we will remember the `Store` struct we implemented in [part 0][0]. We will extend it by adding methods that will facilitate those operations:

```rust
// src/store/users.rs
use sqlx::Row;

impl crate::store::Store {
    #[tracing::instrument(name = "get_user_by_id", fields(user_id = id.to_string()))]
    pub async fn get_user_by_id(&self, id: uuid::Uuid) -> Result<crate::models::User, sqlx::Error> {
        sqlx::query_as::<_, crate::models::User>(
            r#"
        SELECT
            id,
            email,
            password,
            first_name,
            last_name,
            is_active,
            is_staff,
            is_superuser,
            thumbnail,
            date_joined
        FROM users
        WHERE id = $1 AND is_active = true
        "#,
        )
        .bind(id)
        .fetch_one(&self.connection)
        .await
    }

    #[tracing::instrument(name = "get_user_by_email", fields(user_email = email))]
    pub async fn get_user_by_email(&self, email: &str) -> Result<crate::models::User, sqlx::Error> {
        sqlx::query_as::<_, crate::models::User>(
            r#"
        SELECT
            id,
            email,
            password,
            first_name,
            last_name,
            is_active,
            is_staff,
            is_superuser,
            thumbnail,
            date_joined
        FROM users
        WHERE email = $1 AND is_active = true
        "#,
        )
        .bind(email)
        .fetch_one(&self.connection)
        .await
    }

    #[tracing::instrument(name = "create_user", skip(password), fields(user_first_name = first_name, user_last_name = last_name, user_email = email))]
    pub async fn create_user(
        &self,
        first_name: &str,
        last_name: &str,
        email: &str,
        password: &str,
    ) -> Result<crate::models::UserVisible, sqlx::Error> {
        sqlx::query_as::<_, crate::models::UserVisible>(
        r#"
        INSERT INTO users (first_name, last_name, email, password)
            VALUES ($1, $2, $3, $4)
        RETURNING
            id, email, first_name, last_name, is_active, is_staff, is_superuser, thumbnail, date_joined
        "#
    )
    .bind(first_name)
    .bind(last_name)
    .bind(email)
    .bind(password)
    .fetch_one(&self.connection)
    .await
    }

    #[tracing::instrument(name = "activate_user", fields(user_id = id.to_string()))]
    pub async fn activate_user(&self, id: &uuid::Uuid) -> Result<(), sqlx::Error> {
        sqlx::query(
            r#"
        UPDATE users
        SET is_active = true
        WHERE id = $1
        "#,
        )
        .bind(id)
        .execute(&self.connection)
        .await?;

        Ok(())
    }

    #[tracing::instrument(name="create_super_user_in_db.", skip(settings), fields(user_email = settings.superuser.email, user_first_name = settings.superuser.first_name, user_last_name = settings.superuser.last_name))]
    pub async fn create_super_user_in_db(&self, settings: &crate::settings::Settings) {
        let new_super_user = crate::models::NewUser {
            email: settings.superuser.email.clone(),
            password: crate::utils::hash_password(&settings.superuser.password.as_bytes()).await,
            first_name: settings.superuser.first_name.clone(),
            last_name: settings.superuser.last_name.clone(),
        };

        match sqlx::query(
            "INSERT INTO users
                    (email, password, first_name, last_name, is_active, is_staff, is_superuser)
                VALUES ($1, $2, $3, $4, true, true, true)
                ON CONFLICT (email)
                DO UPDATE
                    SET
                        first_name=EXCLUDED.first_name,
                        last_name=EXCLUDED.last_name
                RETURNING id",
        )
        .bind(new_super_user.email)
        .bind(&new_super_user.password)
        .bind(new_super_user.first_name)
        .bind(new_super_user.last_name)
        .map(|row: sqlx::postgres::PgRow| -> uuid::Uuid { row.get("id") })
        .fetch_one(&self.connection)
        .await
        {
            Ok(id) => {
                tracing::info!("Super user created successfully {:#?}.", id);
                id
            }

            Err(e) => {
                tracing::error!("Failed to insert user into DB: {:#?}.", e);
                uuid::Uuid::new_v4()
            }
        };
    }
}
```

We `impl` the `Store` struct so we can use its `connection` attribute to talk to the database directly via the `SQLx` crate. The methods peculiar to user management are the `get_user_by_id`, `get_user_by_email`, `activate_user`, and `create_super_user_in_db` for now. They all mean what their names imply. The last one is an administrative method that creates a user with "superpowers". It will be used in the `build` method later on. All of these methods referenced different data models (typically `struct`s) that we've not defined yet. Let's define them in `src/models/users.rs`:

```rust
// src/models/users.rs
#[derive(serde::Serialize, Debug, sqlx::FromRow)]
pub struct User {
    pub id: uuid::Uuid,
    pub email: String,
    pub password: String,
    pub first_name: String,
    pub last_name: String,
    pub is_active: Option<bool>,
    pub is_staff: Option<bool>,
    pub is_superuser: Option<bool>,
    pub thumbnail: Option<String>,
    pub date_joined: time::OffsetDateTime,
}

#[derive(serde::Serialize, Debug, sqlx::FromRow)]
pub struct UserVisible {
    pub id: uuid::Uuid,
    pub email: String,
    pub first_name: String,
    pub last_name: String,
    pub is_active: Option<bool>,
    pub is_staff: Option<bool>,
    pub is_superuser: Option<bool>,
    pub thumbnail: Option<String>,
    pub date_joined: time::OffsetDateTime,
}

#[derive(serde::Serialize)]
pub struct LoggedInUser {
    pub id: uuid::Uuid,
    pub email: String,
    pub password: String,
    pub is_staff: bool,
    pub is_superuser: bool,
}
#[derive(serde::Deserialize, Debug)]
pub struct NewUser {
    pub email: String,
    pub password: String,
    pub first_name: String,
    pub last_name: String,
}

#[derive(serde::Deserialize, Debug)]
pub struct LoginUser {
    pub email: String,
    pub password: String,
}

#[derive(serde::Deserialize, Debug)]
pub struct ActivateUser {
    pub id: uuid::Uuid,
    pub token: String,
}
```

They are just `struct`s that derived from (or implemented) `serde`'s `Serialize` and/or `Deserialize`. Those that derived `Deserialize` are used for incoming requests while those that derived `Serialize` are going to be used as requests' responses. Most of them derived `Debug` too so that they can be logged. Two implemented `sqlx::FromRow`. This is to allow passing them in `sqlx::query_as`. It is a requirement. For this to work, ensure that the data returned by the SQL statement have the same name as the attributes in the struct.

Having paved the way, let's write the login handler:

```rust
// src/routes/users/login.rs
use crate::models::LoginUser;
use crate::startup::AppState;
use crate::utils::verify_password;
use crate::utils::SuccessResponse;
use crate::utils::{CustomAppError, CustomAppJson, ErrorContext};
use axum::{extract::State, http::StatusCode, response::IntoResponse};
use axum_extra::extract::cookie::{Cookie, PrivateCookieJar, SameSite};
use time::Duration;

#[axum::debug_handler]
#[tracing::instrument(name = "login_user", skip(cookies, state, login))]
pub async fn login_user(
    cookies: PrivateCookieJar,
    State(state): State<AppState>,
    CustomAppJson(login): CustomAppJson<LoginUser>,
) -> Result<(PrivateCookieJar, impl IntoResponse), CustomAppError> {
    // Get user from db by email
    let user = state
        .db_store
        .get_user_by_email(&login.email)
        .await
        .map_err(|_| {
            CustomAppError::from((
                "Invalid email or password".to_string(),
                ErrorContext::BadRequest,
            ))
        })?;

    // Verify password
    tokio::task::spawn_blocking(move || {
        verify_password(&user.password, &login.password.as_bytes())
    })
    .await
    .map_err(|_| {
        CustomAppError::from((
            "Server error occurred".to_string(),
            ErrorContext::InternalServerError,
        ))
    })?
    .map_err(|_| {
        CustomAppError::from((
            "Invalid email or password".to_string(),
            ErrorContext::BadRequest,
        ))
    })?;

    // Generate a truly random session id for the user
    let session_id = uuid::Uuid::new_v4().to_string();

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

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

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

    // Create cookie
    let cookie = Cookie::build(("sessionid", session_id))
        .secure(true)
        .same_site(SameSite::Strict)
        .http_only(true)
        .path("/")
        .max_age(Duration::minutes(cookie_expiration));

    Ok((
        cookies.add(cookie),
        SuccessResponse {
            message: "The authentication process was successful.".to_string(),
            status_code: StatusCode::OK.as_u16(),
        }
        .into_response(),
    ))
}
```

Though long due to error handling, the concept is simple. The handler takes the `PrivateCookieJar` extractor (to extract and help propagate user cookies), `AppState` extractor (to help hold the `AppState` data for the handler to use), and the `CustomAppJson` extractor (to extract the request body). The last argument must be positioned because it consumes the request body. If not, the code will not compile! It returns a `Result` of the tuple (`PrivateCookieJar`, `impl IntoResponse`) or `CustomAppError`. For `PrivateCookieJar` to work, its "value must be returned from the handler as part of the response for the changes to be propagated". In the function body, we first tried to retrieve the requesting user from the database via the email address. Failure to get the user leads to an error being returned. This is where our efforts so far start to shine. Otherwise, we proceeded to verify the user's password. Verification of hashed passwords is CPU-intensive and can block the async runtime. To mitigate this, we spawned a tokio task so that the operation wouldn't be blocked. Next, every user should have a truly random and unique session identification. `uuid` came to the rescue! Since we need to store this session somewhere for subsequent validations as long as the session lives, we chose to store it in redis. Another option is to store it in the PostgreSQL database but this will be slower. A simple key/value store like redis is perfect! The normal Rust's [redis crate][9] doesn't support pooling which is important for a system that will serve a lot of traffic. This made me opt for [bb8-redis][10] which provides `async` and tokio-based redis connection pool. We will set it later but for now, we are already using it to store the `session_id` as the key and `user_id` as the value. After that, we built a cookie which encrypts the generated `session_id`. The cookie is made secure and imposes strict same-site attributes. Of course, we made the cookie HttpOnly to prevent client-side scripts from accessing its embedded data. We used a global path for it and using a configurable expiration period, we set the cookie's `max_age`. As previously stated, it is a requirement to return the cookie as part of the HTTP response.

Next up, let's route this handler:

```rust
// src/routes/users/mod.rs
use axum::{routing::post, Router};

mod login;

pub fn users_routes() -> Router<crate::startup::AppState> {
    Router::new()
        .route("/login", post(login::login_user))
}
```

`users_routes` will build all routes related to user management. Since the route uses the `AppState`, the router returned must specify it. Ensure you make `users_routes` available by exporting it in the main `routes/mod.rs`.

To conclude this long journey, let's set up [bb8-redis][10] connection and include `users_routes` to the main route instance!

```rust
// src/startup.rs
...

#[derive(Clone)]
pub struct AppState {
    ...
    pub redis_store: bb8_redis::bb8::Pool<bb8_redis::RedisConnectionManager>,
}

...

impl Application {
    ...
    pub async fn build(
        settings: crate::settings::Settings,
        test_pool: Option<sqlx::postgres::PgPool>,
    ) -> Result<Self, std::io::Error> {
        ...

        sqlx::migrate!()
            .run(&store.clone().connection)
            .await
            .expect("Failed to migrate");

        // Create superuser if not exists
        store.create_super_user_in_db(&settings).await;
        ...
    }
}

async fn run(
    listener: tokio::net::TcpListener,
    store: crate::store::Store,
    settings: crate::settings::Settings,
) {
    let redis_url = std::env::var("REDIS_URL").expect("Failed to get REDIS_URL.");
    let manager =
        bb8_redis::RedisConnectionManager::new(redis_url).expect("Failed to create redis manager");
    let redis_pool = bb8_redis::bb8::Pool::builder()
        .max_size(15)
        .build(manager)
        .await
        .expect("Failed to create redis pool.");

    ...
    let app_state = AppState {
        ...
        redis_store: redis_pool,
    };
    // build our application with a route
    let app = axum::Router::new()
        ...
        .nest("/api/users", routes::users_routes(app_state.clone()))
        ...;

    ...
}
...
```

Since we want to automatically create the user with "superpowers", we called the utility method for doing that in the `build` method. In the `run` function, we retrieved our machine's redis instance's URL from the `.env` file and created a new bb8 `RedisConnectionManager` from it. From the connection manager, we built a 15-connection pool and added it to the application state. A nice improvement is to make the number of connection pools configurable. Lastly, we added our `users_routes` to the main route using the `nest` method. It helps make our routing composable!

Let's stop here for this part. We'll continue in the next article!

## Outro

Enjoyed this article? I'm a Software Engineer and Technical Writer actively seeking new opportunities, particularly in areas related to web security, finance, health care, and education. If you think my expertise aligns with your team's needs, let's chat! You can find me on LinkedIn: [LinkedIn](https://www.linkedin.com/in/idogun-john-nelson/) and Twitter: [Twitter](https://twitter.com/Sirneij).

If you found this article valuable, consider sharing it with your network to help spread the knowledge!

[0]: https://dev.to/sirneij/cryptoflow-building-a-secure-and-scalable-system-with-axum-and-sveltekit-part-0-mn5 "CryptoFlow: Building a secure and scalable system with Axum and SvelteKit - Part 0"
[1]: https://crates.io/crates/axum-extra "Extra utilities for axum"
[2]: https://crates.io/crates/tower-http "Tower middleware and utilities for HTTP clients and servers"
[3]: https://docs.rs/axum-extra/latest/axum_extra/extract/cookie/struct.PrivateCookieJar.html "Available on crate features cookie-private and cookie only."
[4]: https://en.wikipedia.org/wiki/Argon2 "Argon2"
[5]: https://docs.rs/argon2/latest/argon2/ "Crate argon2"
[6]: https://cheatsheetseries.owasp.org/cheatsheets/Password_Storage_Cheat_Sheet.html "Password Storage Cheat Sheet"
[7]: https://github.com/tokio-rs/axum/tree/main/examples/error-handling "Axum error-handling example"
[8]: https://github.com/Sirneij/cryptoflow/blob/main/backend/src/utils/errors.rs "CryptoFlow error handling"
[9]: https://crates.io/crates/redis "Redis driver for Rust."
[10]: https://crates.io/crates/bb8-redis "Full-featured async (tokio-based) redis connection pool (like r2d2)"