vaultwarden/src/db/models/two_factor_duo_context.rs

use chrono::Utc;
use diesel::prelude::*;

use crate::{
    api::EmptyResult,
    db::{DbConn, schema::twofactor_duo_ctx},
    error::MapResult,
};

#[derive(Identifiable, Queryable, Insertable, AsChangeset)]
#[diesel(table_name = twofactor_duo_ctx)]
#[diesel(primary_key(state))]
pub struct TwoFactorDuoContext {
    pub state: String,
    pub user_email: String,
    pub nonce: String,
    pub exp: i64,
}

impl TwoFactorDuoContext {
    pub async fn find_by_state(state: &str, conn: &DbConn) -> Option<Self> {
        conn.run(move |conn| {
            twofactor_duo_ctx::table.filter(twofactor_duo_ctx::state.eq(state)).first::<Self>(conn).ok()
        })
        .await
    }

    pub async fn save(state: &str, user_email: &str, nonce: &str, ttl: i64, conn: &DbConn) -> EmptyResult {
        // A saved context should never be changed, only created or deleted.
        let exists = Self::find_by_state(state, conn).await;
        if exists.is_some() {
            return Ok(());
        }

        let exp = Utc::now().timestamp() + ttl;

        conn.run(move |conn| {
            diesel::insert_into(twofactor_duo_ctx::table)
                .values((
                    twofactor_duo_ctx::state.eq(state),
                    twofactor_duo_ctx::user_email.eq(user_email),
                    twofactor_duo_ctx::nonce.eq(nonce),
                    twofactor_duo_ctx::exp.eq(exp),
                ))
                .execute(conn)
                .map_res("Error saving context to twofactor_duo_ctx")
        })
        .await
    }

    pub async fn find_expired(conn: &DbConn) -> Vec<Self> {
        let now = Utc::now().timestamp();
        conn.run(move |conn| {
            twofactor_duo_ctx::table
                .filter(twofactor_duo_ctx::exp.lt(now))
                .load::<Self>(conn)
                .expect("Error finding expired contexts in twofactor_duo_ctx")
        })
        .await
    }

    pub async fn delete(&self, conn: &DbConn) -> EmptyResult {
        conn.run(move |conn| {
            diesel::delete(twofactor_duo_ctx::table.filter(twofactor_duo_ctx::state.eq(&self.state)))
                .execute(conn)
                .map_res("Error deleting from twofactor_duo_ctx")
        })
        .await
    }

    pub async fn purge_expired_duo_contexts(conn: &DbConn) {
        for context in Self::find_expired(conn).await {
            context.delete(conn).await.ok();
        }
    }
}