// ABOUTME: Tests for RSA key persistence and cross-process key synchronization
// ABOUTME: Validates CLI and server use the same keys via database persistence
//
// SPDX-License-Identifier: MIT OR Apache-2.0
// Copyright (c) 2025 Pierre Fitness Intelligence
#![allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
#![allow(missing_docs)]
mod common;
use anyhow::Result;
use chrono::Utc;
#[cfg(feature = "postgresql")]
use pierre_mcp_server::config::environment::PostgresPoolConfig;
use pierre_mcp_server::{
admin::{
jwks::JwksManager,
jwt::AdminJwtManager,
models::{AdminPermission, AdminPermissions},
},
auth::AuthManager,
database,
database_plugins::{factory::Database, DatabaseProvider},
models::User,
};
use std::{sync::Arc, time::Duration};
use tokio::time::sleep;
/// Create a test database for RSA key persistence tests
async fn create_rsa_test_database() -> Result<Database> {
common::init_test_logging();
let database_url = "sqlite::memory:";
let encryption_key = database::generate_encryption_key().to_vec();
#[cfg(feature = "postgresql")]
let database =
Database::new(database_url, encryption_key, &PostgresPoolConfig::default()).await?;
#[cfg(not(feature = "postgresql"))]
let database = Database::new(database_url, encryption_key).await?;
Ok(database)
}
/// Test that RSA keys can be saved and loaded from database
#[tokio::test]
async fn test_rsa_keypair_persistence() -> Result<()> {
let database = create_rsa_test_database().await?;
// Generate a key and save it
let mut jwks_manager = JwksManager::new();
let kid = format!("test_key_{}", Utc::now().format("%Y%m%d_%H%M%S"));
jwks_manager.generate_rsa_key_pair_with_size(&kid, 2048)?;
let key_pair = jwks_manager.get_active_key()?;
let private_pem = key_pair.export_private_key_pem()?;
let public_pem = key_pair.export_public_key_pem()?;
let created_at = Utc::now();
// Save to database
database
.save_rsa_keypair(&kid, &private_pem, &public_pem, created_at, true, 2048)
.await?;
// Load from database
let loaded_keys = database.load_rsa_keypairs().await?;
assert_eq!(loaded_keys.len(), 1, "Should have exactly one key");
let (loaded_kid, loaded_private, loaded_public, _, loaded_active) = &loaded_keys[0];
assert_eq!(loaded_kid, &kid, "Key ID should match");
assert_eq!(loaded_private, &private_pem, "Private key PEM should match");
assert_eq!(loaded_public, &public_pem, "Public key PEM should match");
assert!(*loaded_active, "Key should be marked as active");
Ok(())
}
/// Test that a JWKS manager can load keys from database
#[tokio::test]
async fn test_jwks_manager_loads_from_database() -> Result<()> {
let database = create_rsa_test_database().await?;
// Create and save a key (simulates CLI admin-setup)
let mut original_jwks = JwksManager::new();
let kid = format!("cli_key_{}", Utc::now().format("%Y%m%d_%H%M%S"));
original_jwks.generate_rsa_key_pair_with_size(&kid, 2048)?;
let key_pair = original_jwks.get_active_key()?;
let private_pem = key_pair.export_private_key_pem()?;
let public_pem = key_pair.export_public_key_pem()?;
let created_at = Utc::now();
database
.save_rsa_keypair(&kid, &private_pem, &public_pem, created_at, true, 2048)
.await?;
// Create a new JWKS manager and load from database (simulates server startup)
let mut loaded_jwks = JwksManager::new();
let keypairs = database.load_rsa_keypairs().await?;
loaded_jwks.load_keys_from_database(keypairs)?;
// Verify the loaded key matches
let loaded_key = loaded_jwks.get_active_key()?;
assert_eq!(loaded_key.kid, kid, "Loaded key ID should match original");
Ok(())
}
/// Test that admin tokens generated by CLI are valid on server with shared DB keys
/// This is the exact scenario that was broken before the fix
#[tokio::test]
async fn test_cli_generated_admin_token_valid_on_server() -> Result<()> {
let database = create_rsa_test_database().await?;
// === CLI SIDE (admin-setup) ===
// CLI generates key and saves to database
let mut cli_jwks = JwksManager::new();
let kid = format!("admin_key_{}", Utc::now().format("%Y%m%d_%H%M%S"));
cli_jwks.generate_rsa_key_pair_with_size(&kid, 2048)?;
let key_pair = cli_jwks.get_active_key()?;
let private_pem = key_pair.export_private_key_pem()?;
let public_pem = key_pair.export_public_key_pem()?;
let created_at = Utc::now();
database
.save_rsa_keypair(&kid, &private_pem, &public_pem, created_at, true, 2048)
.await?;
// CLI generates admin token
let cli_jwks_arc = Arc::new(cli_jwks);
let jwt_manager = AdminJwtManager::new();
let token_id = "admin_token_123";
let service_name = "test_cli_service";
let permissions = AdminPermissions::default_admin();
let admin_token = jwt_manager.generate_token(
token_id,
service_name,
&permissions,
false,
None,
&cli_jwks_arc,
)?;
// === SERVER SIDE ===
// Server loads keys from database (same database)
let mut server_jwks = JwksManager::new();
let keypairs = database.load_rsa_keypairs().await?;
server_jwks.load_keys_from_database(keypairs)?;
let server_jwks_arc = Arc::new(server_jwks);
// Server validates the CLI-generated token
let validated = jwt_manager.validate_token(&admin_token, &server_jwks_arc)?;
assert_eq!(validated.token_id, token_id, "Token ID should match");
assert_eq!(
validated.service_name, service_name,
"Service name should match"
);
assert!(!validated.is_super_admin, "Should not be super admin");
Ok(())
}
/// Test that user session tokens generated with persisted keys are valid
#[tokio::test]
async fn test_user_session_token_with_persisted_keys() -> Result<()> {
let database = create_rsa_test_database().await?;
// Process 1: Generate and save key
let mut process1_jwks = JwksManager::new();
let kid = format!("session_key_{}", Utc::now().format("%Y%m%d_%H%M%S"));
process1_jwks.generate_rsa_key_pair_with_size(&kid, 2048)?;
let key_pair = process1_jwks.get_active_key()?;
let private_pem = key_pair.export_private_key_pem()?;
let public_pem = key_pair.export_public_key_pem()?;
let created_at = Utc::now();
database
.save_rsa_keypair(&kid, &private_pem, &public_pem, created_at, true, 2048)
.await?;
// Process 1: Generate user token
let process1_jwks_arc = Arc::new(process1_jwks);
let auth_manager = AuthManager::new(24);
let user = User::new(
"cross_process@example.com".to_owned(),
"password_hash".to_owned(),
Some("Cross Process User".to_owned()),
);
let token = auth_manager.generate_token(&user, &process1_jwks_arc)?;
// Process 2: Load keys from database and validate
let mut process2_jwks = JwksManager::new();
let keypairs = database.load_rsa_keypairs().await?;
process2_jwks.load_keys_from_database(keypairs)?;
let process2_jwks_arc = Arc::new(process2_jwks);
// Token generated by process 1 should be valid in process 2
let claims = auth_manager.validate_token(&token, &process2_jwks_arc)?;
assert_eq!(claims.sub, user.id.to_string(), "User ID should match");
assert_eq!(claims.email, user.email, "Email should match");
Ok(())
}
/// Test that tokens fail validation with different keys (not from DB)
/// This documents the exact bug scenario: ephemeral keys don't work cross-process
#[tokio::test]
async fn test_ephemeral_keys_fail_cross_process() -> Result<()> {
// Process 1: Generate ephemeral key (NOT saved to DB)
let mut process1_jwks = JwksManager::new();
process1_jwks.generate_rsa_key_pair_with_size("ephemeral_key_1", 2048)?;
let process1_jwks_arc = Arc::new(process1_jwks);
let auth_manager = AuthManager::new(24);
let user = User::new(
"ephemeral_test@example.com".to_owned(),
"password_hash".to_owned(),
Some("Ephemeral Test User".to_owned()),
);
let token = auth_manager.generate_token(&user, &process1_jwks_arc)?;
// Process 2: Generate different ephemeral key (NOT loaded from DB)
let mut process2_jwks = JwksManager::new();
process2_jwks.generate_rsa_key_pair_with_size("ephemeral_key_2", 2048)?;
let process2_jwks_arc = Arc::new(process2_jwks);
// Token should FAIL validation with different keys
let result = auth_manager.validate_token(&token, &process2_jwks_arc);
assert!(
result.is_err(),
"Token should fail validation with different ephemeral keys"
);
Ok(())
}
/// Test multiple keys can be stored and the correct one is used
#[tokio::test]
async fn test_multiple_persisted_keys() -> Result<()> {
let database = create_rsa_test_database().await?;
// Save first key (inactive)
let mut jwks1 = JwksManager::new();
let kid1 = format!("key1_{}", Utc::now().format("%Y%m%d_%H%M%S"));
jwks1.generate_rsa_key_pair_with_size(&kid1, 2048)?;
let key_pair1 = jwks1.get_active_key()?;
let private_pem1 = key_pair1.export_private_key_pem()?;
let public_pem1 = key_pair1.export_public_key_pem()?;
let created_at1 = Utc::now();
database
.save_rsa_keypair(&kid1, &private_pem1, &public_pem1, created_at1, false, 2048)
.await?;
// Small delay to ensure different timestamps
sleep(Duration::from_millis(100)).await;
// Save second key (active)
let mut jwks2 = JwksManager::new();
let kid2 = format!("key2_{}", Utc::now().format("%Y%m%d_%H%M%S"));
jwks2.generate_rsa_key_pair_with_size(&kid2, 2048)?;
let key_pair2 = jwks2.get_active_key()?;
let private_pem2 = key_pair2.export_private_key_pem()?;
let public_pem2 = key_pair2.export_public_key_pem()?;
let created_at2 = Utc::now();
database
.save_rsa_keypair(&kid2, &private_pem2, &public_pem2, created_at2, true, 2048)
.await?;
// Load all keys
let loaded_keys = database.load_rsa_keypairs().await?;
assert_eq!(loaded_keys.len(), 2, "Should have two keys");
// Load into JWKS manager
let mut loaded_jwks = JwksManager::new();
loaded_jwks.load_keys_from_database(loaded_keys)?;
// Active key should be kid2
let active_key = loaded_jwks.get_active_key()?;
assert_eq!(active_key.kid, kid2, "Active key should be the second key");
// Both keys should be retrievable
assert!(
loaded_jwks.get_key(&kid1).is_some(),
"First key should be retrievable"
);
assert!(
loaded_jwks.get_key(&kid2).is_some(),
"Second key should be retrievable"
);
Ok(())
}
/// Test that super admin tokens work across processes with shared DB keys
#[tokio::test]
async fn test_super_admin_token_persistence() -> Result<()> {
let database = create_rsa_test_database().await?;
// Generate and save key
let mut jwks = JwksManager::new();
let kid = format!("super_admin_key_{}", Utc::now().format("%Y%m%d_%H%M%S"));
jwks.generate_rsa_key_pair_with_size(&kid, 2048)?;
let key_pair = jwks.get_active_key()?;
let private_pem = key_pair.export_private_key_pem()?;
let public_pem = key_pair.export_public_key_pem()?;
let created_at = Utc::now();
database
.save_rsa_keypair(&kid, &private_pem, &public_pem, created_at, true, 2048)
.await?;
// Generate super admin token
let jwks_arc = Arc::new(jwks);
let jwt_manager = AdminJwtManager::new();
let permissions = AdminPermissions::super_admin();
let token = jwt_manager.generate_token(
"super_token",
"super_service",
&permissions,
true, // is_super_admin
None,
&jwks_arc,
)?;
// Load keys in "different process" and validate
let mut loaded_jwks = JwksManager::new();
let keypairs = database.load_rsa_keypairs().await?;
loaded_jwks.load_keys_from_database(keypairs)?;
let loaded_jwks_arc = Arc::new(loaded_jwks);
let validated = jwt_manager.validate_token(&token, &loaded_jwks_arc)?;
assert!(validated.is_super_admin, "Should be super admin");
assert!(
validated
.permissions
.has_permission(&AdminPermission::ManageAdminTokens),
"Should have ManageAdminTokens permission"
);
Ok(())
}
