axum_firebase_middleware/
lib.rs

1//! # Firebase JWT Authentication Middleware for Axum
2//!
3//! A production-ready Firebase JWT token validation middleware for Axum web applications.
4//! This crate provides secure token validation with automatic public key caching,
5//! comprehensive error handling, and built-in security features.
6//!
7//! ## Features
8//!
9//! - **Secure JWT validation** with Firebase-specific claim verification
10//! - **Automatic public key caching** with configurable expiration
11//! - **Production-ready error handling** with detailed error types
12//! - **Security hardening** including token length limits and timing validation
13//! - **Retry logic with exponential backoff** for key fetching
14//! - **Comprehensive logging** for monitoring and debugging
15//!
16//! ## Quick Start
17//!
18//! ```rust,no_run
19//! use axum::{routing::get, Router, Extension, Json};
20//! use axum::middleware::from_fn_with_state;
21//! use serde_json::json;
22//! use axum_firebase_middleware::{FirebaseClaims, FirebaseConfig, firebase_auth_middleware};
23//!
24//! // Your protected handler
25//! async fn protected_handler(
26//!     Extension(claims): Extension<FirebaseClaims>
27//! ) -> Json<serde_json::Value> {
28//!     Json(json!({
29//!         "user_id": claims.user_id,
30//!         "email": claims.email
31//!     }))
32//! }
33//!
34//! #[tokio::main]
35//! async fn main() {
36//!     let config = FirebaseConfig::new("your-firebase-project-id".to_string())
37//!         .expect("Failed to create Firebase config");
38//!
39//!     let app = Router::new()
40//!         .route("/protected", get(protected_handler))
41//!         .route_layer(from_fn_with_state(config.clone(), firebase_auth_middleware))
42//!         .with_state(config);
43//!
44//!     // Run your server...
45//! }
46//! ```
47
48use axum::{
49    extract::{Request, State},
50    http::{HeaderMap, StatusCode},
51    middleware::Next,
52    response::Response,
53};
54use jsonwebtoken::{decode, decode_header, Algorithm, DecodingKey, Validation};
55use log::{debug, error, info, warn};
56use serde::{Deserialize, Serialize};
57use std::{
58    collections::HashMap,
59    sync::Arc,
60    time::{Duration, SystemTime, UNIX_EPOCH},
61};
62use tokio::sync::RwLock;
63use uuid::Uuid;
64
65const GOOGLE_CERTS_URL: &str =
66    "https://www.googleapis.com/robot/v1/metadata/x509/securetoken@system.gserviceaccount.com";
67const DEFAULT_CACHE_DURATION: u64 = 3600;
68const MAX_TOKEN_LENGTH: usize = 4096;
69const HTTP_TIMEOUT: Duration = Duration::from_secs(10);
70const MAX_RETRIES: u32 = 3;
71
72/// Firebase JWT claims structure containing user authentication information.
73///
74/// This struct represents the decoded claims from a Firebase ID token.
75/// All fields follow Firebase's JWT specification.
76#[derive(Debug, Clone, Serialize, Deserialize)]
77pub struct FirebaseClaims {
78    /// Token issuer (should be `https://securetoken.google.com/{project_id}`)
79    pub iss: String,
80    /// Audience (your Firebase project ID)
81    pub aud: String,
82    /// Authentication time (Unix timestamp)
83    pub auth_time: i64,
84    /// Firebase user ID
85    pub user_id: String,
86    /// Subject (should match user_id)
87    pub sub: String,
88    /// Issued at time (Unix timestamp)
89    pub iat: i64,
90    /// Expiration time (Unix timestamp)
91    pub exp: i64,
92    /// User's email address (if available)
93    pub email: Option<String>,
94    /// Whether the email has been verified
95    pub email_verified: Option<bool>,
96    /// Firebase authentication provider information
97    pub firebase: FirebaseAuthProvider,
98}
99
100/// Firebase authentication provider information.
101#[derive(Debug, Clone, Serialize, Deserialize)]
102pub struct FirebaseAuthProvider {
103    /// User identities from various providers
104    pub identities: HashMap<String, Vec<String>>,
105    /// The sign-in provider used for authentication
106    pub sign_in_provider: String,
107}
108
109/// Comprehensive error types for Firebase authentication failures.
110#[derive(Debug, thiserror::Error)]
111pub enum FirebaseAuthError {
112    #[error("Invalid token format: {0}")]
113    InvalidTokenFormat(String),
114
115    #[error("Token validation failed: {0}")]
116    ValidationFailed(String),
117
118    #[error("Public key fetch failed: {0}")]
119    KeyFetchFailed(String),
120
121    #[error("Token expired or invalid timing")]
122    InvalidTiming,
123
124    #[error("Missing required claims")]
125    MissingClaims,
126
127    #[error("Configuration error: {0}")]
128    ConfigError(String),
129
130    #[error("Rate limit exceeded")]
131    RateLimited,
132}
133
134/// Public key cache with automatic refresh and retry logic.
135///
136/// Handles fetching and caching of Firebase's public keys used for JWT verification.
137/// Keys are automatically refreshed based on the configured cache duration.
138#[derive(Clone)]
139pub struct PublicKeyCache {
140    keys: Arc<RwLock<HashMap<String, DecodingKey>>>,
141    last_updated: Arc<RwLock<SystemTime>>,
142    cache_duration: Duration,
143    http_client: reqwest::Client,
144    retry_count: Arc<RwLock<u32>>,
145}
146
147impl PublicKeyCache {
148    /// Creates a new public key cache with the specified cache duration.
149    ///
150    /// # Arguments
151    /// * `cache_duration_seconds` - How long to cache keys before refreshing
152    ///
153    /// # Errors
154    /// Returns `FirebaseAuthError::ConfigError` if HTTP client creation fails.
155    pub fn new(cache_duration_seconds: u64) -> Result<Self, FirebaseAuthError> {
156        let http_client = reqwest::Client::builder()
157            .timeout(HTTP_TIMEOUT)
158            .user_agent("Firebase-JWT-Validator/1.0")
159            .https_only(true)
160            .build()
161            .map_err(|e| {
162                FirebaseAuthError::ConfigError(format!("HTTP client creation failed: {}", e))
163            })?;
164
165        Ok(Self::with_client(cache_duration_seconds, http_client))
166    }
167
168    /// Creates a new public key cache backed by a caller-supplied HTTP client.
169    ///
170    /// Use this when the application needs full control over the `reqwest::Client`
171    /// (TLS backend, proxy, timeouts, connection pooling). Unlike [`Self::new`],
172    /// this never builds a client and therefore cannot fail.
173    ///
174    /// # Arguments
175    /// * `cache_duration_seconds` - How long to cache keys before refreshing
176    /// * `http_client` - The HTTP client used to fetch Google's certificates
177    pub fn with_client(cache_duration_seconds: u64, http_client: reqwest::Client) -> Self {
178        Self {
179            keys: Arc::new(RwLock::new(HashMap::new())),
180            last_updated: Arc::new(RwLock::new(UNIX_EPOCH)),
181            cache_duration: Duration::from_secs(cache_duration_seconds),
182            http_client,
183            retry_count: Arc::new(RwLock::new(0)),
184        }
185    }
186
187    /// Retrieves a decoding key for the given key ID.
188    ///
189    /// Automatically refreshes the cache if needed and handles key rotation.
190    ///
191    /// # Arguments
192    /// * `kid` - The key ID from the JWT header
193    ///
194    /// # Errors
195    /// Returns various `FirebaseAuthError` types for validation or fetch failures.
196    pub async fn get_key(&self, kid: &str) -> Result<DecodingKey, FirebaseAuthError> {
197        if kid.is_empty() || kid.len() > 128 {
198            return Err(FirebaseAuthError::InvalidTokenFormat(
199                "Invalid key ID".to_string(),
200            ));
201        }
202
203        let last_updated = *self.last_updated.read().await;
204        let now = SystemTime::now();
205
206        let needs_refresh =
207            now.duration_since(last_updated).unwrap_or(Duration::MAX) > self.cache_duration;
208
209        if needs_refresh {
210            self.refresh_keys().await?;
211        }
212
213        if let Some(key) = self.keys.read().await.get(kid).cloned() {
214            debug!("Public key cache hit for kid: {}", kid);
215            return Ok(key);
216        }
217
218        if !needs_refresh {
219            warn!("Key {} not found in fresh cache, forcing refresh", kid);
220            self.refresh_keys().await?;
221
222            if let Some(key) = self.keys.read().await.get(kid).cloned() {
223                return Ok(key);
224            }
225        }
226
227        Err(FirebaseAuthError::KeyFetchFailed(format!(
228            "Public key not found for kid: {}",
229            kid
230        )))
231    }
232
233    /// Refreshes public keys with exponential backoff retry logic.
234    async fn refresh_keys(&self) -> Result<(), FirebaseAuthError> {
235        let mut retry_count = *self.retry_count.read().await;
236        let mut delay = Duration::from_millis(100);
237
238        for attempt in 0..MAX_RETRIES {
239            match self.fetch_keys().await {
240                Ok(()) => {
241                    *self.retry_count.write().await = 0;
242                    info!("Successfully refreshed Firebase public keys");
243                    return Ok(());
244                }
245                Err(e) => {
246                    retry_count += 1;
247                    *self.retry_count.write().await = retry_count;
248
249                    if attempt < MAX_RETRIES - 1 {
250                        warn!(
251                            "Failed to fetch keys (attempt {}): {}. Retrying in {:?}",
252                            attempt + 1,
253                            e,
254                            delay
255                        );
256                        tokio::time::sleep(delay).await;
257                        delay *= 2;
258                    } else {
259                        error!("Failed to fetch keys after {} attempts: {}", MAX_RETRIES, e);
260                        return Err(e);
261                    }
262                }
263            }
264        }
265
266        unreachable!()
267    }
268
269    /// Fetches public keys from Google's certificate endpoint.
270    async fn fetch_keys(&self) -> Result<(), FirebaseAuthError> {
271        let response = self
272            .http_client
273            .get(GOOGLE_CERTS_URL)
274            .send()
275            .await
276            .map_err(|e| {
277                FirebaseAuthError::KeyFetchFailed(format!("HTTP request failed: {}", e))
278            })?;
279
280        if !response.status().is_success() {
281            return Err(FirebaseAuthError::KeyFetchFailed(format!(
282                "HTTP {} from Google certificates endpoint",
283                response.status()
284            )));
285        }
286
287        let certs: HashMap<String, String> = response.json().await.map_err(|e| {
288            FirebaseAuthError::KeyFetchFailed(format!("Invalid JSON response: {}", e))
289        })?;
290
291        if certs.is_empty() {
292            return Err(FirebaseAuthError::KeyFetchFailed(
293                "Empty certificates response".to_string(),
294            ));
295        }
296
297        let mut keys = HashMap::new();
298        let mut parse_errors = 0;
299
300        for (kid, cert) in certs {
301            if !cert.starts_with("-----BEGIN CERTIFICATE-----") {
302                warn!("Invalid certificate format for kid: {}", kid);
303                parse_errors += 1;
304                continue;
305            }
306
307            match DecodingKey::from_rsa_pem(cert.as_bytes()) {
308                Ok(key) => {
309                    keys.insert(kid.clone(), key);
310                    debug!("Successfully parsed certificate for kid: {}", kid);
311                }
312                Err(e) => {
313                    warn!("Failed to parse certificate for kid {}: {}", kid, e);
314                    parse_errors += 1;
315                }
316            }
317        }
318
319        if keys.is_empty() {
320            return Err(FirebaseAuthError::KeyFetchFailed(
321                "No valid certificates found".to_string(),
322            ));
323        }
324
325        if parse_errors > 0 {
326            warn!(
327                "Failed to parse {} out of {} certificates",
328                parse_errors,
329                keys.len() + parse_errors
330            );
331        }
332
333        *self.keys.write().await = keys;
334        *self.last_updated.write().await = SystemTime::now();
335
336        Ok(())
337    }
338}
339
340/// Firebase authentication configuration.
341///
342/// Contains all settings needed for JWT validation including project ID,
343/// key cache configuration, and security parameters.
344#[derive(Clone)]
345pub struct FirebaseConfig {
346    /// Firebase project ID
347    pub project_id: String,
348    /// Public key cache instance
349    pub key_cache: PublicKeyCache,
350    /// Maximum allowed token age
351    pub max_token_age: Duration,
352    /// Allowed JWT algorithms (defaults to RS256 only)
353    pub allowed_algorithms: Vec<Algorithm>,
354}
355
356impl FirebaseConfig {
357    /// Creates a new Firebase configuration with secure defaults.
358    ///
359    /// # Arguments
360    /// * `project_id` - Your Firebase project ID
361    ///
362    /// # Errors
363    /// Returns `FirebaseAuthError::ConfigError` for invalid project IDs or setup failures.
364    ///
365    /// # Example
366    /// ```rust,no_run
367    /// use axum_firebase_middleware::FirebaseConfig;
368    ///
369    /// let config = FirebaseConfig::new("my-firebase-project".to_string())?;
370    /// # Ok::<(), axum_firebase_middleware::FirebaseAuthError>(())
371    /// ```
372    pub fn new(project_id: String) -> Result<Self, FirebaseAuthError> {
373        Self::validate_project_id(&project_id)?;
374
375        let key_cache = PublicKeyCache::new(DEFAULT_CACHE_DURATION)?;
376
377        Ok(Self {
378            project_id,
379            key_cache,
380            max_token_age: Duration::from_secs(24 * 3600),
381            allowed_algorithms: vec![Algorithm::RS256],
382        })
383    }
384
385    /// Creates a new Firebase configuration backed by a caller-supplied HTTP client.
386    ///
387    /// This lets the application own the `reqwest::Client` and configure the TLS
388    /// backend, proxy, and timeouts itself, while reusing a single connection pool
389    /// across the app. The client is used to fetch Google's public certificates.
390    ///
391    /// # Arguments
392    /// * `project_id` - Your Firebase project ID
393    /// * `client` - The HTTP client used for certificate fetching
394    ///
395    /// # Errors
396    /// Returns `FirebaseAuthError::ConfigError` for invalid project IDs.
397    ///
398    /// # Example
399    /// ```rust,no_run
400    /// use axum_firebase_middleware::FirebaseConfig;
401    ///
402    /// let client = reqwest::Client::builder()
403    ///     .https_only(true)
404    ///     .build()
405    ///     .expect("failed to build HTTP client");
406    /// let config = FirebaseConfig::with_client("my-firebase-project".to_string(), client)?;
407    /// # Ok::<(), axum_firebase_middleware::FirebaseAuthError>(())
408    /// ```
409    pub fn with_client(
410        project_id: String,
411        client: reqwest::Client,
412    ) -> Result<Self, FirebaseAuthError> {
413        Self::validate_project_id(&project_id)?;
414
415        let key_cache = PublicKeyCache::with_client(DEFAULT_CACHE_DURATION, client);
416
417        Ok(Self {
418            project_id,
419            key_cache,
420            max_token_age: Duration::from_secs(24 * 3600),
421            allowed_algorithms: vec![Algorithm::RS256],
422        })
423    }
424
425    /// Validates a Firebase project ID against the expected format.
426    fn validate_project_id(project_id: &str) -> Result<(), FirebaseAuthError> {
427        if project_id.is_empty() {
428            return Err(FirebaseAuthError::ConfigError(
429                "Project ID cannot be empty".to_string(),
430            ));
431        }
432
433        if !project_id.chars().all(|c| c.is_alphanumeric() || c == '-') || project_id.len() > 30 {
434            return Err(FirebaseAuthError::ConfigError(
435                "Invalid project ID format".to_string(),
436            ));
437        }
438
439        Ok(())
440    }
441
442    /// Sets a custom cache duration for public keys.
443    ///
444    /// The existing HTTP client (including any client supplied via
445    /// [`Self::with_client`]) is preserved.
446    ///
447    /// # Arguments
448    /// * `seconds` - Cache duration in seconds
449    pub fn with_cache_duration(mut self, seconds: u64) -> Result<Self, FirebaseAuthError> {
450        self.key_cache = PublicKeyCache::with_client(seconds, self.key_cache.http_client.clone());
451        Ok(self)
452    }
453
454    /// Sets the maximum allowed token age.
455    ///
456    /// Tokens older than this duration will be rejected even if not expired.
457    ///
458    /// # Arguments
459    /// * `duration` - Maximum token age
460    pub fn with_max_token_age(mut self, duration: Duration) -> Self {
461        self.max_token_age = duration;
462        self
463    }
464}
465
466/// Extracts and validates Bearer token from Authorization header.
467///
468/// Performs comprehensive security checks including length limits,
469/// format validation, and character validation.
470fn extract_bearer_token(headers: &HeaderMap) -> Result<String, FirebaseAuthError> {
471    let auth_header = headers.get("authorization").ok_or_else(|| {
472        FirebaseAuthError::InvalidTokenFormat("Missing Authorization header".to_string())
473    })?;
474
475    let auth_str = auth_header.to_str().map_err(|_| {
476        FirebaseAuthError::InvalidTokenFormat("Invalid Authorization header encoding".to_string())
477    })?;
478
479    if !auth_str.starts_with("Bearer ") {
480        return Err(FirebaseAuthError::InvalidTokenFormat(
481            "Authorization header must use Bearer scheme".to_string(),
482        ));
483    }
484
485    let token = &auth_str[7..];
486
487    if token.is_empty() {
488        return Err(FirebaseAuthError::InvalidTokenFormat(
489            "Empty token".to_string(),
490        ));
491    }
492
493    if token.len() > MAX_TOKEN_LENGTH {
494        return Err(FirebaseAuthError::InvalidTokenFormat(
495            "Token too long".to_string(),
496        ));
497    }
498
499    let parts: Vec<&str> = token.split('.').collect();
500    if parts.len() != 3 {
501        return Err(FirebaseAuthError::InvalidTokenFormat(
502            "Invalid JWT format".to_string(),
503        ));
504    }
505
506    if token.contains('\0') || token.contains('\n') || token.contains('\r') {
507        return Err(FirebaseAuthError::InvalidTokenFormat(
508            "Token contains invalid characters".to_string(),
509        ));
510    }
511
512    Ok(token.to_string())
513}
514
515/// Validates Firebase JWT token with comprehensive security checks.
516///
517/// Performs all necessary validations including:
518/// - Signature verification using Google's public keys
519/// - Claims validation (issuer, audience, timing)
520/// - Firebase-specific claim verification
521/// - Security checks for token age and format
522async fn validate_firebase_token(
523    token: &str,
524    config: &FirebaseConfig,
525) -> Result<FirebaseClaims, FirebaseAuthError> {
526    let header = decode_header(token).map_err(|e| {
527        FirebaseAuthError::InvalidTokenFormat(format!("Invalid token header: {}", e))
528    })?;
529
530    let algorithm = header.alg;
531    if !config.allowed_algorithms.contains(&algorithm) {
532        return Err(FirebaseAuthError::ValidationFailed(format!(
533            "Algorithm {:?} not allowed",
534            algorithm
535        )));
536    }
537
538    let kid = header.kid.ok_or_else(|| {
539        FirebaseAuthError::InvalidTokenFormat("Missing key ID in token header".to_string())
540    })?;
541
542    let decoding_key = config.key_cache.get_key(&kid).await?;
543
544    let mut validation = Validation::new(algorithm);
545    validation.set_audience(&[&config.project_id]);
546    validation.set_issuer(&[&format!(
547        "https://securetoken.google.com/{}",
548        config.project_id
549    )]);
550    validation.validate_exp = true;
551    validation.validate_nbf = false;
552    validation.validate_aud = true;
553    validation.leeway = 60;
554    validation.reject_tokens_expiring_in_less_than = 0;
555
556    let token_data = decode::<FirebaseClaims>(token, &decoding_key, &validation).map_err(|e| {
557        FirebaseAuthError::ValidationFailed(format!("Token validation failed: {}", e))
558    })?;
559
560    let claims = token_data.claims;
561
562    if claims.sub.is_empty() || claims.sub.len() > 128 {
563        return Err(FirebaseAuthError::MissingClaims);
564    }
565
566    if claims.sub != claims.user_id {
567        return Err(FirebaseAuthError::ValidationFailed(
568            "Subject and user_id mismatch".to_string(),
569        ));
570    }
571
572    let now = SystemTime::now()
573        .duration_since(UNIX_EPOCH)
574        .unwrap()
575        .as_secs() as i64;
576
577    if claims.auth_time > now + 60 {
578        return Err(FirebaseAuthError::InvalidTiming);
579    }
580
581    let token_age = Duration::from_secs((now - claims.iat) as u64);
582    if token_age > config.max_token_age {
583        return Err(FirebaseAuthError::InvalidTiming);
584    }
585
586    let expected_issuer = format!("https://securetoken.google.com/{}", config.project_id);
587    if claims.iss != expected_issuer {
588        return Err(FirebaseAuthError::ValidationFailed(
589            "Invalid issuer".to_string(),
590        ));
591    }
592
593    if claims.aud != config.project_id {
594        return Err(FirebaseAuthError::ValidationFailed(
595            "Invalid audience".to_string(),
596        ));
597    }
598
599    let auth_age = Duration::from_secs((now - claims.auth_time) as u64);
600    if auth_age > Duration::from_secs(7 * 24 * 3600) {
601        return Err(FirebaseAuthError::InvalidTiming);
602    }
603
604    debug!(
605        "Successfully validated Firebase token for user: {}",
606        claims.user_id
607    );
608    Ok(claims)
609}
610
611/// Axum middleware for Firebase JWT authentication.
612///
613/// This middleware validates Firebase ID tokens and adds the decoded claims
614/// to the request extensions for use by downstream handlers.
615///
616/// # Security Features
617/// - Validates JWT signatures using Google's public keys
618/// - Enforces token expiration and timing constraints
619/// - Checks Firebase-specific claims and issuer
620/// - Prevents common attacks (oversized requests, malformed tokens)
621/// - Adds request IDs for tracing
622///
623/// # Usage
624/// ```rust,no_run
625/// use axum::{Router, routing::get, Json, response::IntoResponse, Extension, middleware::from_fn_with_state};
626/// use axum_firebase_middleware::{firebase_auth_middleware, FirebaseConfig, FirebaseClaims};
627/// use jsonwebtoken::errors::ErrorKind::Json as OtherJson;
628///
629///
630/// let config = FirebaseConfig::new("project-id".to_string())?;
631///
632/// async fn protected_handler(Extension(claims): Extension<FirebaseClaims>) -> impl IntoResponse {
633///     Json(serde_json::json!({
634///         "message": "Successfully authenticated",
635///         "user_id": claims.user_id,
636///         "email": claims.email
637///     }))
638///  }
639///
640/// let app = Router::new()
641///     .route("/protected", get(protected_handler))
642///     .route_layer(from_fn_with_state(config.clone(), firebase_auth_middleware))
643///     .with_state(config);
644/// # Ok::<(), axum_firebase_middleware::FirebaseAuthError>(())
645/// ```
646pub async fn firebase_auth_middleware(
647    State(config): State<FirebaseConfig>,
648    mut request: Request,
649    next: Next,
650) -> Result<Response, StatusCode> {
651    if let Some(content_length) = request.headers().get("content-length") {
652        if let Ok(length_str) = content_length.to_str() {
653            if let Ok(length) = length_str.parse::<usize>() {
654                if length > 10_485_760 {
655                    warn!("Request body too large: {} bytes", length);
656                    return Err(StatusCode::PAYLOAD_TOO_LARGE);
657                }
658            }
659        }
660    }
661
662    let token = match extract_bearer_token(request.headers()) {
663        Ok(token) => token,
664        Err(e) => {
665            warn!("Token extraction failed: {}", e);
666            return Err(StatusCode::UNAUTHORIZED);
667        }
668    };
669
670    let claims = match validate_firebase_token(&token, &config).await {
671        Ok(claims) => {
672            debug!("Successfully authenticated user: {}", claims.user_id);
673            claims
674        }
675        Err(e) => match e {
676            FirebaseAuthError::InvalidTokenFormat(_) | FirebaseAuthError::MissingClaims => {
677                warn!("Invalid token format: {}", e);
678                return Err(StatusCode::UNAUTHORIZED);
679            }
680            FirebaseAuthError::ValidationFailed(_) | FirebaseAuthError::InvalidTiming => {
681                warn!("Token validation failed: {}", e);
682                return Err(StatusCode::UNAUTHORIZED);
683            }
684            FirebaseAuthError::KeyFetchFailed(_) => {
685                error!("Key fetch failed: {}", e);
686                return Err(StatusCode::SERVICE_UNAVAILABLE);
687            }
688            FirebaseAuthError::RateLimited => {
689                warn!("Rate limit exceeded");
690                return Err(StatusCode::TOO_MANY_REQUESTS);
691            }
692            FirebaseAuthError::ConfigError(_) => {
693                error!("Configuration error: {}", e);
694                return Err(StatusCode::INTERNAL_SERVER_ERROR);
695            }
696        },
697    };
698
699    request.extensions_mut().insert(claims);
700
701    if request.extensions().get::<String>().is_none() {
702        let request_id = Uuid::new_v4().to_string();
703        request.extensions_mut().insert(request_id);
704    }
705
706    Ok(next.run(request).await)
707}
708
709#[cfg(test)]
710mod tests {
711    use super::*;
712    use axum::body::Body;
713    use axum::extract::FromRef;
714    use axum::http::Request;
715    use axum::middleware::from_fn_with_state;
716    use axum::response::IntoResponse;
717    use axum::routing::get;
718    use axum::{Extension, Json, Router};
719    use tower::ServiceExt;
720
721    #[derive(Clone, FromRef)]
722    struct AppStateMock {
723        fb: FirebaseConfig,
724    }
725
726    async fn health_check(
727        State(config): State<FirebaseConfig>,
728    ) -> Result<Json<serde_json::Value>, StatusCode> {
729        match config.key_cache.fetch_keys().await {
730            Ok(()) => Ok(Json(serde_json::json!({
731                "status": "healthy",
732                "firebase_keys": "accessible",
733                "timestamp": SystemTime::now()
734                    .duration_since(UNIX_EPOCH)
735                    .unwrap()
736                    .as_secs()
737            }))),
738            Err(_) => Err(StatusCode::SERVICE_UNAVAILABLE),
739        }
740    }
741
742    async fn protected_handler(Extension(claims): Extension<FirebaseClaims>) -> impl IntoResponse {
743        Json(serde_json::json!({
744            "message": "Successfully authenticated",
745            "user_id": claims.user_id,
746            "email": claims.email
747        }))
748    }
749
750    async fn create_route() -> Router {
751        let app_state = AppStateMock {
752            fb: FirebaseConfig::new("test-project-id".to_string()).unwrap(),
753        };
754
755        Router::new()
756            .route("/health", get(health_check))
757            .nest(
758                "/api/v1",
759                Router::new()
760                    .route("/protected", get(protected_handler))
761                    .route_layer(from_fn_with_state(
762                        app_state.fb.clone(),
763                        firebase_auth_middleware,
764                    )),
765            )
766            .with_state(app_state)
767    }
768
769    #[test]
770    fn test_extract_bearer_token() {
771        let mut headers = HeaderMap::new();
772        headers.insert("authorization", "Bearer test.token.123".parse().unwrap());
773        let token = extract_bearer_token(&headers).unwrap();
774        assert_eq!(token, "test.token.123");
775
776        let mut headers = HeaderMap::new();
777        headers.insert("authorization", "Basic invalid".parse().unwrap());
778        let result = extract_bearer_token(&headers);
779        assert!(result.is_err());
780        assert!(matches!(
781            result.unwrap_err(),
782            FirebaseAuthError::InvalidTokenFormat(_)
783        ));
784
785        let headers = HeaderMap::new();
786        let result = extract_bearer_token(&headers);
787        assert!(result.is_err());
788        assert!(matches!(
789            result.unwrap_err(),
790            FirebaseAuthError::InvalidTokenFormat(_)
791        ));
792
793        let mut headers = HeaderMap::new();
794        headers.insert("authorization", "Bearer ".parse().unwrap());
795        let result = extract_bearer_token(&headers);
796        assert!(result.is_err());
797        assert!(matches!(
798            result.unwrap_err(),
799            FirebaseAuthError::InvalidTokenFormat(_)
800        ));
801
802        let mut headers = HeaderMap::new();
803        let long_token = "Bearer ".to_string() + &"a".repeat(MAX_TOKEN_LENGTH + 1);
804        headers.insert("authorization", long_token.parse().unwrap());
805        let result = extract_bearer_token(&headers);
806        assert!(result.is_err());
807        assert!(matches!(
808            result.unwrap_err(),
809            FirebaseAuthError::InvalidTokenFormat(_)
810        ));
811
812        let mut headers = HeaderMap::new();
813        headers.insert("authorization", "Bearer part1.part2".parse().unwrap());
814        let result = extract_bearer_token(&headers);
815        assert!(result.is_err());
816        assert!(matches!(
817            result.unwrap_err(),
818            FirebaseAuthError::InvalidTokenFormat(_)
819        ));
820
821        let mut headers = HeaderMap::new();
822        headers.insert("authorization", "Bearer part1.part2.part3".parse().unwrap());
823        let result = extract_bearer_token(&headers);
824        assert!(result.is_ok());
825        assert_eq!(result.unwrap(), "part1.part2.part3");
826
827        let invalid_chars = ["token\0null", "token\nline", "token\rreturn"];
828        for invalid_token in invalid_chars {
829            let has_invalid_chars = invalid_token.contains('\0')
830                || invalid_token.contains('\n')
831                || invalid_token.contains('\r');
832            assert!(has_invalid_chars, "Token should contain invalid characters");
833        }
834    }
835
836    #[tokio::test]
837    async fn test_public_key_cache_creation() {
838        let cache = PublicKeyCache::new(3600);
839        assert!(cache.is_ok());
840        assert!(cache.unwrap().keys.read().await.is_empty());
841    }
842
843    #[tokio::test]
844    async fn test_firebase_config_creation() {
845        let project_id = "test-project-id".to_string();
846        let config = FirebaseConfig::new(project_id.clone());
847        assert!(config.is_ok());
848
849        let config = config.unwrap();
850        assert_eq!(config.project_id, project_id);
851        assert_eq!(config.allowed_algorithms, vec![Algorithm::RS256]);
852    }
853
854    #[tokio::test]
855    async fn test_firebase_auth_middleware_no_token() {
856        let app = create_route().await;
857
858        let request = Request::builder()
859            .uri("/api/v1/protected")
860            .body(Body::empty())
861            .unwrap();
862
863        let response = app.oneshot(request).await.unwrap();
864        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
865    }
866
867    #[tokio::test]
868    async fn test_firebase_auth_middleware_invalid_token() {
869        let app = create_route().await;
870
871        let request = Request::builder()
872            .uri("/api/v1/protected")
873            .header("Authorization", "Bearer invalid.token.format")
874            .body(Body::empty())
875            .unwrap();
876
877        let response = app.oneshot(request).await.unwrap();
878        assert_eq!(response.status(), StatusCode::UNAUTHORIZED);
879    }
880
881    #[tokio::test]
882    async fn test_firebase_auth_without_middleware() {
883        let app = create_route().await;
884
885        let request = Request::builder()
886            .uri("/health")
887            .body(Body::empty())
888            .unwrap();
889
890        let response = app.oneshot(request).await.unwrap();
891        assert_eq!(response.status(), StatusCode::OK);
892    }
893}
axum_firebase_middleware/lib.rs

axum_firebase_middleware/
lib.rs
