Something went wrong. Try again.
A lexicon-driven AppView for ATProto. forked from gamesgamesgamesgames.games/happyview
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197use atrium_api::types::string::Did;use atrium_common::store::Store;use atrium_oauth::store::session::{Session, SessionStore};use atrium_oauth::store::state::{InternalStateData, StateStore};use sqlx::AnyPool;use std::fmt;use std::sync::{Arc, Mutex};
use crate::db::{DatabaseBackend, adapt_sql};
#[derive(Debug)]pub enum StoreError { Sqlx(sqlx::Error), Json(serde_json::Error),}
impl fmt::Display for StoreError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { StoreError::Sqlx(e) => write!(f, "database error: {e}"), StoreError::Json(e) => write!(f, "json error: {e}"), } }}
impl std::error::Error for StoreError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { StoreError::Sqlx(e) => Some(e), StoreError::Json(e) => Some(e), } }}
impl From<sqlx::Error> for StoreError { fn from(e: sqlx::Error) -> Self { StoreError::Sqlx(e) }}
impl From<serde_json::Error> for StoreError { fn from(e: serde_json::Error) -> Self { StoreError::Json(e) }}
// --- DbSessionStore ---
#[derive(Clone)]pub struct DbSessionStore { pool: AnyPool, backend: DatabaseBackend,}
impl DbSessionStore { pub fn new(pool: AnyPool, backend: DatabaseBackend) -> Self { Self { pool, backend } }}
impl Store<Did, Session> for DbSessionStore { type Error = StoreError;
async fn get(&self, key: &Did) -> Result<Option<Session>, Self::Error> { let row: Option<(String,)> = sqlx::query_as(&adapt_sql( "SELECT session_data FROM oauth_sessions WHERE did = ?", self.backend, )) .bind(key.as_ref()) .fetch_optional(&self.pool) .await?;
match row { Some((data,)) => Ok(Some(serde_json::from_str(&data)?)), None => Ok(None), } }
async fn set(&self, key: Did, value: Session) -> Result<(), Self::Error> { let json = serde_json::to_string(&value)?; sqlx::query(&adapt_sql( "INSERT INTO oauth_sessions (did, session_data, updated_at) VALUES (?, ?, datetime('now')) ON CONFLICT (did) DO UPDATE SET session_data = EXCLUDED.session_data, updated_at = datetime('now')", self.backend, )) .bind(key.as_ref()) .bind(&json) .execute(&self.pool) .await?; Ok(()) }
async fn del(&self, key: &Did) -> Result<(), Self::Error> { sqlx::query(&adapt_sql( "DELETE FROM oauth_sessions WHERE did = ?", self.backend, )) .bind(key.as_ref()) .execute(&self.pool) .await?; Ok(()) }
async fn clear(&self) -> Result<(), Self::Error> { sqlx::query("DELETE FROM oauth_sessions") .execute(&self.pool) .await?; Ok(()) }}
impl SessionStore for DbSessionStore {}
// --- DbStateStore ---
#[derive(Clone)]pub struct DbStateStore { pool: AnyPool, backend: DatabaseBackend, /// Captures the most recently stored state key so callers can associate /// additional data (e.g., redirect URIs) with the OAuth state. last_state_key: Arc<Mutex<Option<String>>>, /// Serializes authorize() + take_last_state_key() pairs so concurrent /// logins cannot interleave and swap each other's state keys. pub authorize_lock: Arc<tokio::sync::Mutex<()>>,}
impl DbStateStore { pub fn new(pool: AnyPool, backend: DatabaseBackend) -> Self { Self { pool, backend, last_state_key: Arc::new(Mutex::new(None)), authorize_lock: Arc::new(tokio::sync::Mutex::new(())), } }
/// Returns the state key from the most recent `set()` call, clearing it. pub fn take_last_state_key(&self) -> Option<String> { self.last_state_key.lock().unwrap().take() }}
impl Store<String, InternalStateData> for DbStateStore { type Error = StoreError;
async fn get(&self, key: &String) -> Result<Option<InternalStateData>, Self::Error> { let row: Option<(String,)> = sqlx::query_as(&adapt_sql( "SELECT state_data FROM oauth_state WHERE state_key = ?", self.backend, )) .bind(key) .fetch_optional(&self.pool) .await?;
match row { Some((data,)) => Ok(Some(serde_json::from_str(&data)?)), None => Ok(None), } }
async fn set(&self, key: String, value: InternalStateData) -> Result<(), Self::Error> { let json = serde_json::to_string(&value)?; sqlx::query(&adapt_sql( "INSERT INTO oauth_state (state_key, state_data) VALUES (?, ?) ON CONFLICT (state_key) DO UPDATE SET state_data = EXCLUDED.state_data", self.backend, )) .bind(&key) .bind(&json) .execute(&self.pool) .await?; *self.last_state_key.lock().unwrap() = Some(key); Ok(()) }
async fn del(&self, key: &String) -> Result<(), Self::Error> { sqlx::query(&adapt_sql( "DELETE FROM oauth_state WHERE state_key = ?", self.backend, )) .bind(key) .execute(&self.pool) .await?; Ok(()) }
async fn clear(&self) -> Result<(), Self::Error> { sqlx::query("DELETE FROM oauth_state") .execute(&self.pool) .await?; Ok(()) }}
impl StateStore for DbStateStore {}