use std::{fmt::Debug, str::FromStr}; use axum::async_trait; use secrecy::ExposeSecret; use sqlx::{query, query_as, sqlite::SqliteConnectOptions, Connection, SqliteConnection}; use tracing::instrument; use crate::secrets::UserToken; // TODO: check again if it's possible to pass this as a parameter to query! // const TABLE_USER_TOKENS: &str = "user_tokens"; pub enum UserTokenEntry { Valid( UserToken, time::OffsetDateTime, Option, ), Invalid( UserToken, time::OffsetDateTime, Option, ), } #[async_trait] pub trait Database: Debug { async fn get_token(&mut self, username: &str) -> Result, sqlx::Error>; async fn save_token(&mut self, username: &str, token: &UserToken) -> Result<(), sqlx::Error>; async fn update_token_last_used( &mut self, username: &str, token: &UserToken, ) -> Result<(), sqlx::Error>; } #[derive(Debug)] pub struct SqliteDatabase { conn: SqliteConnection, } impl SqliteDatabase { #[instrument] pub async fn open(connection_string: &str) -> Self { let options = SqliteConnectOptions::from_str(connection_string) .expect("Invalid database URI") .create_if_missing(true); let mut db = Self { conn: SqliteConnection::connect_with(&options) .await .expect("Failed to open SQLite database"), }; db.init().await; db } #[instrument] pub async fn init(&mut self) { query!( "CREATE TABLE IF NOT EXISTS user_tokens ( id INTEGER PRIMARY KEY AUTOINCREMENT, username VARCHAR(255) NOT NULL, token NCHAR(30) NOT NULL, valid BOOLEAN NOT NULL, created DATETIME NOT NULL, last_used DATETIME )" ) .execute(&mut self.conn) .await .expect("Failed to initialize table user_tokens"); } } #[async_trait] impl Database for SqliteDatabase { #[instrument] async fn get_token(&mut self, username: &str) -> Result, sqlx::Error> { struct TokenRow { token: String, valid: bool, created: time::OffsetDateTime, last_used: Option, } let row = query_as!( TokenRow, "SELECT token, valid, created, last_used FROM user_tokens WHERE username = $1 ORDER BY created DESC", username ) .fetch_optional(&mut self.conn) .await?; #[allow(clippy::match_bool)] Ok(row.map( |TokenRow { token, valid, created, last_used, }| match valid { true => UserTokenEntry::Valid(UserToken::from(token), created, last_used), false => UserTokenEntry::Invalid(UserToken::from(token), created, last_used), }, )) } #[instrument] async fn save_token(&mut self, username: &str, token: &UserToken) -> Result<(), sqlx::Error> { let token_inner = token.0.expose_secret(); query!( "INSERT INTO user_tokens (username, token, created, valid) VALUES ($1, $2, DATETIME('NOW'), TRUE) ", username, token_inner ) .execute(&mut self.conn) .await?; Ok(()) } #[instrument] async fn update_token_last_used( &mut self, username: &str, token: &UserToken, ) -> Result<(), sqlx::Error> { let token_inner = token.0.expose_secret(); query!( "UPDATE user_tokens SET last_used = DATETIME('NOW') WHERE username = $1 AND token = $2", username, token_inner ) .execute(&mut self.conn) .await?; Ok(()) } }