llm_orchestrator_auth/
lib.rs

1//! # LLM Orchestrator Authentication & Authorization
2//!
3//! This crate provides comprehensive authentication and authorization for the LLM Orchestrator.
4//!
5//! ## Features
6//!
7//! - **JWT Authentication**: Stateless token-based authentication with short-lived access tokens
8//!   and long-lived refresh tokens
9//! - **API Key Management**: Secure API key generation, hashing, and validation
10//! - **Role-Based Access Control (RBAC)**: Fine-grained permission system with predefined roles
11//! - **Auth Middleware**: Ready-to-use middleware for authenticating requests
12//!
13//! ## Quick Start
14//!
15//! ```rust,no_run
16//! use llm_orchestrator_auth::*;
17//! use std::sync::Arc;
18//!
19//! #[tokio::main]
20//! async fn main() -> Result<(), Box<dyn std::error::Error>> {
21//!     // Create JWT auth
22//!     let jwt_auth = Arc::new(JwtAuth::new(b"your-secret-key-at-least-32-bytes".to_vec()));
23//!
24//!     // Create API key manager
25//!     let api_key_store = Arc::new(InMemoryApiKeyStore::new());
26//!     let api_key_manager = Arc::new(ApiKeyManager::new(api_key_store));
27//!
28//!     // Create RBAC engine
29//!     let rbac = Arc::new(RbacEngine::new());
30//!
31//!     // Create auth middleware
32//!     let auth = AuthMiddleware::new(jwt_auth.clone(), api_key_manager.clone(), rbac.clone());
33//!
34//!     // Generate a JWT token
35//!     let token = jwt_auth.generate_token("user123", vec!["developer".to_string()])?;
36//!     println!("JWT Token: {}", token);
37//!
38//!     // Authenticate a request
39//!     let auth_header = format!("Bearer {}", token);
40//!     let ctx = auth.authenticate(Some(&auth_header)).await?;
41//!     println!("Authenticated user: {}", ctx.user_id);
42//!
43//!     // Check permissions
44//!     ctx.require_permission(&Permission::WorkflowExecute)?;
45//!     println!("User has permission to execute workflows");
46//!
47//!     Ok(())
48//! }
49//! ```
50//!
51//! ## Predefined Roles
52//!
53//! - **viewer**: Read-only access to workflows and executions
54//! - **executor**: Can read and execute workflows
55//! - **developer**: Full access to workflows, can create/update/delete
56//! - **admin**: Full administrative access to all resources
57//!
58//! ## Security Features
59//!
60//! - JWT tokens expire after 15 minutes by default
61//! - Refresh tokens expire after 7 days by default
62//! - API keys are hashed using SHA-256 before storage
63//! - Cryptographically secure random key generation
64//! - Token expiration validation
65//! - Permission-based authorization
66
67pub mod api_keys;
68pub mod jwt;
69pub mod middleware;
70pub mod models;
71pub mod rbac;
72
73// Re-export main types for convenience
74pub use api_keys::{ApiKeyManager, ApiKeyStore, InMemoryApiKeyStore};
75pub use jwt::JwtAuth;
76pub use middleware::AuthMiddleware;
77pub use models::{
78    ApiKey, ApiKeyInfo, AuthContext, AuthError, AuthResult, AuthType, Claims, Permission,
79    RolePolicy,
80};
81pub use rbac::RbacEngine;
82
83/// Version information
84pub const VERSION: &str = env!("CARGO_PKG_VERSION");
85
86#[cfg(test)]
87mod integration_tests {
88    use super::*;
89    use std::sync::Arc;
90
91    #[tokio::test]
92    async fn test_full_jwt_flow() {
93        // Setup
94        let jwt_auth = Arc::new(JwtAuth::new(
95            b"test-secret-key-at-least-32-bytes-long".to_vec(),
96        ));
97        let api_key_store = Arc::new(InMemoryApiKeyStore::new());
98        let api_key_manager = Arc::new(ApiKeyManager::new(api_key_store));
99        let rbac = Arc::new(RbacEngine::new());
100        let auth = AuthMiddleware::new(jwt_auth.clone(), api_key_manager, rbac);
101
102        // Generate token
103        let token = jwt_auth
104            .generate_token("user123", vec!["developer".to_string()])
105            .unwrap();
106
107        // Authenticate
108        let auth_header = format!("Bearer {}", token);
109        let ctx = auth.authenticate(Some(&auth_header)).await.unwrap();
110
111        // Verify context
112        assert_eq!(ctx.user_id, "user123");
113        assert_eq!(ctx.roles, vec!["developer"]);
114        assert!(ctx.has_permission(&Permission::WorkflowWrite));
115
116        // Check authorization
117        assert!(auth
118            .authorize(&ctx, &Permission::WorkflowWrite)
119            .is_ok());
120    }
121
122    #[tokio::test]
123    async fn test_full_api_key_flow() {
124        // Setup
125        let jwt_auth = Arc::new(JwtAuth::new(
126            b"test-secret-key-at-least-32-bytes-long".to_vec(),
127        ));
128        let api_key_store = Arc::new(InMemoryApiKeyStore::new());
129        let api_key_manager = Arc::new(ApiKeyManager::new(api_key_store));
130        let rbac = Arc::new(RbacEngine::new());
131        let auth = AuthMiddleware::new(jwt_auth, api_key_manager.clone(), rbac);
132
133        // Create API key
134        let api_key = api_key_manager
135            .create_key(
136                "user456",
137                vec!["workflow:read".to_string(), "workflow:execute".to_string()],
138                Some("My Test Key".to_string()),
139                Some(30),
140            )
141            .await
142            .unwrap();
143
144        // Authenticate with API key
145        let auth_header = format!("ApiKey {}", api_key.key);
146        let ctx = auth.authenticate(Some(&auth_header)).await.unwrap();
147
148        // Verify context
149        assert_eq!(ctx.user_id, "user456");
150        assert!(ctx.has_permission(&Permission::WorkflowRead));
151        assert!(ctx.has_permission(&Permission::WorkflowExecute));
152        assert!(!ctx.has_permission(&Permission::WorkflowWrite));
153
154        // List user's keys
155        let keys = api_key_manager.list_keys("user456").await.unwrap();
156        assert_eq!(keys.len(), 1);
157        assert_eq!(keys[0].name, Some("My Test Key".to_string()));
158
159        // Revoke key
160        api_key_manager.revoke_key(&api_key.id).await.unwrap();
161
162        // Verify key is revoked
163        let result = auth.authenticate(Some(&auth_header)).await;
164        assert!(matches!(result, Err(AuthError::ApiKeyNotFound)));
165    }
166
167    #[tokio::test]
168    async fn test_refresh_token_flow() {
169        // Setup
170        let jwt_auth = JwtAuth::new(b"test-secret-key-at-least-32-bytes-long".to_vec());
171
172        // Generate refresh token
173        let refresh_token = jwt_auth.generate_refresh_token("user789").unwrap();
174
175        // Verify refresh token
176        let user_id = jwt_auth.verify_refresh_token(&refresh_token).unwrap();
177        assert_eq!(user_id, "user789");
178
179        // Use refresh token to get new access token
180        let access_token = jwt_auth
181            .refresh_access_token(&refresh_token, vec!["executor".to_string()])
182            .unwrap();
183
184        // Verify new access token
185        let claims = jwt_auth.verify_token(&access_token).unwrap();
186        assert_eq!(claims.sub, "user789");
187        assert_eq!(claims.roles, vec!["executor"]);
188    }
189
190    #[tokio::test]
191    async fn test_rbac_permission_checks() {
192        let rbac = RbacEngine::new();
193
194        // Test viewer permissions
195        assert!(rbac.check_permission(&["viewer".to_string()], &Permission::WorkflowRead));
196        assert!(!rbac.check_permission(&["viewer".to_string()], &Permission::WorkflowWrite));
197
198        // Test executor permissions
199        assert!(rbac.check_permission(
200            &["executor".to_string()],
201            &Permission::WorkflowExecute
202        ));
203
204        // Test developer permissions
205        assert!(rbac.check_permission(
206            &["developer".to_string()],
207            &Permission::WorkflowWrite
208        ));
209
210        // Test admin permissions
211        for permission in Permission::all() {
212            assert!(rbac.check_permission(&["admin".to_string()], &permission));
213        }
214    }
215
216    #[tokio::test]
217    async fn test_multiple_roles() {
218        let jwt_auth = Arc::new(JwtAuth::new(
219            b"test-secret-key-at-least-32-bytes-long".to_vec(),
220        ));
221        let rbac = Arc::new(RbacEngine::new());
222
223        // Create token with multiple roles
224        let token = jwt_auth
225            .generate_token(
226                "user123",
227                vec!["viewer".to_string(), "executor".to_string()],
228            )
229            .unwrap();
230
231        let claims = jwt_auth.verify_token(&token).unwrap();
232
233        // Compute combined permissions
234        let permissions = rbac.compute_permissions(&claims.roles);
235
236        // Should have union of both roles
237        assert!(permissions.contains(&Permission::WorkflowRead));
238        assert!(permissions.contains(&Permission::WorkflowExecute));
239        assert!(permissions.contains(&Permission::ExecutionRead));
240    }
241}