Something went wrong. Try again.
A downstream fork of Tranquil PDS for use in margin.cafe
Something went wrong. Try again.
20 kB · 604 lines
Rust
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605use async_trait::async_trait;use chrono::{DateTime, Duration, Utc};use sqlx::PgPool;use tranquil_db_traits::{ AppPasswordCreate, AppPasswordPrivilege, AppPasswordRecord, DbError, LoginType, REFRESH_GRACE_PERIOD_SECS, RefreshGraceLookup, RefreshGraceReplay, RefreshSessionResult, SessionForRefresh, SessionId, SessionListItem, SessionMfaStatus, SessionRefreshData, SessionRepository, SessionToken, SessionTokenCreate,};use tranquil_types::{Did, Jti, PasswordHash};use uuid::Uuid;
use super::col;use super::user::map_sqlx_error;use super::{column, opt_column};
pub struct PostgresSessionRepository { pool: PgPool,}
impl PostgresSessionRepository { pub fn new(pool: PgPool) -> Self { Self { pool } }}
#[async_trait]impl SessionRepository for PostgresSessionRepository { async fn create_session(&self, data: &SessionTokenCreate) -> Result<SessionId, DbError> { let row = sqlx::query!( r#" INSERT INTO session_tokens (did, access_jti, refresh_jti, access_expires_at, refresh_expires_at, legacy_login, mfa_verified, scope, controller_did, app_password_name) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) RETURNING id "#, data.did.as_str(), data.access_jti.as_str(), data.refresh_jti.as_str(), data.access_expires_at, data.refresh_expires_at, data.login_type.is_legacy(), data.mfa_verified, data.scope, data.controller_did.as_ref().map(|d| d.as_str()), data.app_password_name ) .fetch_one(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(SessionId::new(row.id)) }
async fn get_session_by_access_jti( &self, access_jti: &Jti, ) -> Result<Option<SessionToken>, DbError> { let row = sqlx::query!( r#" SELECT id, did, access_jti, refresh_jti, access_expires_at, refresh_expires_at, legacy_login, mfa_verified, scope, controller_did, app_password_name, created_at, updated_at FROM session_tokens WHERE access_jti = $1 "#, access_jti.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
row.map(|r| { Ok(SessionToken { id: SessionId::new(r.id), did: column(r.did, col::SESSION_TOKENS_DID)?, access_jti: Jti::from(r.access_jti), refresh_jti: Jti::from(r.refresh_jti), access_expires_at: r.access_expires_at, refresh_expires_at: r.refresh_expires_at, login_type: LoginType::from_legacy_flag(r.legacy_login), mfa_verified: r.mfa_verified, scope: r.scope, controller_did: opt_column(r.controller_did, col::SESSION_TOKENS_CONTROLLER_DID)?, app_password_name: r.app_password_name, created_at: r.created_at, updated_at: r.updated_at, }) }) .transpose() }
async fn get_session_for_refresh( &self, refresh_jti: &Jti, ) -> Result<Option<SessionForRefresh>, DbError> { let row = sqlx::query!( r#" SELECT st.id, st.did, st.scope, st.controller_did, k.key_bytes, k.encryption_version FROM session_tokens st JOIN users u ON st.did = u.did JOIN user_keys k ON u.id = k.user_id WHERE st.refresh_jti = $1 AND st.refresh_expires_at > NOW() "#, refresh_jti.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
row.map(|r| { Ok(SessionForRefresh { id: SessionId::new(r.id), did: column(r.did, col::SESSION_TOKENS_DID)?, scope: r.scope, controller_did: opt_column(r.controller_did, col::SESSION_TOKENS_CONTROLLER_DID)?, key_bytes: r.key_bytes, encryption_version: r.encryption_version.unwrap_or(0), }) }) .transpose() }
async fn delete_session_by_access_jti( &self, access_jti: &Jti, did: &Did, ) -> Result<u64, DbError> { let result = sqlx::query!( "DELETE FROM session_tokens WHERE access_jti = $1 AND did = $2", access_jti.as_str(), did.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn delete_session_by_id(&self, session_id: SessionId, did: &Did) -> Result<u64, DbError> { let result = sqlx::query!( "DELETE FROM session_tokens WHERE id = $1 AND did = $2", session_id.as_i32(), did.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn delete_sessions_by_did(&self, did: &Did) -> Result<u64, DbError> { let result = sqlx::query!("DELETE FROM session_tokens WHERE did = $1", did.as_str()) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn delete_sessions_by_did_except_jti( &self, did: &Did, except_jti: &Jti, ) -> Result<u64, DbError> { let result = sqlx::query!( "DELETE FROM session_tokens WHERE did = $1 AND access_jti != $2", did.as_str(), except_jti.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn list_sessions_by_did(&self, did: &Did) -> Result<Vec<SessionListItem>, DbError> { let rows = sqlx::query!( r#" SELECT id, access_jti, created_at, refresh_expires_at FROM session_tokens WHERE did = $1 AND refresh_expires_at > NOW() ORDER BY created_at DESC "#, did.as_str() ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(rows .into_iter() .map(|r| SessionListItem { id: SessionId::new(r.id), access_jti: Jti::from(r.access_jti), created_at: r.created_at, refresh_expires_at: r.refresh_expires_at, }) .collect()) }
async fn get_session_access_jti_by_id( &self, session_id: SessionId, did: &Did, ) -> Result<Option<Jti>, DbError> { let row = sqlx::query_scalar!( "SELECT access_jti FROM session_tokens WHERE id = $1 AND did = $2", session_id.as_i32(), did.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(row.map(Jti::from)) }
async fn delete_sessions_by_app_password( &self, did: &Did, app_password_name: &str, ) -> Result<u64, DbError> { let result = sqlx::query!( "DELETE FROM session_tokens WHERE did = $1 AND app_password_name = $2", did.as_str(), app_password_name ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn get_session_jtis_by_app_password( &self, did: &Did, app_password_name: &str, ) -> Result<Vec<Jti>, DbError> { let rows = sqlx::query_scalar!( "SELECT access_jti FROM session_tokens WHERE did = $1 AND app_password_name = $2", did.as_str(), app_password_name ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(rows.into_iter().map(Jti::from).collect()) }
async fn lookup_refresh_grace(&self, refresh_jti: &Jti) -> Result<RefreshGraceLookup, DbError> { let row = sqlx::query!( r#" SELECT u.used_at, st.id AS session_id, st.did, st.scope, st.controller_did, st.access_jti, st.refresh_jti, st.access_expires_at, st.refresh_expires_at, k.key_bytes, k.encryption_version FROM used_refresh_tokens u JOIN session_tokens st ON st.id = u.session_id JOIN users us ON st.did = us.did JOIN user_keys k ON us.id = k.user_id WHERE u.refresh_jti = $1 "#, refresh_jti.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
// No marker (or a missing users/user_keys join row) degrades to NotUsed. // That is safe: the normal refresh path then fails closed with "Invalid // refresh token" without mutating any state. let Some(r) = row else { return Ok(RefreshGraceLookup::NotUsed); };
let grace_cutoff = Utc::now() - Duration::seconds(REFRESH_GRACE_PERIOD_SECS); if r.used_at > grace_cutoff { Ok(RefreshGraceLookup::Replay(RefreshGraceReplay { did: column(r.did, col::SESSION_TOKENS_DID)?, scope: r.scope, controller_did: opt_column(r.controller_did, col::SESSION_TOKENS_CONTROLLER_DID)?, access_jti: Jti::from(r.access_jti), refresh_jti: Jti::from(r.refresh_jti), access_expires_at: r.access_expires_at, refresh_expires_at: r.refresh_expires_at, key_bytes: r.key_bytes, encryption_version: r.encryption_version.unwrap_or(0), })) } else { Ok(RefreshGraceLookup::Compromised { did: column(r.did, col::SESSION_TOKENS_DID)?, session_id: SessionId::new(r.session_id), key_bytes: r.key_bytes, encryption_version: r.encryption_version.unwrap_or(0), }) } }
async fn list_app_passwords(&self, user_id: Uuid) -> Result<Vec<AppPasswordRecord>, DbError> { let rows = sqlx::query!( r#" SELECT id, user_id, name, password_hash, created_at, privileged, scopes, created_by_controller_did FROM app_passwords WHERE user_id = $1 ORDER BY created_at DESC "#, user_id ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
rows.into_iter() .map(|r| { Ok(AppPasswordRecord { id: r.id, user_id: r.user_id, name: r.name, password_hash: PasswordHash::new(r.password_hash), created_at: r.created_at, privilege: AppPasswordPrivilege::from_privileged_flag(r.privileged), scopes: r.scopes, created_by_controller_did: opt_column( r.created_by_controller_did, col::APP_PASSWORDS_CREATED_BY_CONTROLLER_DID, )?, }) }) .collect() }
async fn get_app_passwords_for_login( &self, user_id: Uuid, ) -> Result<Vec<AppPasswordRecord>, DbError> { let rows = sqlx::query!( r#" SELECT id, user_id, name, password_hash, created_at, privileged, scopes, created_by_controller_did FROM app_passwords WHERE user_id = $1 ORDER BY created_at DESC LIMIT 20 "#, user_id ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
rows.into_iter() .map(|r| { Ok(AppPasswordRecord { id: r.id, user_id: r.user_id, name: r.name, password_hash: PasswordHash::new(r.password_hash), created_at: r.created_at, privilege: AppPasswordPrivilege::from_privileged_flag(r.privileged), scopes: r.scopes, created_by_controller_did: opt_column( r.created_by_controller_did, col::APP_PASSWORDS_CREATED_BY_CONTROLLER_DID, )?, }) }) .collect() }
async fn get_app_password_by_name( &self, user_id: Uuid, name: &str, ) -> Result<Option<AppPasswordRecord>, DbError> { let row = sqlx::query!( r#" SELECT id, user_id, name, password_hash, created_at, privileged, scopes, created_by_controller_did FROM app_passwords WHERE user_id = $1 AND name = $2 "#, user_id, name ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
row.map(|r| { Ok(AppPasswordRecord { id: r.id, user_id: r.user_id, name: r.name, password_hash: PasswordHash::new(r.password_hash), created_at: r.created_at, privilege: AppPasswordPrivilege::from_privileged_flag(r.privileged), scopes: r.scopes, created_by_controller_did: opt_column( r.created_by_controller_did, col::APP_PASSWORDS_CREATED_BY_CONTROLLER_DID, )?, }) }) .transpose() }
async fn create_app_password(&self, data: &AppPasswordCreate) -> Result<Uuid, DbError> { let row = sqlx::query!( r#" INSERT INTO app_passwords (user_id, name, password_hash, privileged, scopes, created_by_controller_did) VALUES ($1, $2, $3, $4, $5, $6) RETURNING id "#, data.user_id, data.name, data.password_hash.as_str(), data.privilege.is_privileged(), data.scopes, data.created_by_controller_did.as_ref().map(|d| d.as_str()) ) .fetch_one(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(row.id) }
async fn delete_app_password(&self, user_id: Uuid, name: &str) -> Result<u64, DbError> { let result = sqlx::query!( "DELETE FROM app_passwords WHERE user_id = $1 AND name = $2", user_id, name ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn delete_app_passwords_by_controller( &self, did: &Did, controller_did: &Did, ) -> Result<u64, DbError> { let result = sqlx::query!( r#"DELETE FROM app_passwords WHERE user_id = (SELECT id FROM users WHERE did = $1) AND created_by_controller_did = $2"#, did.as_str(), controller_did.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(result.rows_affected()) }
async fn get_last_reauth_at(&self, did: &Did) -> Result<Option<DateTime<Utc>>, DbError> { let row = sqlx::query_scalar!( r#"SELECT last_reauth_at FROM session_tokens WHERE did = $1 ORDER BY created_at DESC LIMIT 1"#, did.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(row.flatten()) }
async fn update_last_reauth(&self, did: &Did) -> Result<DateTime<Utc>, DbError> { let now = Utc::now(); sqlx::query!( "UPDATE session_tokens SET last_reauth_at = $1, mfa_verified = TRUE WHERE did = $2", now, did.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(now) }
async fn get_session_mfa_status(&self, did: &Did) -> Result<Option<SessionMfaStatus>, DbError> { let row = sqlx::query!( r#"SELECT legacy_login, mfa_verified, last_reauth_at FROM session_tokens WHERE did = $1 ORDER BY created_at DESC LIMIT 1"#, did.as_str() ) .fetch_optional(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(row.map(|r| SessionMfaStatus { login_type: LoginType::from_legacy_flag(r.legacy_login), mfa_verified: r.mfa_verified, last_reauth_at: r.last_reauth_at, })) }
async fn update_mfa_verified(&self, did: &Did) -> Result<(), DbError> { sqlx::query!( "UPDATE session_tokens SET mfa_verified = TRUE, last_reauth_at = NOW() WHERE did = $1", did.as_str() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(()) }
async fn get_app_password_hashes_by_did( &self, did: &Did, ) -> Result<Vec<PasswordHash>, DbError> { let rows = sqlx::query_scalar!( r#"SELECT ap.password_hash FROM app_passwords ap JOIN users u ON ap.user_id = u.id WHERE u.did = $1"#, did.as_str() ) .fetch_all(&self.pool) .await .map_err(map_sqlx_error)?;
Ok(rows.into_iter().map(PasswordHash::new).collect()) }
async fn refresh_session_atomic( &self, data: &SessionRefreshData, ) -> Result<RefreshSessionResult, DbError> { let mut tx = self.pool.begin().await.map_err(map_sqlx_error)?;
// Atomically claim the old refresh jti. The INSERT serializes concurrent // rotations of the same token: exactly one request inserts the row, the // rest see `rows_affected == 0`. let claimed = sqlx::query!( "INSERT INTO used_refresh_tokens (refresh_jti, session_id) VALUES ($1, $2) ON CONFLICT (refresh_jti) DO NOTHING", data.old_refresh_jti.as_str(), data.session_id.as_i32() ) .execute(&mut *tx) .await .map_err(map_sqlx_error)?;
if claimed.rows_affected() == 0 { // Another request already rotated this token. Nothing to write, so // end our transaction before reading the winner's committed row. tx.rollback().await.map_err(map_sqlx_error)?;
// Within the grace window (measured from this token's own rotation // time) we replay the session's current tokens so a benignly-racing // client keeps a working session instead of being revoked. match self.lookup_refresh_grace(&data.old_refresh_jti).await? { RefreshGraceLookup::Replay(replay) => { return Ok(RefreshSessionResult::GraceReplay(replay)); } RefreshGraceLookup::Compromised { .. } | RefreshGraceLookup::NotUsed => { // Outside the grace window, or the marker/session vanished // concurrently: genuine reuse. Revoke the session (delete is // idempotent). sqlx::query!( "DELETE FROM session_tokens WHERE id = $1", data.session_id.as_i32() ) .execute(&self.pool) .await .map_err(map_sqlx_error)?; return Ok(RefreshSessionResult::Compromise); } } }
// We won the rotation. sqlx::query!( r#" UPDATE session_tokens SET access_jti = $1, refresh_jti = $2, access_expires_at = $3, refresh_expires_at = $4, updated_at = NOW() WHERE id = $5 "#, data.new_access_jti.as_str(), data.new_refresh_jti.as_str(), data.new_access_expires_at, data.new_refresh_expires_at, data.session_id.as_i32() ) .execute(&mut *tx) .await .map_err(map_sqlx_error)?;
tx.commit().await.map_err(map_sqlx_error)?; Ok(RefreshSessionResult::Success) }}