authkestra_engine/
engine.rs

1use crate::auth::session::{Session, SessionConfig, SessionStore};
2use crate::auth::{AuthError, AuthInput, AuthMethod, AuthResult, ErasedOAuthFlow, Identity};
3#[cfg(feature = "token")]
4use crate::token::TokenManager;
5use std::collections::HashMap;
6use std::sync::Arc;
7
8/// Marker for a missing component in the typestate pattern.
9#[derive(Clone, Default, Debug)]
10pub struct Missing;
11
12/// Marker for a configured component in the typestate pattern.
13#[derive(Clone, Debug)]
14pub struct Configured<T>(pub T);
15
16/// Trait for the session store state in the `Engine`.
17pub trait SessionStoreState: Send + Sync + Clone {
18    /// Returns the session store if configured.
19    fn get_store(&self) -> Arc<dyn SessionStore>;
20}
21
22impl SessionStoreState for Configured<Arc<dyn SessionStore>> {
23    fn get_store(&self) -> Arc<dyn SessionStore> {
24        self.0.clone()
25    }
26}
27
28/// Trait for the token manager state in the `Engine`.
29pub trait TokenManagerState: Send + Sync + Clone {
30    /// Returns the token manager if configured.
31    #[cfg(feature = "token")]
32    fn get_manager(&self) -> Arc<TokenManager>;
33}
34
35#[cfg(feature = "token")]
36impl TokenManagerState for Configured<Arc<TokenManager>> {
37    fn get_manager(&self) -> Arc<TokenManager> {
38        self.0.clone()
39    }
40}
41
42/// The central orchestrator for Engine.
43///
44/// `Engine` ties together authentication methods, session management, and flows.
45/// It is constructed using the [`EngineBuilder`] which uses the Typestate pattern
46/// to ensure that certain methods are only available when the necessary components are configured.
47pub struct Engine<S = Missing, T = Missing> {
48    /// Map of registered OAuth providers.
49    pub providers: HashMap<String, Arc<dyn ErasedOAuthFlow>>,
50    /// Map of registered local authentication methods.
51    pub auth_methods: HashMap<String, Arc<dyn AuthMethod>>,
52    /// Map of methods explicitly registered for step-up (MFA) use.
53    pub mfa_methods: HashMap<String, Arc<dyn AuthMethod>>,
54    /// Internal secret for signing temporary MFA JWT tokens.
55    pub mfa_jwt_secret: [u8; 32],
56    /// The session storage backend.
57    pub session_store: S,
58    /// Configuration for session cookies.
59    pub session_config: SessionConfig,
60    /// Manager for JWT signing and verification.
61    #[cfg(feature = "token")]
62    pub token_manager: T,
63}
64
65impl<S, T> Clone for Engine<S, T>
66where
67    S: Clone,
68    T: Clone,
69{
70    fn clone(&self) -> Self {
71        Self {
72            providers: self.providers.clone(),
73            auth_methods: self.auth_methods.clone(),
74            mfa_methods: self.mfa_methods.clone(),
75            mfa_jwt_secret: self.mfa_jwt_secret,
76            session_store: self.session_store.clone(),
77            session_config: self.session_config.clone(),
78            #[cfg(feature = "token")]
79            token_manager: self.token_manager.clone(),
80        }
81    }
82}
83
84impl Engine<Missing, Missing> {
85    /// Start building a new `Engine`.
86    pub fn builder() -> EngineBuilder<Missing, Missing> {
87        let mut secret = [0u8; 32];
88        rand::RngCore::fill_bytes(&mut rand::rng(), &mut secret);
89
90        EngineBuilder {
91            providers: HashMap::new(),
92            auth_methods: HashMap::new(),
93            mfa_methods: HashMap::new(),
94            mfa_jwt_secret: secret,
95            session_store: Missing,
96            session_config: SessionConfig::default(),
97            #[cfg(feature = "token")]
98            token_manager: Missing,
99        }
100    }
101}
102
103/// A builder for configuring and creating an [`Engine`] instance.
104pub struct EngineBuilder<S = Missing, T = Missing> {
105    providers: HashMap<String, Arc<dyn ErasedOAuthFlow>>,
106    auth_methods: HashMap<String, Arc<dyn AuthMethod>>,
107    mfa_methods: HashMap<String, Arc<dyn AuthMethod>>,
108    mfa_jwt_secret: [u8; 32],
109    session_store: S,
110    session_config: SessionConfig,
111    #[cfg(feature = "token")]
112    token_manager: T,
113}
114
115impl<S, T> EngineBuilder<S, T> {
116    /// Register an OAuth provider flow.
117    pub fn provider<F>(mut self, flow: F) -> Self
118    where
119        F: ErasedOAuthFlow + 'static,
120    {
121        let id = flow.provider_id();
122        self.providers.insert(id, Arc::new(flow));
123        self
124    }
125
126    /// Register a local authentication method.
127    pub fn with_auth_method<M>(mut self, method: M) -> Self
128    where
129        M: AuthMethod + 'static,
130    {
131        self.auth_methods
132            .insert(method.name().to_string(), Arc::new(method));
133        self
134    }
135
136    /// Register a local authentication method to be used EXCLUSIVELY for step-up MFA challenges.
137    pub fn with_mfa_method<M>(mut self, method: M) -> Self
138    where
139        M: AuthMethod + 'static,
140    {
141        self.mfa_methods
142            .insert(method.name().to_string(), Arc::new(method));
143        self
144    }
145
146    /// Register the TOTP authentication method.
147    #[cfg(feature = "totp")]
148    pub fn with_totp<C>(self, store: C) -> Self
149    where
150        C: crate::CredentialStore + 'static,
151    {
152        self.with_auth_method(crate::auth::totp::TotpAuthMethod::new(store))
153    }
154
155    /// Register the WebAuthn authentication method.
156    #[cfg(feature = "webauthn")]
157    pub fn with_webauthn<C>(self, webauthn: Arc<webauthn_rs::prelude::Webauthn>, store: C) -> Self
158    where
159        C: crate::CredentialStore + 'static,
160    {
161        self.with_auth_method(crate::auth::webauthn::WebAuthnAuthMethod::new(
162            webauthn, store,
163        ))
164    }
165
166    /// Set the session store.
167    pub fn session_store(
168        self,
169        store: Arc<dyn SessionStore>,
170    ) -> EngineBuilder<Configured<Arc<dyn SessionStore>>, T> {
171        EngineBuilder {
172            providers: self.providers,
173            auth_methods: self.auth_methods,
174            mfa_methods: self.mfa_methods,
175            mfa_jwt_secret: self.mfa_jwt_secret,
176            session_store: Configured(store),
177            session_config: self.session_config,
178            #[cfg(feature = "token")]
179            token_manager: self.token_manager,
180        }
181    }
182
183    /// Set the token manager.
184    #[cfg(feature = "token")]
185    pub fn token_manager(
186        self,
187        manager: Arc<TokenManager>,
188    ) -> EngineBuilder<S, Configured<Arc<TokenManager>>> {
189        EngineBuilder {
190            providers: self.providers,
191            auth_methods: self.auth_methods,
192            mfa_methods: self.mfa_methods,
193            mfa_jwt_secret: self.mfa_jwt_secret,
194            session_store: self.session_store,
195            session_config: self.session_config,
196            token_manager: Configured(manager),
197        }
198    }
199
200    /// Set the JWT secret for the default token manager.
201    #[cfg(feature = "token")]
202    pub fn jwt_secret(self, secret: &[u8]) -> EngineBuilder<S, Configured<Arc<TokenManager>>> {
203        self.token_manager(Arc::new(TokenManager::new(secret, None)))
204    }
205
206    /// Set the session configuration.
207    pub fn session_config(mut self, config: SessionConfig) -> Self {
208        self.session_config = config;
209        self
210    }
211
212    /// Build the `Engine`.
213    pub fn build(self) -> Engine<S, T> {
214        Engine {
215            providers: self.providers,
216            auth_methods: self.auth_methods,
217            mfa_methods: self.mfa_methods,
218            mfa_jwt_secret: self.mfa_jwt_secret,
219            session_store: self.session_store,
220            session_config: self.session_config,
221            #[cfg(feature = "token")]
222            token_manager: self.token_manager,
223        }
224    }
225}
226
227impl<S, T> Engine<S, T> {
228    /// Attempt to authenticate a user.
229    /// Returns `AuthResult::Success` if authentication is fully complete,
230    /// or `AuthResult::MfaRequired` if a second factor is needed.
231    pub async fn authenticate(&self, input: AuthInput) -> Result<AuthResult, AuthError> {
232        // Handle MFA Challenge Continuation
233        if let AuthInput::MfaChallenge {
234            mfa_token,
235            challenge_input,
236        } = input
237        {
238            // Verify MFA Token
239            let token_data = jsonwebtoken::decode::<crate::auth::state::MfaTokenClaims>(
240                &mfa_token,
241                &jsonwebtoken::DecodingKey::from_secret(&self.mfa_jwt_secret),
242                &jsonwebtoken::Validation::new(jsonwebtoken::Algorithm::HS256),
243            )
244            .map_err(|_| AuthError::InvalidInput)?;
245
246            if !token_data.claims.mfa_pending {
247                return Err(AuthError::InvalidInput);
248            }
249
250            let method_name = match &*challenge_input {
251                #[cfg(feature = "totp")]
252                AuthInput::Totp { .. } => "totp",
253                #[cfg(feature = "webauthn")]
254                AuthInput::WebAuthnAuthentication { .. } => "webauthn",
255                _ => "",
256            };
257
258            if method_name.is_empty() {
259                return Err(AuthError::InvalidInput);
260            }
261
262            let method = self
263                .auth_methods
264                .get(method_name)
265                .or_else(|| self.mfa_methods.get(method_name))
266                .ok_or_else(|| {
267                    AuthError::Internal(format!("MFA method {} not registered", method_name))
268                })?;
269
270            let identity = method.authenticate(*challenge_input).await?;
271
272            if identity.external_id != token_data.claims.sub {
273                return Err(AuthError::Credentials("MFA token user mismatch".into()));
274            }
275
276            return Ok(AuthResult::Success(identity));
277        }
278
279        // Primary Authentication
280        let method_name = match &input {
281            AuthInput::Password { .. } => "password",
282            #[cfg(feature = "totp")]
283            AuthInput::Totp { .. } => "totp", // Could theoretically be used as primary
284            #[cfg(feature = "webauthn")]
285            AuthInput::WebAuthnAuthentication { .. } => "webauthn",
286            _ => "",
287        };
288
289        if method_name.is_empty() {
290            return Err(AuthError::InvalidInput);
291        }
292
293        let method = self.auth_methods.get(method_name).ok_or_else(|| {
294            AuthError::Internal(format!(
295                "Primary auth method {} not registered or is step-up only",
296                method_name
297            ))
298        })?;
299
300        let identity = method.authenticate(input).await?;
301
302        // Check if user has MFA enrolled
303        let mut enrolled_methods = Vec::new();
304        for (name, m) in self.auth_methods.iter().chain(self.mfa_methods.iter()) {
305            if name == "password" || name == method_name {
306                continue;
307            }
308            if !enrolled_methods.contains(name)
309                && m.has_enrolled(&identity.external_id).await.unwrap_or(false)
310            {
311                enrolled_methods.push(name.clone());
312            }
313        }
314
315        // If this method was already an MFA method (e.g. WebAuthn primary), we don't prompt for MFA again.
316        // Or if the user has no other MFA methods enrolled.
317        if enrolled_methods.is_empty() || method.is_mfa_equivalent() {
318            Ok(AuthResult::Success(identity))
319        } else {
320            // Issue MFA Token
321            let exp = chrono::Utc::now() + chrono::Duration::minutes(15);
322            let claims = crate::auth::state::MfaTokenClaims {
323                sub: identity.external_id.clone(),
324                mfa_pending: true,
325                exp: exp.timestamp() as usize,
326            };
327
328            let mfa_token = jsonwebtoken::encode(
329                &jsonwebtoken::Header::default(),
330                &claims,
331                &jsonwebtoken::EncodingKey::from_secret(&self.mfa_jwt_secret),
332            )
333            .map_err(|e| AuthError::Internal(e.to_string()))?;
334
335            Ok(AuthResult::MfaRequired {
336                mfa_token,
337                user_id: identity.external_id,
338                allowed_methods: enrolled_methods,
339            })
340        }
341    }
342
343    /// Starts a WebAuthn authentication ceremony for a list of enrolled passkeys.
344    #[cfg(feature = "webauthn")]
345    pub fn start_webauthn(
346        &self,
347        passkeys: &[webauthn_rs::prelude::Passkey],
348    ) -> Result<
349        (
350            webauthn_rs::prelude::RequestChallengeResponse,
351            webauthn_rs::prelude::PasskeyAuthentication,
352        ),
353        AuthError,
354    > {
355        let method = self
356            .auth_methods
357            .get("webauthn")
358            .ok_or_else(|| AuthError::Internal("WebAuthn method not registered".into()))?;
359
360        let webauthn_starter = method.as_webauthn_starter().ok_or_else(|| {
361            AuthError::Internal("WebAuthn method doesn't implement starter".into())
362        })?;
363
364        webauthn_starter.start_authentication(passkeys)
365    }
366}
367
368// Methods available only when a session store is present
369impl<T> Engine<Configured<Arc<dyn SessionStore>>, T> {
370    /// Get the session store.
371    pub fn session_store(&self) -> Arc<dyn SessionStore> {
372        self.session_store.0.clone()
373    }
374
375    /// Create a new session for the given identity.
376    #[tracing::instrument(skip(self, identity), fields(user_id = %identity.external_id))]
377    pub async fn create_session(&self, identity: Identity) -> Result<Session, AuthError> {
378        let session_duration = self
379            .session_config
380            .max_age
381            .unwrap_or(chrono::Duration::hours(24));
382        let session = Session {
383            id: uuid::Uuid::new_v4().to_string(),
384            identity,
385            expires_at: chrono::Utc::now() + session_duration,
386        };
387
388        tracing::debug!(session_id = %session.id, "creating new session");
389
390        self.session_store
391            .0
392            .save_session(&session)
393            .await
394            .map_err(|e| {
395                tracing::error!(error = %e, "failed to save session");
396                AuthError::Session(e.to_string())
397            })?;
398
399        tracing::info!(session_id = %session.id, "session created successfully");
400        Ok(session)
401    }
402}
403
404#[cfg(feature = "token")]
405impl<S> Engine<S, Configured<Arc<TokenManager>>> {
406    /// Get the token manager.
407    pub fn token_manager(&self) -> Arc<TokenManager> {
408        self.token_manager.0.clone()
409    }
410
411    /// Issue a JWT for the given identity.
412    #[tracing::instrument(skip(self, identity), fields(user_id = %identity.external_id))]
413    pub fn issue_token(
414        &self,
415        identity: Identity,
416        expires_in_secs: u64,
417    ) -> Result<String, AuthError> {
418        tracing::debug!("issuing token for user");
419        self.token_manager
420            .0
421            .issue_user_token(identity, expires_in_secs, None, None)
422            .map_err(|e| {
423                tracing::error!(error = %e, "failed to issue token");
424                AuthError::Token(e.to_string())
425            })
426            .inspect(|_| {
427                tracing::info!("token issued successfully");
428            })
429    }
430}
431
432/// Trait for Engine instances that have a session store configured.
433pub trait HasSessionStore {
434    /// Returns the session store.
435    fn session_store(&self) -> Arc<dyn SessionStore>;
436}
437
438impl<T> HasSessionStore for Engine<Configured<Arc<dyn SessionStore>>, T> {
439    fn session_store(&self) -> Arc<dyn SessionStore> {
440        self.session_store.0.clone()
441    }
442}
443
444/// Trait for Engine instances that have a token manager configured.
445#[cfg(feature = "token")]
446pub trait HasTokenManager {
447    /// Returns the token manager.
448    fn token_manager(&self) -> Arc<TokenManager>;
449}
450
451#[cfg(feature = "token")]
452impl<S> HasTokenManager for Engine<S, Configured<Arc<TokenManager>>> {
453    fn token_manager(&self) -> Arc<TokenManager> {
454        self.token_manager.0.clone()
455    }
456}
authkestra_engine/engine.rs

authkestra_engine/
engine.rs
