Rust Story: Simple Auth Token generator in Rust using tokio

Shobhit chaturvedi
6 min readDec 3, 2023

In the digital realm, where security is paramount and time is of the essence, managing access tokens effectively becomes an art as much as a science. Picture a world where each digital handshake is governed by a time-bound passcode, a token that holds the key to secure and authorised communication. This is a tale of such tokens, their intricate dance with time, and the Rusty gears that keep this dance in perfect harmony.

Our story unfolds in the server-client landscape, a place where trust is paramount and every interaction is scrutinised. Here, a server must send a response to the client with a special charm — an authentication token. But this token is no ordinary key; it’s akin to Cinderella’s carriage, magical yet time-bound. As the clock ticks and reaches the 20-minute mark, the token, like the carriage, transforms, ensuring that no expired token ever overstays its welcome.

In this dance of tokens, we turn to Rust, the language known for its robustness and reliability, to choreograph our steps. Our tools? Tokio for creating tasks that keep an eye on the ticking clock, and Axum to construct a server, a stage for our performance.

Imagine a new cargo project, token_manager, our backstage, where the magic of code turns into a performance. Here, dependencies like Axum and Tokio are the stagehands, setting the scene for our play.

Our script is written in routes.rs, where AuthToken and TokenManager take the lead roles. AuthToken is a meticulous character, always aware of the ticking clock, while TokenManager is the mastermind, ensuring the tokens are fresh and the performance seamless.

As our story progresses, TokenManager comes to life, juggling AuthTokens with precision. It's a spectacle of asynchronous tasks, where tokens are generated, monitored, and refreshed – all in a rhythmic dance orchestrated by Tokio's ticks.

Our stage is set with lazy_static!, ensuring our TokenManager is always ready, always alert. The default expiration time and tick count are like the tempo of our music, guiding the pace of our performance.

In the grand finale, our main.rs, the curtain rises on an Axum server. Routes are exposed like pathways on our stage, each leading to a different act of our token tale. As the audience, you can witness the token generation, a marvel in itself, refresh every 5 seconds, a testament to the precision of our code.

As you navigate to http://localhost:3000, you're not just accessing a route; you're stepping into a world where tokens dance to the beats of Rust, where security is a performance, and where every token's expiry is a cue for the next to take the stage.

In this world, the another_route is your window to the ongoing show, a peek behind the curtains of our token theater. Here, you see the active token, a star in its own right, performing its role flawlessly until it's time to pass the baton.

And so, our story of token management in Rust concludes, but the performance never truly ends. It’s a cycle, a continuous loop of security and efficiency, a ballet of bytes and tokens, all playing their part in the vast digital theater. Welcome to the world of token management in Rust — secure, efficient, and always on time.

lets create a new cargo project token_manager and update the toml as follows —

[package]
name = "token_manager"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html

[dependencies]
axum = "0.6.20"
hyper = { version = "0.14.27", features = ["full"] }
chrono = {version = "0.4.31" }

tokio = { version = "1.34.0", features = ["full"] }
lazy_static = "1.4.0"

create a crate routes.rs which willhave actual token struct and token_manager struct to expose token monitoring APIs as follows —

//routes.rs
use chrono::{DateTime, Utc};
use std::sync::{Arc, RwLock};
use tokio::time::{interval, Duration};

lazy_static! {
pub static ref GLOBAL_TOKEN_MANAGER: TokenManager = TokenManager::new(5, 1);
}

struct AuthToken {
token: String,
exp_time: DateTime<Utc>,
expire_seconds: i64,
}

impl AuthToken {
pub fn new(expire_seconds: i64) -> Self {
AuthToken {
token: String::new(),
exp_time: Utc::now() + chrono::Duration::seconds(expire_seconds),
expire_seconds,
}
}

pub async fn update_token(&mut self) {
println!("Generating new token!");
self.token = AuthToken::generate_secure_token().await;
self.exp_time = Utc::now() + chrono::Duration::seconds(self.expire_seconds);
}

pub fn refresh_token(&mut self, token: &str) {
println!("Generating new token!");
self.token = token.to_owned();
self.exp_time = Utc::now() + chrono::Duration::seconds(self.expire_seconds);
}

pub fn is_expired(&self) -> bool {
Utc::now() > self.exp_time
}

pub fn get_token(&self) -> String {
self.token.clone()
}

pub async fn generate_secure_token() -> String {
let mut interval = interval(Duration::from_secs(1));
interval.tick().await;
Utc::now().to_string()
}
}

pub struct TokenManager {
auth_token: Arc<RwLock<AuthToken>>,
tick_interval_secs: u64,
}

impl TokenManager {
pub fn new(expire_seconds: i64, tick_interval_secs: u64) -> Self {
Self {
auth_token: Arc::new(RwLock::new(AuthToken::new(expire_seconds))),
tick_interval_secs,
}
}

pub fn fetch_token(&self) -> Option<String> {
let token = self.auth_token.clone();
let token_string = match token.read() {
Ok(t) => Some(t.get_token()),
Err(e) => {
eprintln!("Failed to acquire lock: {}", e);
None
}
};
token_string
}

pub async fn refresh_token(&self) {
let token_ref = self.auth_token.clone();
let new_token = AuthToken::generate_secure_token().await;
match token_ref.write() {
Ok(mut t) => t.refresh_token(&new_token),
Err(e) => {
eprintln!("Failed to acquire lock: {}", e);
}
};
}

pub async fn generate_tokens(&self) -> tokio::task::JoinHandle<()> {
let token_ref = self.auth_token.clone();
let interval_secs = self.tick_interval_secs;
tokio::spawn(async move {
let mut interval = interval(Duration::from_secs(interval_secs));
loop {
interval.tick().await;

// Check if the token needs to be refreshed
let needs_refresh = {
let token = match token_ref.read() {
Ok(t) => t,
Err(e) => {
eprintln!("Failed to acquire lock: {}", e);
continue;
}
};
token.is_expired()
};

// Refresh the token if necessary
if needs_refresh {
let new_token = AuthToken::generate_secure_token().await;
let mut token = match token_ref.write() {
Ok(t) => t,
Err(e) => {
eprintln!("Failed to acquire lock: {}", e);
continue;
}
};
token.refresh_token(&new_token);
}
}
})
}
}

struct AuthToken has following attributes

struct AuthToken {
token: String, // this will contain actual token
exp_time: DateTime<Utc>, // this will maintain expire time of current active token
expire_seconds: i64,// this will keep information of expiry time of a token .
}

now create a TokenManager , which will contain AuthToken RwLock and tick count to wait before checking token expiration time for an active token

pub struct TokenManager {
auth_token: Arc<RwLock<AuthToken>>, // actual token struct object
tick_interval_secs: u64,// wait period to check expire time of an active token
}

TokenManager will have helper public method to generate and fetch token whenever needed, it is important to understand how generate_token method works based on our requirement,

generate_token spawn a tokio async task and make a read call to auth_token object to check if token is expired or not. if token is expired then it will return true to need_refresh flag and once need_refresh flag set to true ,write lock on auth_token acquired and new token will be generated before calling refresh_token method.

pub async fn generate_tokens(&self) -> tokio::task::JoinHandle<()> {
let token_ref = self.auth_token.clone();
let interval_secs = self.tick_interval_secs;
tokio::spawn(async move {
let mut interval = interval(Duration::from_secs(interval_secs));
loop {
interval.tick().await;

// Check if the token needs to be refreshed
let needs_refresh = {
let token = match token_ref.read() {
Ok(t) => t,
Err(e) => {
eprintln!("Failed to acquire lock: {}", e);
continue;
}
};
token.is_expired()
};

// Refresh the token if necessary
if needs_refresh {
let new_token = AuthToken::generate_secure_token().await;
let mut token = match token_ref.write() {
Ok(t) => t,
Err(e) => {
eprintln!("Failed to acquire lock: {}", e);
continue;
}
};
token.refresh_token(&new_token);
}
}
})
}
}

lets create static variable to initialise and access the TokenManger globally

lazy_static! {
pub static ref GLOBAL_TOKEN_MANAGER: TokenManager = TokenManager::new(5, 1);
}

we can set default expiration_time and tick_count.

Now, lets create a main function to demo the requirement —

//main.rs

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

// use axum_server::Server;
#[macro_use]
extern crate lazy_static;

mod routes;
use crate::routes::route;


async fn token_logic() -> tokio::task::JoinHandle<()> {
route::GLOBAL_TOKEN_MANAGER.refresh_token().await;
let handle = route::GLOBAL_TOKEN_MANAGER.generate_tokens().await;

handle
}

pub async fn another_route() -> String {
if let Some(token) = route::GLOBAL_TOKEN_MANAGER.fetch_token() {
String::from(token)
} else {
String::new()
}
// "This is another route."
}

#[tokio::main]
async fn main() {
// build our application with a single route
let app = Router::new()
.route("/", get(|| async { "Hello, World!" }))
.route("/another", axum::routing::get(another_route));

let handle = token_logic().await;

// run it with hyper on localhost:3000
axum::Server::bind(&"0.0.0.0:3000".parse().unwrap())
.serve(app.into_make_service())
.await
.unwrap();
// handle.abort();
// println!("token aborted");
}

in main , we are creating an app which will generate token and start axum server with exposed routes, if you run http://localhost:3000 in your browser you will see token generation is started and it is refreshing token in every 5 sec. after checking expiration time of generated token continuously in every 1 sec.

you see active token value by accessing route “another_route” http://localhost:3000/another_route.

you can update generate_secure_token logic to create token based on your requirement

Thats how we can create access token application in Rust !

--

--