use std::collections::HashMap;
use std::sync::Arc;
use async_trait::async_trait;
use parking_lot::RwLock;
use crate::error::{Result, RiftError};
use crate::protocol::hello::AuthMode;
use crate::session::session::ClientId;
#[derive(Debug, Clone)]
pub struct AuthContext {
pub client_id: ClientId,
pub claims: serde_json::Value,
pub mode: AuthMode,
pub hints: AuthHints,
}
#[derive(Debug, Clone, Default)]
pub struct AuthHints {
pub region: Option<String>,
pub device: Option<String>,
pub risk: Option<u32>,
}
#[async_trait]
pub trait AuthProvider: Send + Sync {
async fn authenticate(&self, mode: AuthMode, token: Option<&str>) -> Result<AuthContext>;
async fn revoke(&self, client_id: &ClientId) -> Result<()>;
}
pub struct TokenAuth {
tokens: RwLock<HashMap<String, AuthContext>>,
by_client: RwLock<HashMap<String, Vec<String>>>,
}
impl TokenAuth {
pub fn new() -> Self {
Self {
tokens: RwLock::new(HashMap::new()),
by_client: RwLock::new(HashMap::new()),
}
}
pub fn register(&self, token: impl Into<String>, ctx: AuthContext) {
let token = token.into();
let client_id = ctx.client_id.0.clone();
self.by_client
.write()
.entry(client_id)
.or_default()
.push(token.clone());
self.tokens.write().insert(token, ctx);
}
}
impl Default for TokenAuth {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl AuthProvider for TokenAuth {
async fn authenticate(&self, mode: AuthMode, token: Option<&str>) -> Result<AuthContext> {
if mode == AuthMode::Anonymous {
return Ok(AuthContext {
client_id: ClientId::new(format!("anon-{}", rand_suffix())),
claims: serde_json::json!({}),
mode,
hints: AuthHints::default(),
});
}
let token = token.ok_or_else(|| RiftError::Auth(crate::error::AuthReject::Required))?;
self.tokens.read().get(token).cloned().ok_or_else(|| {
RiftError::Auth(crate::error::AuthReject::Invalid("unknown token".into()))
})
}
async fn revoke(&self, client_id: &ClientId) -> Result<()> {
let tokens = {
let mut bc = self.by_client.write();
bc.remove(client_id.as_str()).unwrap_or_default()
};
let mut t = self.tokens.write();
for token in &tokens {
t.remove(token);
}
Ok(())
}
}
fn rand_suffix() -> String {
ulid::Ulid::new().to_string()
}
pub struct AllowAllAuth;
#[async_trait]
impl AuthProvider for AllowAllAuth {
async fn authenticate(&self, _mode: AuthMode, _token: Option<&str>) -> Result<AuthContext> {
Ok(AuthContext {
client_id: ClientId::new("anonymous"),
claims: serde_json::json!({}),
mode: AuthMode::Anonymous,
hints: AuthHints::default(),
})
}
async fn revoke(&self, _client_id: &ClientId) -> Result<()> {
Ok(())
}
}
pub fn shared<A: AuthProvider + 'static>(a: A) -> Arc<dyn AuthProvider> {
Arc::new(a)
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn token_auth_round_trip() {
let auth = TokenAuth::new();
auth.register(
"tok-1",
AuthContext {
client_id: ClientId::new("user-1"),
claims: serde_json::json!({"sub": "user-1"}),
mode: AuthMode::Bearer,
hints: AuthHints::default(),
},
);
let ctx = auth
.authenticate(AuthMode::Bearer, Some("tok-1"))
.await
.unwrap();
assert_eq!(ctx.client_id.as_str(), "user-1");
assert!(
auth.authenticate(AuthMode::Bearer, Some("nope"))
.await
.is_err()
);
}
#[tokio::test]
async fn anonymous_auth_works_without_token() {
let auth = TokenAuth::new();
let ctx = auth.authenticate(AuthMode::Anonymous, None).await.unwrap();
assert_eq!(ctx.mode, AuthMode::Anonymous);
}
#[tokio::test]
async fn revoke_removes_ability_to_authenticate() {
let auth = TokenAuth::new();
let client = ClientId::new("user-1");
auth.register(
"tok-1",
AuthContext {
client_id: client.clone(),
claims: serde_json::json!({}),
mode: AuthMode::Bearer,
hints: AuthHints::default(),
},
);
assert!(
auth.authenticate(AuthMode::Bearer, Some("tok-1"))
.await
.is_ok()
);
auth.revoke(&client).await.unwrap();
assert!(
auth.authenticate(AuthMode::Bearer, Some("tok-1"))
.await
.is_err()
);
}
}