use crate::sql::Pool;
use crate::totp::TotpSecret;
use crate::Model;
#[derive(Model, Clone)]
#[rustango(table = "rustango_admin_totp", managed = false)]
#[allow(dead_code)]
pub struct AdminTotp {
#[rustango(primary_key)]
pub user_id: i64,
#[rustango(max_length = 64)]
pub secret_base32: String,
#[rustango(default = "false")]
pub confirmed: bool,
#[rustango(default = "now()")]
pub created_at: chrono::DateTime<chrono::Utc>,
pub last_used_step: Option<i64>,
#[rustango(max_length = 64)]
pub pending_secret_base32: Option<String>,
}
impl std::fmt::Debug for AdminTotp {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("AdminTotp")
.field("user_id", &self.user_id)
.field("secret_base32", &"<redacted>")
.field("confirmed", &self.confirmed)
.field("created_at", &self.created_at)
.field("last_used_step", &self.last_used_step)
.field(
"pending_secret_base32",
&self.pending_secret_base32.as_ref().map(|_| "<redacted>"),
)
.finish()
}
}
impl AdminTotp {
pub fn pending_secret(&self) -> Option<TotpSecret> {
match &self.pending_secret_base32 {
Some(s) => TotpSecret::from_base32(s),
None if !self.confirmed => TotpSecret::from_base32(&self.secret_base32),
None => None,
}
}
}
const STEP_SECS: u64 = 30;
const DIGITS: u32 = 6;
const WINDOW: i64 = 1;
pub async fn ensure_table(pool: &Pool) -> Result<(), sqlx::Error> {
use crate::core::Model as _;
let snapshot = crate::migrate::SchemaSnapshot::from_models_forced(&[AdminTotp::SCHEMA]);
crate::migrate::apply_idempotent(pool, &snapshot).await?;
add_new_columns(pool).await
}
async fn add_new_columns(pool: &Pool) -> Result<(), sqlx::Error> {
use crate::core::Model as _;
let snapshot = crate::migrate::SchemaSnapshot::from_models_forced(&[AdminTotp::SCHEMA]);
crate::migrate::add_columns_idempotent(
pool,
&snapshot,
&[
(AdminTotp::SCHEMA.table, "last_used_step"),
(AdminTotp::SCHEMA.table, "pending_secret_base32"),
],
)
.await
}
pub async fn redeem_code(
pool: &Pool,
user_id: i64,
secret: &TotpSecret,
code: &str,
) -> Result<bool, crate::sql::ExecError> {
redeem(pool, user_id, secret, code, false).await
}
pub async fn confirm_with_code(
pool: &Pool,
user_id: i64,
secret: &TotpSecret,
code: &str,
) -> Result<bool, crate::sql::ExecError> {
redeem(pool, user_id, secret, code, true).await
}
async fn redeem(
pool: &Pool,
user_id: i64,
secret: &TotpSecret,
code: &str,
confirm: bool,
) -> Result<bool, crate::sql::ExecError> {
let Some(step) = crate::totp::matched_step(secret, code, STEP_SECS, DIGITS, WINDOW) else {
return Ok(false);
};
let step = i64::try_from(step).unwrap_or(i64::MAX);
let updated = match write_step(pool, user_id, secret, step, confirm).await {
Ok(n) => n,
Err(first) => match add_new_columns(pool).await {
Ok(()) => write_step(pool, user_id, secret, step, confirm).await?,
Err(_) => return Err(first),
},
};
Ok(updated == 1)
}
async fn write_step(
pool: &Pool,
user_id: i64,
secret: &TotpSecret,
step: i64,
confirm: bool,
) -> Result<u64, crate::sql::ExecError> {
use crate::query::Q;
use crate::sql::UpdaterPool as _;
if confirm {
let promoted = AdminTotp::objects()
.filter("user_id", user_id)
.filter("pending_secret_base32", secret.to_base32())
.update()
.set("secret_base32", secret.to_base32())
.set("pending_secret_base32", None::<String>)
.set("confirmed", true)
.set("last_used_step", step)
.execute_pool(pool)
.await?;
if promoted == 1 {
return Ok(promoted);
}
}
let q = AdminTotp::objects()
.filter("user_id", user_id)
.filter("secret_base32", secret.to_base32());
let q = if confirm {
q.filter("confirmed", false)
} else {
q
};
let q = q
.where_(Q::is_null("last_used_step") | Q::lt("last_used_step", step))
.update()
.set("last_used_step", step);
let q = if confirm {
q.set("confirmed", true)
} else {
q.set("pending_secret_base32", None::<String>)
};
q.execute_pool(pool).await
}
pub async fn device(pool: &Pool, user_id: i64) -> Option<AdminTotp> {
use crate::sql::FetcherPool as _;
AdminTotp::objects()
.filter("user_id", user_id)
.fetch(pool)
.await
.ok()
.and_then(|rows| rows.into_iter().next())
}
#[deprecated(
since = "0.59.17",
note = "a read error looks like no device; use confirmed_secret_checked"
)]
pub async fn confirmed_secret(pool: &Pool, user_id: i64) -> Option<TotpSecret> {
confirmed_secret_checked(pool, user_id).await.ok().flatten()
}
pub async fn confirmed_secret_checked(
pool: &Pool,
user_id: i64,
) -> Result<Option<TotpSecret>, crate::sql::ExecError> {
let secrets: Vec<String> = AdminTotp::objects()
.filter("user_id", user_id)
.filter("confirmed", true)
.values_list_flat("secret_base32")
.fetch(pool)
.await?;
secrets
.first()
.map(|s| {
TotpSecret::from_base32(s).ok_or_else(|| {
crate::sql::ExecError::Driver(crate::sql::sqlx::Error::Decode(
"rustango_admin_totp.secret_base32 is not valid base32".into(),
))
})
})
.transpose()
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum Reenroll {
Verified,
Refused,
NoDevice,
}
pub async fn start_reenrollment(
pool: &Pool,
user_id: i64,
current_code: &str,
secret: &TotpSecret,
) -> Result<Reenroll, crate::sql::ExecError> {
let outcome = match confirmed_secret_checked(pool, user_id).await? {
Some(current) if !redeem_code(pool, user_id, ¤t, current_code).await? => {
return Ok(Reenroll::Refused);
}
Some(_) => Reenroll::Verified,
None => Reenroll::NoDevice,
};
start_enrollment(pool, user_id, secret).await?;
Ok(outcome)
}
pub async fn start_enrollment(
pool: &Pool,
user_id: i64,
secret: &TotpSecret,
) -> Result<(), crate::sql::ExecError> {
use crate::sql::UpdaterPool as _;
let staged = AdminTotp::objects()
.filter("user_id", user_id)
.filter("confirmed", true)
.update()
.set("pending_secret_base32", secret.to_base32())
.execute_pool(pool)
.await?;
if staged == 1 {
return Ok(());
}
let del = AdminTotp::objects()
.filter("user_id", user_id)
.filter("confirmed", false)
.compile_delete()?;
crate::sql::delete_pool(pool, &del).await?;
let row = AdminTotp {
user_id,
secret_base32: secret.to_base32(),
confirmed: false,
created_at: chrono::Utc::now(),
last_used_step: None,
pending_secret_base32: None,
};
row.insert_pool(pool).await?;
Ok(())
}
pub async fn confirm(pool: &Pool, user_id: i64) -> Result<(), crate::sql::ExecError> {
use crate::sql::UpdaterPool as _;
let q = AdminTotp::objects().filter("user_id", user_id);
let pending = device(pool, user_id)
.await
.and_then(|d| d.pending_secret_base32);
let q = match pending {
Some(p) => q
.filter("pending_secret_base32", p.clone())
.update()
.set("secret_base32", p)
.set("pending_secret_base32", None::<String>),
None => q.update(),
};
q.set("confirmed", true).execute_pool(pool).await?;
Ok(())
}