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::{ServerHash, ServerPadlock, 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_user_token( &mut self, username: &str, ) -> Result, sqlx::Error>; async fn save_user_token( &mut self, username: &str, token: &UserToken, ) -> Result<(), sqlx::Error>; async fn update_user_token_last_used( &mut self, username: &str, token: &UserToken, ) -> Result<(), sqlx::Error>; async fn get_server_padlock( &mut self, server_hash: &ServerHash, ) -> Result, sqlx::Error>; async fn save_server_padlock( &mut self, server_hash: &ServerHash, server_padlock: &ServerPadlock, ) -> Result<(), sqlx::Error>; } #[derive(Debug)] pub struct SqliteDatabase { conn: SqliteConnection, } impl SqliteDatabase { #[instrument] pub async fn open(connection_string: &str) -> Result { let options = SqliteConnectOptions::from_str(connection_string)?.create_if_missing(true); let mut db = Self { conn: SqliteConnection::connect_with(&options).await?, }; db.init().await?; Ok(db) } #[instrument] async fn init(&mut self) -> Result<(), sqlx::Error> { 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?; query!( "CREATE TABLE IF NOT EXISTS server_padlocks ( hash TEXT NOT NULL UNIQUE, padlock TEXT NOT NULL ) STRICT" ) .execute(&mut self.conn) .await?; Ok(()) } } #[async_trait] impl Database for SqliteDatabase { #[instrument] async fn get_user_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_user_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_user_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(()) } #[instrument] async fn get_server_padlock( &mut self, server_hash: &ServerHash, ) -> Result, sqlx::Error> { let server_hash = &server_hash.0; let padlock = query!( "SELECT padlock FROM server_padlocks WHERE hash = $1", server_hash ) .fetch_optional(&mut self.conn) .await?; Ok(padlock.map(|d| ServerPadlock::from(d.padlock))) } #[instrument] async fn save_server_padlock( &mut self, server_hash: &ServerHash, server_padlock: &ServerPadlock, ) -> Result<(), sqlx::Error> { let server_hash = &server_hash.0; let server_padlock = server_padlock.0.expose_secret(); query!( "INSERT INTO server_padlocks (hash, padlock) VALUES ($1, $2)", server_hash, server_padlock ) .execute(&mut self.conn) .await?; Ok(()) } }