factoriauth/src/db.rs

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<time::OffsetDateTime>,
    ),
    Invalid(
        UserToken,
        time::OffsetDateTime,
        Option<time::OffsetDateTime>,
    ),
}

#[async_trait]
pub trait Database: Debug {
    async fn get_user_token(
        &mut self,
        username: &str,
    ) -> Result<Option<UserTokenEntry>, 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<Option<ServerPadlock>, 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<Self, sqlx::Error> {
        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<Option<UserTokenEntry>, sqlx::Error> {
        struct TokenRow {
            token: String,
            valid: bool,
            created: time::OffsetDateTime,
            last_used: Option<time::OffsetDateTime>,
        }
        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<Option<ServerPadlock>, 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(())
    }
}