use async_trait::async_trait;
use etwin_core::clock::Clock;
use etwin_core::core::Instant;
use etwin_core::dinoparc::{DinoparcServer, DinoparcSessionKey, DinoparcUserIdRef, StoredDinoparcSession};
use etwin_core::hammerfest::{HammerfestServer, HammerfestSessionKey, HammerfestUserIdRef, StoredHammerfestSession};
use etwin_core::oauth::{RfcOauthAccessTokenKey, RfcOauthRefreshTokenKey, TwinoidAccessToken, TwinoidRefreshToken};
use etwin_core::token::{TokenStore, TouchOauthTokenOptions, TwinoidOauth};
use etwin_core::twinoid::{TwinoidUserId, TwinoidUserIdRef};
use etwin_core::types::AnyError;
use std::collections::HashMap;
use std::fmt::Debug;
use std::sync::RwLock;
struct MemDinoparcServers {
dinoparc_com: MemSessions<StoredDinoparcSession>,
en_dinoparc_com: MemSessions<StoredDinoparcSession>,
sp_dinoparc_com: MemSessions<StoredDinoparcSession>,
}
impl MemDinoparcServers {
fn new() -> Self {
Self {
dinoparc_com: MemSessions::new(),
en_dinoparc_com: MemSessions::new(),
sp_dinoparc_com: MemSessions::new(),
}
}
fn get(&self, server: DinoparcServer) -> &MemSessions<StoredDinoparcSession> {
match server {
DinoparcServer::DinoparcCom => &self.dinoparc_com,
DinoparcServer::EnDinoparcCom => &self.en_dinoparc_com,
DinoparcServer::SpDinoparcCom => &self.sp_dinoparc_com,
}
}
fn get_mut(&mut self, server: DinoparcServer) -> &mut MemSessions<StoredDinoparcSession> {
match server {
DinoparcServer::DinoparcCom => &mut self.dinoparc_com,
DinoparcServer::EnDinoparcCom => &mut self.en_dinoparc_com,
DinoparcServer::SpDinoparcCom => &mut self.sp_dinoparc_com,
}
}
}
struct MemHammerfestServers {
hammerfest_fr: MemSessions<StoredHammerfestSession>,
hammerfest_es: MemSessions<StoredHammerfestSession>,
hfest_net: MemSessions<StoredHammerfestSession>,
}
impl MemHammerfestServers {
fn new() -> Self {
Self {
hammerfest_fr: MemSessions::new(),
hammerfest_es: MemSessions::new(),
hfest_net: MemSessions::new(),
}
}
fn get(&self, server: HammerfestServer) -> &MemSessions<StoredHammerfestSession> {
match server {
HammerfestServer::HammerfestFr => &self.hammerfest_fr,
HammerfestServer::HfestNet => &self.hfest_net,
HammerfestServer::HammerfestEs => &self.hammerfest_es,
}
}
fn get_mut(&mut self, server: HammerfestServer) -> &mut MemSessions<StoredHammerfestSession> {
match server {
HammerfestServer::HammerfestFr => &mut self.hammerfest_fr,
HammerfestServer::HfestNet => &mut self.hfest_net,
HammerfestServer::HammerfestEs => &mut self.hammerfest_es,
}
}
}
struct MemSessions<Session: MemSession> {
sessions: HashMap<Session::SessionKey, Session>,
session_by_user: HashMap<Session::UserId, Session::SessionKey>,
}
impl<Session: MemSession> MemSessions<Session> {
fn new() -> Self {
Self {
sessions: HashMap::new(),
session_by_user: HashMap::new(),
}
}
fn touch(&mut self, now: Instant, key: &Session::SessionKey, user: Session::UserId) -> Session {
touch_session(&mut self.sessions, &mut self.session_by_user, now, key, user)
}
fn get(&self, user: Session::UserId) -> Option<Session> {
get_session(&self.sessions, &self.session_by_user, &user).cloned()
}
fn revoke(&mut self, key: &Session::SessionKey) {
revoke_session::<Session>(&mut self.sessions, &mut self.session_by_user, key)
}
}
trait MemSession: Clone {
type SessionKey: Debug + Clone + Eq + core::hash::Hash;
type UserId: Debug + Clone + Eq + core::hash::Hash;
fn new(now: Instant, key: Self::SessionKey, user_id: Self::UserId) -> Self;
fn user_id(&self) -> &Self::UserId;
fn atime_mut(&mut self) -> &mut Instant;
}
fn touch_session<Session: MemSession>(
sessions: &mut HashMap<Session::SessionKey, Session>,
session_by_user: &mut HashMap<Session::UserId, Session::SessionKey>,
now: Instant,
key: &Session::SessionKey,
user_id: Session::UserId,
) -> Session {
let old_session = sessions.get_mut(key);
if let Some(old_session) = old_session {
let old_user_id = old_session.user_id();
if old_user_id == &user_id {
*old_session.atime_mut() = now;
old_session.clone()
} else {
session_by_user.remove(old_user_id);
session_by_user.insert(user_id.clone(), key.clone());
let session = Session::new(now, key.clone(), user_id);
sessions.insert(key.clone(), session.clone());
session
}
} else {
session_by_user.insert(user_id.clone(), key.clone());
let session = Session::new(now, key.clone(), user_id);
sessions.insert(key.clone(), session.clone());
session
}
}
fn get_session<'a, Session: MemSession>(
sessions: &'a HashMap<Session::SessionKey, Session>,
session_by_user: &HashMap<Session::UserId, Session::SessionKey>,
user_id: &Session::UserId,
) -> Option<&'a Session> {
let key = session_by_user.get(user_id)?;
Some(sessions.get(key).unwrap())
}
impl MemSession for StoredHammerfestSession {
type SessionKey = HammerfestSessionKey;
type UserId = HammerfestUserIdRef;
fn new(now: Instant, key: Self::SessionKey, user_id: Self::UserId) -> Self {
Self {
key,
user: user_id,
ctime: now,
atime: now,
}
}
fn user_id(&self) -> &Self::UserId {
&self.user
}
fn atime_mut(&mut self) -> &mut Instant {
&mut self.atime
}
}
impl MemSession for StoredDinoparcSession {
type SessionKey = DinoparcSessionKey;
type UserId = DinoparcUserIdRef;
fn new(now: Instant, key: Self::SessionKey, user_id: Self::UserId) -> Self {
Self {
key,
user: user_id,
ctime: now,
atime: now,
}
}
fn user_id(&self) -> &Self::UserId {
&self.user
}
fn atime_mut(&mut self) -> &mut Instant {
&mut self.atime
}
}
fn revoke_session<Session: MemSession>(
sessions: &mut HashMap<Session::SessionKey, Session>,
session_by_user: &mut HashMap<Session::UserId, Session::SessionKey>,
key: &Session::SessionKey,
) {
let session = sessions.remove(key);
if let Some(session) = session {
let old_key = session_by_user.remove(session.user_id());
debug_assert_eq!(old_key.as_ref(), Some(key))
}
}
struct StoreState {
twinoid_access_tokens: HashMap<RfcOauthAccessTokenKey, TwinoidAccessToken>,
twinoid_refresh_tokens: HashMap<RfcOauthRefreshTokenKey, TwinoidRefreshToken>,
twinoid_user_to_access_token: HashMap<TwinoidUserId, RfcOauthAccessTokenKey>,
twinoid_user_to_refresh_token: HashMap<TwinoidUserId, RfcOauthRefreshTokenKey>,
dinoparc: MemDinoparcServers,
hammerfest: MemHammerfestServers,
}
impl StoreState {
fn new() -> Self {
Self {
twinoid_access_tokens: HashMap::new(),
twinoid_refresh_tokens: HashMap::new(),
twinoid_user_to_access_token: HashMap::new(),
twinoid_user_to_refresh_token: HashMap::new(),
dinoparc: MemDinoparcServers::new(),
hammerfest: MemHammerfestServers::new(),
}
}
fn get_twinoid_oauth(&self, now: Instant, id: TwinoidUserIdRef) -> TwinoidOauth {
let refresh_token = {
let refresh_key = self.twinoid_user_to_refresh_token.get(&id.id);
refresh_key.map(|refresh_key| self.twinoid_refresh_tokens.get(refresh_key).unwrap().clone())
};
let access_token = {
let access_key = self.twinoid_user_to_access_token.get(&id.id);
match access_key {
Some(access_key) => {
let at = self.twinoid_access_tokens.get(access_key).unwrap();
if now >= at.expires_at {
None
} else {
Some(at.clone())
}
}
None => None,
}
};
TwinoidOauth {
access_token,
refresh_token,
}
}
fn touch_twinoid_oauth(&mut self, now: Instant, options: &TouchOauthTokenOptions) {
if let Some(refresh_token) = options.refresh_token.as_ref() {
let old_user_id = self
.twinoid_refresh_tokens
.get(refresh_token)
.map(|t| t.twinoid_user_id);
if let Some(old_user_id) = old_user_id {
if old_user_id != options.twinoid_user_id {
self.twinoid_refresh_tokens.remove(refresh_token);
}
}
}
{
let old_user_id = self
.twinoid_access_tokens
.get(&options.access_token)
.map(|t| t.twinoid_user_id);
if let Some(old_user_id) = old_user_id {
if old_user_id != options.twinoid_user_id {
self.twinoid_access_tokens.remove(&options.access_token);
}
}
}
if let Some(refresh_token) = options.refresh_token.clone() {
self
.twinoid_refresh_tokens
.entry(refresh_token.clone())
.and_modify(|rt| rt.accessed_at = now)
.or_insert_with(|| TwinoidRefreshToken {
key: refresh_token,
created_at: now,
accessed_at: now,
twinoid_user_id: options.twinoid_user_id,
});
}
self
.twinoid_access_tokens
.entry(options.access_token.clone())
.and_modify(|rt| rt.accessed_at = now)
.or_insert_with(|| TwinoidAccessToken {
key: options.access_token.clone(),
created_at: now,
accessed_at: now,
expires_at: options.expiration_time,
twinoid_user_id: options.twinoid_user_id,
});
if let Some(refresh_token) = options.refresh_token.clone() {
let old_refresh_key = self
.twinoid_user_to_refresh_token
.insert(options.twinoid_user_id, refresh_token.clone());
if let Some(old_refresh_key) = old_refresh_key {
if old_refresh_key != refresh_token {
self.twinoid_refresh_tokens.remove(&old_refresh_key);
}
}
}
{
let old_access_key = self
.twinoid_user_to_access_token
.insert(options.twinoid_user_id, options.access_token.clone());
if let Some(old_access_key) = old_access_key {
if old_access_key != options.access_token {
self.twinoid_access_tokens.remove(&old_access_key);
}
}
}
}
fn revoke_twinoid_access_token(&mut self, options: &RfcOauthAccessTokenKey) {
let token = self.twinoid_access_tokens.remove(options);
if let Some(token) = token {
let old = self.twinoid_user_to_access_token.remove(&token.twinoid_user_id);
debug_assert_eq!(old.as_ref(), Some(options));
}
}
fn revoke_twinoid_refresh_token(&mut self, options: &RfcOauthRefreshTokenKey) {
let token = self.twinoid_refresh_tokens.remove(options);
if let Some(token) = token {
let old = self.twinoid_user_to_refresh_token.remove(&token.twinoid_user_id);
debug_assert_eq!(old.as_ref(), Some(options));
}
}
}
pub struct MemTokenStore<TyClock: Clock> {
clock: TyClock,
state: RwLock<StoreState>,
}
impl<TyClock> MemTokenStore<TyClock>
where
TyClock: Clock,
{
pub fn new(clock: TyClock) -> Self {
Self {
clock,
state: RwLock::new(StoreState::new()),
}
}
}
#[async_trait]
impl<TyClock> TokenStore for MemTokenStore<TyClock>
where
TyClock: Clock,
{
async fn touch_twinoid_oauth(&self, options: &TouchOauthTokenOptions) -> Result<(), AnyError> {
let mut state = self.state.write().unwrap();
let now = self.clock.now();
state.touch_twinoid_oauth(now, options);
Ok(())
}
async fn revoke_twinoid_access_token(&self, options: &RfcOauthAccessTokenKey) -> Result<(), AnyError> {
let mut state = self.state.write().unwrap();
state.revoke_twinoid_access_token(options);
Ok(())
}
async fn revoke_twinoid_refresh_token(&self, options: &RfcOauthRefreshTokenKey) -> Result<(), AnyError> {
let mut state = self.state.write().unwrap();
state.revoke_twinoid_refresh_token(options);
Ok(())
}
async fn get_twinoid_oauth(&self, options: TwinoidUserIdRef) -> Result<TwinoidOauth, AnyError> {
let state = self.state.read().unwrap();
let now = self.clock.now();
Ok(state.get_twinoid_oauth(now, options))
}
async fn touch_dinoparc(
&self,
user: DinoparcUserIdRef,
key: &DinoparcSessionKey,
) -> Result<StoredDinoparcSession, AnyError> {
let mut state = self.state.write().unwrap();
let server = state.dinoparc.get_mut(user.server);
let now = self.clock.now();
Ok(server.touch(now, key, user))
}
async fn revoke_dinoparc(&self, server: DinoparcServer, key: &DinoparcSessionKey) -> Result<(), AnyError> {
let mut state = self.state.write().unwrap();
let server = state.dinoparc.get_mut(server);
server.revoke(key);
Ok(())
}
async fn get_dinoparc(&self, user: DinoparcUserIdRef) -> Result<Option<StoredDinoparcSession>, AnyError> {
let state = self.state.read().unwrap();
let server = state.dinoparc.get(user.server);
Ok(server.get(user))
}
async fn touch_hammerfest(
&self,
user: HammerfestUserIdRef,
key: &HammerfestSessionKey,
) -> Result<StoredHammerfestSession, AnyError> {
let mut state = self.state.write().unwrap();
let server = state.hammerfest.get_mut(user.server);
let now = self.clock.now();
Ok(server.touch(now, key, user))
}
async fn revoke_hammerfest(&self, server: HammerfestServer, key: &HammerfestSessionKey) -> Result<(), AnyError> {
let mut state = self.state.write().unwrap();
let server = state.hammerfest.get_mut(server);
server.revoke(key);
Ok(())
}
async fn get_hammerfest(&self, user: HammerfestUserIdRef) -> Result<Option<StoredHammerfestSession>, AnyError> {
let state = self.state.read().unwrap();
let server = state.hammerfest.get(user.server);
Ok(server.get(user))
}
}
#[cfg(feature = "neon")]
impl<TyClock> neon::prelude::Finalize for MemTokenStore<TyClock> where TyClock: Clock {}
#[cfg(test)]
mod test {
use crate::mem::MemTokenStore;
use crate::test::TestApi;
use etwin_core::clock::VirtualClock;
use etwin_core::core::Instant;
use etwin_core::dinoparc::DinoparcStore;
use etwin_core::hammerfest::HammerfestStore;
use etwin_core::token::TokenStore;
use etwin_core::twinoid::TwinoidStore;
use etwin_dinoparc_store::mem::MemDinoparcStore;
use etwin_hammerfest_store::mem::MemHammerfestStore;
use etwin_twinoid_store::mem::MemTwinoidStore;
use std::sync::Arc;
#[allow(clippy::type_complexity)]
fn make_test_api() -> TestApi<
Arc<VirtualClock>,
Arc<dyn DinoparcStore>,
Arc<dyn HammerfestStore>,
Arc<dyn TokenStore>,
Arc<dyn TwinoidStore>,
> {
let clock = Arc::new(VirtualClock::new(Instant::ymd_hms(2020, 1, 1, 0, 0, 0)));
let dinoparc_store: Arc<dyn DinoparcStore> = Arc::new(MemDinoparcStore::new(Arc::clone(&clock)));
let hammerfest_store: Arc<dyn HammerfestStore> = Arc::new(MemHammerfestStore::new(Arc::clone(&clock)));
let twinoid_store: Arc<dyn TwinoidStore> = Arc::new(MemTwinoidStore::new(Arc::clone(&clock)));
let token_store: Arc<dyn TokenStore> = Arc::new(MemTokenStore::new(Arc::clone(&clock)));
TestApi {
clock,
dinoparc_store,
hammerfest_store,
token_store,
twinoid_store,
}
}
test_token_store!(|| make_test_api());
}