use std::collections::HashMap;
use std::sync::RwLock;
use std::time::{SystemTime, UNIX_EPOCH};
use async_trait::async_trait;
use rand::RngCore;
use sha2::{Digest, Sha256};
use subtle::ConstantTimeEq;
use crate::error::{IdentityError, Result};
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum OttPurpose {
EmailVerification,
PhoneVerification,
PasswordReset,
MagicLink,
LoginOtp,
}
impl OttPurpose {
pub fn as_str(self) -> &'static str {
match self {
Self::EmailVerification => "email_verification",
Self::PhoneVerification => "phone_verification",
Self::PasswordReset => "password_reset",
Self::MagicLink => "magic_link",
Self::LoginOtp => "login_otp",
}
}
}
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct OttRecord {
pub user_id: String,
pub purpose: OttPurpose,
pub token_hash: [u8; 32],
pub expires_at: u64,
}
#[async_trait]
pub trait OttStore: Send + Sync + 'static {
async fn put(&self, record: OttRecord) -> Result<()>;
async fn take(&self, token_hash: &[u8; 32]) -> Result<Option<OttRecord>>;
}
fn unix_now() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
fn hash_token(token: &str) -> [u8; 32] {
let mut h = Sha256::new();
h.update(token.as_bytes());
h.finalize().into()
}
pub struct OneTimeTokenService {
store: std::sync::Arc<dyn OttStore>,
}
impl OneTimeTokenService {
pub fn new(store: std::sync::Arc<dyn OttStore>) -> Self {
Self { store }
}
pub async fn issue(&self, user_id: &str, purpose: OttPurpose, ttl_secs: u64) -> Result<String> {
let mut bytes = [0u8; 32];
rand::thread_rng().fill_bytes(&mut bytes);
let token = base64_url(&bytes);
let record = OttRecord {
user_id: user_id.to_owned(),
purpose,
token_hash: hash_token(&token),
expires_at: unix_now() + ttl_secs,
};
self.store.put(record).await?;
Ok(token)
}
pub async fn issue_numeric(
&self,
user_id: &str,
purpose: OttPurpose,
digits: u32,
ttl_secs: u64,
) -> Result<String> {
let modulo = 10u64.pow(digits);
let n = (rand::thread_rng().next_u64()) % modulo;
let code = format!("{n:0width$}", width = digits as usize);
let record = OttRecord {
user_id: user_id.to_owned(),
purpose,
token_hash: hash_token(&code),
expires_at: unix_now() + ttl_secs,
};
self.store.put(record).await?;
Ok(code)
}
pub async fn consume(&self, token: &str, purpose: OttPurpose) -> Result<String> {
let presented = hash_token(token);
let record = self
.store
.take(&presented)
.await?
.ok_or(IdentityError::InvalidToken)?;
let ok: bool = record.token_hash.ct_eq(&presented).into();
if !ok || record.purpose != purpose || unix_now() >= record.expires_at {
return Err(IdentityError::InvalidToken);
}
Ok(record.user_id)
}
}
fn base64_url(bytes: &[u8]) -> String {
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use base64::Engine;
URL_SAFE_NO_PAD.encode(bytes)
}
#[derive(Default)]
pub struct MemoryOttStore {
inner: RwLock<HashMap<[u8; 32], OttRecord>>,
}
impl MemoryOttStore {
pub fn new() -> Self {
Self::default()
}
}
#[async_trait]
impl OttStore for MemoryOttStore {
async fn put(&self, record: OttRecord) -> Result<()> {
self.inner
.write()
.unwrap()
.insert(record.token_hash, record);
Ok(())
}
async fn take(&self, token_hash: &[u8; 32]) -> Result<Option<OttRecord>> {
Ok(self.inner.write().unwrap().remove(token_hash))
}
}