use std::collections::{HashMap, HashSet};
use std::time::Duration;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone)]
pub struct AuthConfig {
pub enabled: bool,
pub jwt: Option<JwtConfig>,
pub oauth: Option<OAuthConfig>,
pub ldap: Option<LdapConfig>,
pub api_keys: Option<ApiKeyConfig>,
pub role_mapping: Vec<RoleMappingRule>,
pub default_role: Option<String>,
pub credentials: CredentialConfig,
pub session: SessionConfig,
pub rate_limit: AuthRateLimitConfig,
pub auth_methods: Vec<AuthMethod>,
}
impl Default for AuthConfig {
fn default() -> Self {
Self {
enabled: false,
jwt: None,
oauth: None,
ldap: None,
api_keys: None,
role_mapping: Vec::new(),
default_role: Some("db_minimal".to_string()),
credentials: CredentialConfig::default(),
session: SessionConfig::default(),
rate_limit: AuthRateLimitConfig::default(),
auth_methods: Vec::new(),
}
}
}
impl AuthConfig {
pub fn jwt(jwks_url: impl Into<String>) -> Self {
Self {
enabled: true,
jwt: Some(JwtConfig::new(jwks_url)),
..Default::default()
}
}
pub fn api_keys() -> Self {
Self {
enabled: true,
api_keys: Some(ApiKeyConfig::default()),
..Default::default()
}
}
pub fn builder() -> AuthConfigBuilder {
AuthConfigBuilder::new()
}
}
#[derive(Default)]
pub struct AuthConfigBuilder {
config: AuthConfig,
}
impl AuthConfigBuilder {
pub fn new() -> Self {
Self {
config: AuthConfig {
enabled: true,
..Default::default()
},
}
}
pub fn jwt(mut self, config: JwtConfig) -> Self {
self.config.jwt = Some(config);
self
}
pub fn oauth(mut self, config: OAuthConfig) -> Self {
self.config.oauth = Some(config);
self
}
pub fn ldap(mut self, config: LdapConfig) -> Self {
self.config.ldap = Some(config);
self
}
pub fn api_keys(mut self, config: ApiKeyConfig) -> Self {
self.config.api_keys = Some(config);
self
}
pub fn add_role_mapping(mut self, rule: RoleMappingRule) -> Self {
self.config.role_mapping.push(rule);
self
}
pub fn default_role(mut self, role: impl Into<String>) -> Self {
self.config.default_role = Some(role.into());
self
}
pub fn credentials(mut self, config: CredentialConfig) -> Self {
self.config.credentials = config;
self
}
pub fn session(mut self, config: SessionConfig) -> Self {
self.config.session = config;
self
}
pub fn build(self) -> AuthConfig {
self.config
}
}
#[derive(Debug, Clone)]
pub struct JwtConfig {
pub jwks_url: String,
pub jwks_refresh_interval: Duration,
pub allowed_issuers: HashSet<String>,
pub required_audience: Option<String>,
pub clock_skew: Duration,
pub user_id_claim: String,
pub roles_claim: Option<String>,
pub allowed_algorithms: Vec<String>,
}
impl Default for JwtConfig {
fn default() -> Self {
Self {
jwks_url: String::new(),
jwks_refresh_interval: Duration::from_secs(3600),
allowed_issuers: HashSet::new(),
required_audience: None,
clock_skew: Duration::from_secs(60),
user_id_claim: "sub".to_string(),
roles_claim: Some("roles".to_string()),
allowed_algorithms: vec!["RS256".to_string(), "ES256".to_string()],
}
}
}
impl JwtConfig {
pub fn new(jwks_url: impl Into<String>) -> Self {
Self {
jwks_url: jwks_url.into(),
..Default::default()
}
}
pub fn with_issuer(mut self, issuer: impl Into<String>) -> Self {
self.allowed_issuers.insert(issuer.into());
self
}
pub fn with_audience(mut self, audience: impl Into<String>) -> Self {
self.required_audience = Some(audience.into());
self
}
}
#[derive(Debug, Clone)]
pub struct OAuthConfig {
pub introspection_url: String,
pub client_id: String,
pub client_secret: String,
pub token_url: Option<String>,
pub scopes: Vec<String>,
pub cache_ttl: Duration,
pub required_scopes: Vec<String>,
pub issuer: String,
pub authorization_url: Option<String>,
pub audience: Option<String>,
}
impl Default for OAuthConfig {
fn default() -> Self {
Self {
introspection_url: String::new(),
client_id: String::new(),
client_secret: String::new(),
token_url: None,
scopes: Vec::new(),
cache_ttl: Duration::from_secs(60),
required_scopes: Vec::new(),
issuer: String::new(),
authorization_url: None,
audience: None,
}
}
}
impl OAuthConfig {
pub fn new(
introspection_url: impl Into<String>,
client_id: impl Into<String>,
client_secret: impl Into<String>,
) -> Self {
Self {
introspection_url: introspection_url.into(),
client_id: client_id.into(),
client_secret: client_secret.into(),
..Default::default()
}
}
}
#[derive(Debug, Clone)]
pub struct LdapConfig {
pub server_url: String,
pub bind_dn: String,
pub bind_password: String,
pub user_search_base: String,
pub user_filter: String,
pub group_search_base: Option<String>,
pub group_attribute: String,
pub timeout: Duration,
pub starttls: bool,
}
impl Default for LdapConfig {
fn default() -> Self {
Self {
server_url: "ldap://localhost:389".to_string(),
bind_dn: String::new(),
bind_password: String::new(),
user_search_base: String::new(),
user_filter: "(uid={0})".to_string(),
group_search_base: None,
group_attribute: "memberOf".to_string(),
timeout: Duration::from_secs(10),
starttls: false,
}
}
}
#[derive(Debug, Clone)]
pub struct ApiKeyConfig {
pub header_name: String,
pub query_param: Option<String>,
pub prefix: Option<String>,
pub hash_algorithm: String,
}
impl Default for ApiKeyConfig {
fn default() -> Self {
Self {
header_name: "X-API-Key".to_string(),
query_param: None,
prefix: Some("hpk_".to_string()),
hash_algorithm: "sha256".to_string(),
}
}
}
#[derive(Debug, Clone)]
pub struct RoleMappingRule {
pub name: String,
pub condition: RoleCondition,
pub db_role: String,
pub priority: i32,
pub assign_roles: Vec<String>,
pub permissions: Vec<String>,
pub conditions: Vec<RoleMappingCondition>,
}
impl RoleMappingRule {
pub fn new(condition: RoleCondition, db_role: impl Into<String>) -> Self {
Self {
name: String::new(),
condition,
db_role: db_role.into(),
priority: 0,
assign_roles: Vec::new(),
permissions: Vec::new(),
conditions: Vec::new(),
}
}
pub fn with_priority(mut self, priority: i32) -> Self {
self.priority = priority;
self
}
}
#[derive(Debug, Clone)]
pub enum RoleCondition {
JwtClaim { name: String, value: String },
JwtClaimAny { name: String, values: Vec<String> },
OAuthScope(String),
Group(String),
EmailDomain(String),
TenantId(String),
And(Vec<RoleCondition>),
Or(Vec<RoleCondition>),
Always,
}
#[derive(Debug, Clone)]
pub enum RoleMappingCondition {
HasClaim { claim: String, value: Option<String> },
InGroup { group: String },
HasRole { role: String },
FromTenant { tenant_id: String },
AuthMethod { method: String },
EmailDomain { domain: String },
UsernamePattern { pattern: String },
And { conditions: Vec<RoleMappingCondition> },
Or { conditions: Vec<RoleMappingCondition> },
Not { condition: Box<RoleMappingCondition> },
}
impl RoleMappingCondition {
pub fn has_claim(claim: impl Into<String>, value: Option<String>) -> Self {
Self::HasClaim {
claim: claim.into(),
value,
}
}
pub fn in_group(group: impl Into<String>) -> Self {
Self::InGroup {
group: group.into(),
}
}
pub fn has_role(role: impl Into<String>) -> Self {
Self::HasRole {
role: role.into(),
}
}
pub fn auth_method(method: impl Into<String>) -> Self {
Self::AuthMethod {
method: method.into(),
}
}
}
impl RoleCondition {
pub fn jwt_claim(name: impl Into<String>, value: impl Into<String>) -> Self {
Self::JwtClaim {
name: name.into(),
value: value.into(),
}
}
pub fn group(name: impl Into<String>) -> Self {
Self::Group(name.into())
}
pub fn email_domain(domain: impl Into<String>) -> Self {
Self::EmailDomain(domain.into())
}
}
#[derive(Debug, Clone)]
pub struct CredentialConfig {
pub default_provider: CredentialProvider,
pub static_credentials: HashMap<String, Credentials>,
pub vault: Option<VaultConfig>,
pub aws_secrets: Option<AwsSecretsConfig>,
pub cache_ttl: Duration,
}
impl Default for CredentialConfig {
fn default() -> Self {
Self {
default_provider: CredentialProvider::Static,
static_credentials: HashMap::new(),
vault: None,
aws_secrets: None,
cache_ttl: Duration::from_secs(300),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CredentialProvider {
Static,
Vault,
AwsSecrets,
}
#[derive(Debug, Clone)]
pub struct Credentials {
pub username: String,
pub password: String,
pub ttl: Option<Duration>,
pub options: HashMap<String, String>,
}
impl Credentials {
pub fn new(username: impl Into<String>, password: impl Into<String>) -> Self {
Self {
username: username.into(),
password: password.into(),
ttl: None,
options: HashMap::new(),
}
}
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.ttl = Some(ttl);
self
}
}
#[derive(Debug, Clone)]
pub struct VaultConfig {
pub address: String,
pub auth_method: VaultAuthMethod,
pub role: String,
pub secret_path: String,
pub tls_verify: bool,
}
#[derive(Debug, Clone)]
pub enum VaultAuthMethod {
Token(String),
Kubernetes { role: String },
AppRole { role_id: String, secret_id: String },
}
#[derive(Debug, Clone)]
pub struct AwsSecretsConfig {
pub region: String,
pub secret_prefix: String,
pub use_iam_role: bool,
}
#[derive(Debug, Clone)]
pub struct SessionConfig {
pub timeout: Duration,
pub max_sessions_per_identity: usize,
pub max_sessions_per_user: usize,
pub idle_timeout: Duration,
pub absolute_timeout: Duration,
pub secure_cookies: bool,
pub session_vars: HashMap<String, String>,
pub extend_on_activity: bool,
}
impl Default for SessionConfig {
fn default() -> Self {
Self {
timeout: Duration::from_secs(3600),
max_sessions_per_identity: 10,
max_sessions_per_user: 10,
idle_timeout: Duration::from_secs(1800),
absolute_timeout: Duration::from_secs(86400),
secure_cookies: true,
session_vars: HashMap::new(),
extend_on_activity: true,
}
}
}
#[derive(Debug, Clone)]
pub struct AuthRateLimitConfig {
pub enabled: bool,
pub max_attempts_per_ip: u32,
pub max_failures_per_ip: u32,
pub lockout_duration: Duration,
pub window_seconds: u64,
pub max_requests_per_user: u32,
pub max_requests_per_ip: u32,
}
impl Default for AuthRateLimitConfig {
fn default() -> Self {
Self {
enabled: true,
max_attempts_per_ip: 60,
max_failures_per_ip: 10,
lockout_duration: Duration::from_secs(300),
window_seconds: 60,
max_requests_per_user: 120,
max_requests_per_ip: 60,
}
}
}
#[derive(Debug, Clone)]
pub enum AuthType {
Jwt(String),
OAuth(String),
Basic { username: String, password: String },
ApiKey(String),
None,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum AuthMethod {
Jwt,
OAuth,
Ldap,
ApiKey,
Basic,
Trust,
AgentToken,
Session,
Anonymous,
}
impl std::fmt::Display for AuthMethod {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Jwt => write!(f, "jwt"),
Self::OAuth => write!(f, "oauth"),
Self::Ldap => write!(f, "ldap"),
Self::ApiKey => write!(f, "api_key"),
Self::Basic => write!(f, "basic"),
Self::Trust => write!(f, "trust"),
Self::AgentToken => write!(f, "agent_token"),
Self::Session => write!(f, "session"),
Self::Anonymous => write!(f, "anonymous"),
}
}
}
impl AuthMethod {
pub fn from_str(s: &str) -> Option<Self> {
match s.to_lowercase().as_str() {
"jwt" => Some(Self::Jwt),
"oauth" => Some(Self::OAuth),
"ldap" => Some(Self::Ldap),
"api_key" | "apikey" => Some(Self::ApiKey),
"basic" => Some(Self::Basic),
"trust" => Some(Self::Trust),
"agent_token" | "agent" => Some(Self::AgentToken),
"session" => Some(Self::Session),
"anonymous" | "none" => Some(Self::Anonymous),
_ => None,
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Identity {
pub user_id: String,
pub name: Option<String>,
pub email: Option<String>,
pub roles: Vec<String>,
pub groups: Vec<String>,
pub tenant_id: Option<String>,
pub claims: HashMap<String, serde_json::Value>,
pub auth_method: String,
pub authenticated_at: chrono::DateTime<chrono::Utc>,
}
impl Identity {
pub fn new(user_id: impl Into<String>, auth_method: impl Into<String>) -> Self {
Self {
user_id: user_id.into(),
name: None,
email: None,
roles: Vec::new(),
groups: Vec::new(),
tenant_id: None,
claims: HashMap::new(),
auth_method: auth_method.into(),
authenticated_at: chrono::Utc::now(),
}
}
pub fn from_jwt_claims(claims: &JwtClaims) -> Self {
let mut identity = Self::new(&claims.sub, "jwt");
identity.name = claims.name.clone();
identity.email = claims.email.clone();
identity.roles = claims.roles.clone();
identity.tenant_id = claims.tenant_id.clone();
identity.claims = claims.custom.clone();
identity
}
pub fn has_role(&self, role: &str) -> bool {
self.roles.iter().any(|r| r == role)
}
pub fn in_group(&self, group: &str) -> bool {
self.groups.iter().any(|g| g == group)
}
pub fn is_admin(&self) -> bool {
self.has_role("admin") || self.has_role("db_admin")
}
pub fn get_claim(&self, name: &str) -> Option<&serde_json::Value> {
self.claims.get(name)
}
pub fn email_domain(&self) -> Option<&str> {
self.email.as_ref().and_then(|e| e.split('@').nth(1))
}
pub fn anonymous() -> Self {
Self {
user_id: "anonymous".to_string(),
name: None,
email: None,
roles: Vec::new(),
groups: Vec::new(),
tenant_id: None,
claims: HashMap::new(),
auth_method: "anonymous".to_string(),
authenticated_at: chrono::Utc::now(),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct JwtClaims {
pub sub: String,
pub iss: String,
pub aud: Option<Vec<String>>,
pub exp: i64,
pub iat: i64,
pub nbf: Option<i64>,
pub jti: Option<String>,
pub name: Option<String>,
pub email: Option<String>,
#[serde(default)]
pub roles: Vec<String>,
pub tenant_id: Option<String>,
#[serde(flatten)]
pub custom: HashMap<String, serde_json::Value>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentIdentity {
pub id: String,
pub agent_type: String,
pub allowed_tools: Vec<String>,
pub quota: AgentQuota,
pub conversation_id: Option<String>,
pub parent_identity: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AgentQuota {
pub max_queries_per_conversation: u32,
pub max_rows_per_query: u32,
pub token_budget: u64,
pub allowed_tables: Option<Vec<String>>,
}
impl Default for AgentQuota {
fn default() -> Self {
Self {
max_queries_per_conversation: 100,
max_rows_per_query: 1000,
token_budget: 100000,
allowed_tables: None,
}
}
}
#[derive(Debug, Clone)]
pub struct ToolPermission {
pub db_role: String,
pub allowed_tables: Vec<String>,
pub read_only: bool,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_auth_config_builder() {
let config = AuthConfig::builder()
.jwt(JwtConfig::new("https://auth.example.com/.well-known/jwks.json"))
.add_role_mapping(RoleMappingRule::new(
RoleCondition::jwt_claim("role", "admin"),
"db_admin",
).with_priority(100))
.default_role("db_minimal")
.build();
assert!(config.enabled);
assert!(config.jwt.is_some());
assert_eq!(config.role_mapping.len(), 1);
}
#[test]
fn test_identity() {
let identity = Identity::new("user123", "jwt");
assert_eq!(identity.user_id, "user123");
assert_eq!(identity.auth_method, "jwt");
assert!(!identity.is_admin());
}
#[test]
fn test_identity_roles() {
let mut identity = Identity::new("admin123", "jwt");
identity.roles = vec!["admin".to_string(), "db_readwrite".to_string()];
assert!(identity.is_admin());
assert!(identity.has_role("admin"));
assert!(identity.has_role("db_readwrite"));
assert!(!identity.has_role("superuser"));
}
#[test]
fn test_email_domain() {
let mut identity = Identity::new("user", "jwt");
identity.email = Some("alice@example.com".to_string());
assert_eq!(identity.email_domain(), Some("example.com"));
}
#[test]
fn test_credentials() {
let creds = Credentials::new("dbuser", "password123")
.with_ttl(Duration::from_secs(3600));
assert_eq!(creds.username, "dbuser");
assert!(creds.ttl.is_some());
}
#[test]
fn test_role_mapping() {
let rule = RoleMappingRule::new(
RoleCondition::group("developers"),
"db_readwrite",
).with_priority(50);
assert_eq!(rule.db_role, "db_readwrite");
assert_eq!(rule.priority, 50);
}
}