use std::{ops::Deref, sync::Arc}; use async_sqlx_session::PostgresSessionStore; use async_trait::async_trait; use axum_sessions::async_session::{Session as AxumSession, SessionStore as AxumSessionStore}; use serde::{Deserialize, Serialize}; use crate::{AuthClap, SessionBackend}; #[derive(clap::Args, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct SessionClap { #[clap(flatten)] pub postgresql: PostgresqlSessionClap, } #[derive(clap::Args, Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct PostgresqlSessionClap { #[arg(env = "SESSION_POSTGRES_CONN", long = "session-postgres-conn")] pub conn: Option, } #[async_trait] pub trait Session { async fn insert_user(&self, id: &str, user_id: &str) -> anyhow::Result; async fn get_user(&self, cookie: &str) -> anyhow::Result>; } pub struct SessionService(Arc); impl SessionService { pub async fn new(config: &AuthClap) -> anyhow::Result { match config.session_backend { SessionBackend::InMemory => Ok(Self(Arc::new(InMemorySessionService {}))), SessionBackend::Postgresql => { let postgres_session = PostgresSessionStore::new( config .session .postgresql .conn .as_ref() .expect("SESSION_POSTGRES_CONN to be set"), ) .await?; Ok(Self(Arc::new(PostgresSessionService { store: postgres_session, }))) } } } } impl Deref for SessionService { type Target = Arc; fn deref(&self) -> &Self::Target { &self.0 } } pub struct PostgresSessionService { store: PostgresSessionStore, } #[derive(Serialize, Deserialize, Clone, Debug)] pub struct User { pub id: String, } #[async_trait] impl Session for PostgresSessionService { async fn insert_user(&self, _id: &str, user_id: &str) -> anyhow::Result { let mut session = AxumSession::new(); session.insert( "user", User { id: user_id.to_string(), }, )?; let cookie = self .store .store_session(session) .await? .ok_or(anyhow::anyhow!("failed to store session"))?; Ok(cookie) } async fn get_user(&self, cookie: &str) -> anyhow::Result> { if let Some(session) = self.store.load_session(cookie.to_string()).await.unwrap() { if let Some(user) = session.get::("user") { tracing::debug!( "UserFromSession: session decoded success, user_id={:?}", user.id ); Ok(Some(user.id)) } else { Ok(None) } } else { tracing::debug!( "UserIdFromSession: err session not exists in store, {}", cookie ); Err(anyhow::anyhow!("No session found for cookie")) } } } pub struct InMemorySessionService {} #[async_trait] impl Session for InMemorySessionService { async fn insert_user(&self, _id: &str, _user_id: &str) -> anyhow::Result { todo!() } async fn get_user(&self, _cookie: &str) -> anyhow::Result> { todo!() } }