astraea_server/
auth.rs

1//! Authentication and Role-Based Access Control (RBAC) for AstraeaDB.
2//!
3//! Supports API key authentication with three roles:
4//! - `Admin`: full access to all operations
5//! - `Writer`: read + write operations (no admin)
6//! - `Reader`: read-only operations
7
8use std::collections::HashMap;
9use std::sync::RwLock;
10use std::time::SystemTime;
11
12use serde::{Deserialize, Serialize};
13
14/// User roles with increasing privilege levels.
15#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
16pub enum Role {
17    /// Read-only access: get, query, search, traverse.
18    Reader,
19    /// Read + write: create, update, delete nodes/edges.
20    Writer,
21    /// Full access: all operations including server management.
22    Admin,
23}
24
25impl std::fmt::Display for Role {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            Role::Reader => write!(f, "reader"),
29            Role::Writer => write!(f, "writer"),
30            Role::Admin => write!(f, "admin"),
31        }
32    }
33}
34
35/// An API key entry with associated metadata.
36#[derive(Debug, Clone, Serialize, Deserialize)]
37pub struct ApiKeyEntry {
38    /// The API key string.
39    pub key: String,
40    /// The role assigned to this key.
41    pub role: Role,
42    /// Human-readable description (e.g., "CI pipeline key").
43    pub description: String,
44    /// Whether this key is currently active.
45    pub active: bool,
46}
47
48/// An entry in the audit log.
49#[derive(Debug, Clone, Serialize, Deserialize)]
50pub struct AuditEntry {
51    /// Timestamp of the operation (epoch seconds).
52    pub timestamp: u64,
53    /// The API key used (truncated for security).
54    pub api_key_prefix: String,
55    /// The role of the authenticated user.
56    pub role: Role,
57    /// The operation type (e.g., "CreateNode", "Query").
58    pub operation: String,
59    /// Whether the operation was allowed.
60    pub allowed: bool,
61}
62
63/// Authentication and authorization manager.
64pub struct AuthManager {
65    /// Map from API key string to entry.
66    keys: RwLock<HashMap<String, ApiKeyEntry>>,
67    /// Whether authentication is enabled. If false, all requests are allowed.
68    enabled: bool,
69    /// Audit log (bounded circular buffer).
70    audit_log: RwLock<Vec<AuditEntry>>,
71    /// Maximum audit log entries before truncation.
72    max_audit_entries: usize,
73}
74
75impl AuthManager {
76    /// Create a new auth manager with authentication disabled.
77    pub fn disabled() -> Self {
78        Self {
79            keys: RwLock::new(HashMap::new()),
80            enabled: false,
81            audit_log: RwLock::new(Vec::new()),
82            max_audit_entries: 10000,
83        }
84    }
85
86    /// Create a new auth manager with authentication enabled.
87    pub fn new(keys: Vec<ApiKeyEntry>) -> Self {
88        let key_map: HashMap<String, ApiKeyEntry> =
89            keys.into_iter().map(|k| (k.key.clone(), k)).collect();
90        Self {
91            keys: RwLock::new(key_map),
92            enabled: true,
93            audit_log: RwLock::new(Vec::new()),
94            max_audit_entries: 10000,
95        }
96    }
97
98    /// Check if authentication is enabled.
99    pub fn is_enabled(&self) -> bool {
100        self.enabled
101    }
102
103    /// Authenticate an API key. Returns the role if valid.
104    pub fn authenticate(&self, api_key: &str) -> Option<Role> {
105        if !self.enabled {
106            return Some(Role::Admin); // no auth = full access
107        }
108
109        let keys = self.keys.read().unwrap();
110        keys.get(api_key)
111            .filter(|entry| entry.active)
112            .map(|entry| entry.role)
113    }
114
115    /// Check if a role is authorized for a given operation.
116    pub fn authorize(role: Role, operation: &str) -> bool {
117        match role {
118            Role::Admin => true,
119            Role::Writer => !Self::is_admin_operation(operation),
120            Role::Reader => Self::is_read_operation(operation),
121        }
122    }
123
124    /// Check if an operation is read-only.
125    fn is_read_operation(operation: &str) -> bool {
126        matches!(
127            operation,
128            "GetNode"
129                | "GetEdge"
130                | "Neighbors"
131                | "NeighborsAt"
132                | "Bfs"
133                | "BfsAt"
134                | "ShortestPath"
135                | "ShortestPathAt"
136                | "VectorSearch"
137                | "HybridSearch"
138                | "SemanticNeighbors"
139                | "SemanticWalk"
140                | "Query"
141                | "ExtractSubgraph"
142                | "GraphRag"
143                | "Ping"
144        )
145    }
146
147    /// Check if an operation requires admin role.
148    fn is_admin_operation(_operation: &str) -> bool {
149        // Currently no admin-only operations beyond normal CRUD.
150        // This is a hook for future server management commands.
151        false
152    }
153
154    /// Record an operation in the audit log.
155    pub fn audit(&self, api_key: &str, role: Role, operation: &str, allowed: bool) {
156        let entry = AuditEntry {
157            timestamp: SystemTime::now()
158                .duration_since(SystemTime::UNIX_EPOCH)
159                .unwrap_or_default()
160                .as_secs(),
161            api_key_prefix: if api_key.len() >= 8 {
162                format!("{}...", &api_key[..8])
163            } else {
164                api_key.to_string()
165            },
166            role,
167            operation: operation.to_string(),
168            allowed,
169        };
170
171        let mut log = self.audit_log.write().unwrap();
172        log.push(entry);
173        if log.len() > self.max_audit_entries {
174            // Remove oldest 10% to avoid constant shifting.
175            let drain_count = self.max_audit_entries / 10;
176            log.drain(..drain_count);
177        }
178    }
179
180    /// Get recent audit log entries.
181    pub fn recent_audit(&self, count: usize) -> Vec<AuditEntry> {
182        let log = self.audit_log.read().unwrap();
183        log.iter().rev().take(count).cloned().collect()
184    }
185
186    /// Add a new API key.
187    pub fn add_key(&self, entry: ApiKeyEntry) {
188        let mut keys = self.keys.write().unwrap();
189        keys.insert(entry.key.clone(), entry);
190    }
191
192    /// Revoke (deactivate) an API key.
193    pub fn revoke_key(&self, api_key: &str) -> bool {
194        let mut keys = self.keys.write().unwrap();
195        if let Some(entry) = keys.get_mut(api_key) {
196            entry.active = false;
197            true
198        } else {
199            false
200        }
201    }
202
203    /// Get the operation name from a request type string.
204    pub fn operation_name(request_json: &str) -> &str {
205        // Quick extraction of the "type" field from JSON without full parsing.
206        if let Some(start) = request_json.find("\"type\":\"") {
207            let rest = &request_json[start + 8..];
208            if let Some(end) = rest.find('"') {
209                return &rest[..end];
210            }
211        }
212        "Unknown"
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219
220    fn make_keys() -> Vec<ApiKeyEntry> {
221        vec![
222            ApiKeyEntry {
223                key: "admin-key-12345678".into(),
224                role: Role::Admin,
225                description: "Admin key".into(),
226                active: true,
227            },
228            ApiKeyEntry {
229                key: "writer-key-12345678".into(),
230                role: Role::Writer,
231                description: "Writer key".into(),
232                active: true,
233            },
234            ApiKeyEntry {
235                key: "reader-key-12345678".into(),
236                role: Role::Reader,
237                description: "Reader key".into(),
238                active: true,
239            },
240            ApiKeyEntry {
241                key: "inactive-key-12345678".into(),
242                role: Role::Admin,
243                description: "Inactive key".into(),
244                active: false,
245            },
246        ]
247    }
248
249    #[test]
250    fn disabled_auth_allows_all() {
251        let auth = AuthManager::disabled();
252        assert!(!auth.is_enabled());
253        assert_eq!(auth.authenticate("anything"), Some(Role::Admin));
254    }
255
256    #[test]
257    fn valid_key_returns_role() {
258        let auth = AuthManager::new(make_keys());
259        assert_eq!(auth.authenticate("admin-key-12345678"), Some(Role::Admin));
260        assert_eq!(auth.authenticate("writer-key-12345678"), Some(Role::Writer));
261        assert_eq!(auth.authenticate("reader-key-12345678"), Some(Role::Reader));
262    }
263
264    #[test]
265    fn invalid_key_returns_none() {
266        let auth = AuthManager::new(make_keys());
267        assert_eq!(auth.authenticate("bad-key"), None);
268    }
269
270    #[test]
271    fn inactive_key_returns_none() {
272        let auth = AuthManager::new(make_keys());
273        assert_eq!(auth.authenticate("inactive-key-12345678"), None);
274    }
275
276    #[test]
277    fn admin_can_do_everything() {
278        assert!(AuthManager::authorize(Role::Admin, "CreateNode"));
279        assert!(AuthManager::authorize(Role::Admin, "DeleteNode"));
280        assert!(AuthManager::authorize(Role::Admin, "GetNode"));
281        assert!(AuthManager::authorize(Role::Admin, "Ping"));
282    }
283
284    #[test]
285    fn writer_can_read_and_write() {
286        assert!(AuthManager::authorize(Role::Writer, "CreateNode"));
287        assert!(AuthManager::authorize(Role::Writer, "DeleteNode"));
288        assert!(AuthManager::authorize(Role::Writer, "GetNode"));
289        assert!(AuthManager::authorize(Role::Writer, "Query"));
290    }
291
292    #[test]
293    fn reader_cannot_write() {
294        assert!(!AuthManager::authorize(Role::Reader, "CreateNode"));
295        assert!(!AuthManager::authorize(Role::Reader, "DeleteNode"));
296        assert!(!AuthManager::authorize(Role::Reader, "UpdateNode"));
297        assert!(AuthManager::authorize(Role::Reader, "GetNode"));
298        assert!(AuthManager::authorize(Role::Reader, "Query"));
299        assert!(AuthManager::authorize(Role::Reader, "VectorSearch"));
300        assert!(AuthManager::authorize(Role::Reader, "Ping"));
301    }
302
303    #[test]
304    fn audit_log_records_entries() {
305        let auth = AuthManager::new(make_keys());
306        auth.audit("admin-key-12345678", Role::Admin, "CreateNode", true);
307        auth.audit("reader-key-12345678", Role::Reader, "GetNode", true);
308        auth.audit("reader-key-12345678", Role::Reader, "CreateNode", false);
309
310        let recent = auth.recent_audit(10);
311        assert_eq!(recent.len(), 3);
312        assert!(!recent[0].allowed); // most recent first
313        assert!(recent[1].allowed);
314    }
315
316    #[test]
317    fn revoke_key_prevents_auth() {
318        let auth = AuthManager::new(make_keys());
319        assert_eq!(auth.authenticate("writer-key-12345678"), Some(Role::Writer));
320        assert!(auth.revoke_key("writer-key-12345678"));
321        assert_eq!(auth.authenticate("writer-key-12345678"), None);
322    }
323
324    #[test]
325    fn add_key_works() {
326        let auth = AuthManager::new(vec![]);
327        assert_eq!(auth.authenticate("new-key"), None);
328        auth.add_key(ApiKeyEntry {
329            key: "new-key".into(),
330            role: Role::Writer,
331            description: "test".into(),
332            active: true,
333        });
334        assert_eq!(auth.authenticate("new-key"), Some(Role::Writer));
335    }
336
337    #[test]
338    fn role_ordering() {
339        assert!(Role::Reader < Role::Writer);
340        assert!(Role::Writer < Role::Admin);
341    }
342}
astraea_server/auth.rs

astraea_server/
auth.rs
