path: root/src/keys.rs

use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
use jsonwebtoken::{DecodingKey, EncodingKey};
use rsa::pkcs8::{EncodePrivateKey, EncodePublicKey};
use rsa::traits::PublicKeyParts;
use rsa::{RsaPrivateKey, RsaPublicKey};
use serde::Serialize;
use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};
use uuid::Uuid;

#[derive(Clone)]
pub struct KeyPair {
    pub kid: String,
    pub private_key: RsaPrivateKey,
    pub public_key: RsaPublicKey,
    pub created_at: u64,
    pub encoding_key: EncodingKey,
    pub decoding_key: DecodingKey,
}

#[derive(Debug, Serialize)]
pub struct JwkKey {
    pub kty: String,
    #[serde(rename = "use")]
    pub use_: String,
    pub kid: String,
    pub alg: String,
    pub n: String,
    pub e: String,
}

#[derive(Debug, Serialize)]
pub struct Jwks {
    pub keys: Vec<JwkKey>,
}

pub struct KeyManager {
    keys: HashMap<String, KeyPair>,
    current_key_id: Option<String>,
    key_rotation_interval: u64, // seconds
}

impl KeyManager {
    pub fn new() -> Result<Self, Box<dyn std::error::Error>> {
        let mut manager = Self {
            keys: HashMap::new(),
            current_key_id: None,
            key_rotation_interval: 86400, // 24 hours
        };

        manager.generate_new_key()?;
        Ok(manager)
    }

    pub fn generate_new_key(&mut self) -> Result<String, Box<dyn std::error::Error>> {
        let mut rng = rand::thread_rng();
        let private_key = RsaPrivateKey::new(&mut rng, 2048)?;
        let public_key = RsaPublicKey::from(&private_key);

        let kid = Uuid::new_v4().to_string();
        let created_at = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();

        let encoding_key = EncodingKey::from_rsa_pem(
            &private_key
                .to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)?
                .as_bytes(),
        )?;
        let decoding_key = DecodingKey::from_rsa_pem(
            &public_key
                .to_public_key_pem(rsa::pkcs8::LineEnding::LF)?
                .as_bytes(),
        )?;

        let key_pair = KeyPair {
            kid: kid.clone(),
            private_key,
            public_key,
            created_at,
            encoding_key,
            decoding_key,
        };

        self.keys.insert(kid.clone(), key_pair);
        self.current_key_id = Some(kid.clone());

        Ok(kid)
    }

    pub fn get_current_key(&self) -> Option<&KeyPair> {
        self.current_key_id
            .as_ref()
            .and_then(|kid| self.keys.get(kid))
    }

    pub fn get_key(&self, kid: &str) -> Option<&KeyPair> {
        self.keys.get(kid)
    }

    pub fn should_rotate(&self) -> bool {
        if let Some(current_key) = self.get_current_key() {
            let now = SystemTime::now()
                .duration_since(UNIX_EPOCH)
                .unwrap()
                .as_secs();

            now - current_key.created_at > self.key_rotation_interval
        } else {
            true
        }
    }

    pub fn rotate_keys(&mut self) -> Result<(), Box<dyn std::error::Error>> {
        self.generate_new_key()?;
        Ok(())
    }

    pub fn get_jwks(&self) -> Result<Jwks, Box<dyn std::error::Error>> {
        let mut keys = Vec::new();

        for key_pair in self.keys.values() {
            let n = URL_SAFE_NO_PAD.encode(&key_pair.public_key.n().to_bytes_be());
            let e = URL_SAFE_NO_PAD.encode(&key_pair.public_key.e().to_bytes_be());

            keys.push(JwkKey {
                kty: "RSA".to_string(),
                use_: "sig".to_string(),
                kid: key_pair.kid.clone(),
                alg: "RS256".to_string(),
                n,
                e,
            });
        }

        Ok(Jwks { keys })
    }

    pub fn cleanup_old_keys(&mut self, max_age: u64) {
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_secs();

        let current_kid = self.current_key_id.clone();

        self.keys.retain(|kid, key_pair| {
            // Always keep the current key
            if Some(kid) == current_kid.as_ref() {
                return true;
            }

            // Keep keys that are not too old
            now - key_pair.created_at <= max_age
        });
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_key_manager_creation() {
        let manager = KeyManager::new().expect("Failed to create key manager");
        assert!(manager.get_current_key().is_some());
        assert_eq!(manager.keys.len(), 1);
    }

    #[test]
    fn test_key_generation() {
        let mut manager = KeyManager::new().expect("Failed to create key manager");
        let initial_key_count = manager.keys.len();
        
        let new_kid = manager.generate_new_key().expect("Failed to generate new key");
        assert_eq!(manager.keys.len(), initial_key_count + 1);
        assert_eq!(manager.current_key_id, Some(new_kid.clone()));
        assert!(manager.get_key(&new_kid).is_some());
    }

    #[test]
    fn test_jwks_generation() {
        let manager = KeyManager::new().expect("Failed to create key manager");
        let jwks = manager.get_jwks().expect("Failed to get JWKS");
        
        assert_eq!(jwks.keys.len(), 1);
        let key = &jwks.keys[0];
        assert_eq!(key.kty, "RSA");
        assert_eq!(key.use_, "sig");
        assert_eq!(key.alg, "RS256");
        assert!(!key.n.is_empty());
        assert!(!key.e.is_empty());
        assert!(!key.kid.is_empty());
    }

    #[test]
    fn test_key_rotation() {
        let mut manager = KeyManager::new().expect("Failed to create key manager");
        let original_kid = manager.current_key_id.clone().unwrap();
        
        manager.rotate_keys().expect("Failed to rotate keys");
        
        let new_kid = manager.current_key_id.clone().unwrap();
        assert_ne!(original_kid, new_kid);
        assert_eq!(manager.keys.len(), 2); // Should have both old and new keys
        assert!(manager.get_key(&original_kid).is_some());
        assert!(manager.get_key(&new_kid).is_some());
    }

    #[test]
    fn test_should_rotate_new_key() {
        let manager = KeyManager::new().expect("Failed to create key manager");
        // New key should not need rotation
        assert!(!manager.should_rotate());
    }

    #[test]
    fn test_should_rotate_old_key() {
        let mut manager = KeyManager::new().expect("Failed to create key manager");
        
        // Manually modify the current key's creation time to be old
        if let Some(current_kid) = manager.current_key_id.clone() {
            if let Some(key_pair) = manager.keys.get_mut(&current_kid) {
                let old_time = SystemTime::now()
                    .duration_since(UNIX_EPOCH)
                    .unwrap()
                    .as_secs() - 86401; // 1 day + 1 second ago
                    
                // We need to recreate the key pair with the old timestamp
                let mut old_key_pair = key_pair.clone();
                old_key_pair.created_at = old_time;
                manager.keys.insert(current_kid, old_key_pair);
            }
        }
        
        // Should need rotation since key is older than rotation interval
        assert!(manager.should_rotate());
    }

    #[test]
    fn test_cleanup_old_keys() {
        let mut manager = KeyManager::new().expect("Failed to create key manager");
        let original_kid = manager.current_key_id.clone().unwrap();
        
        // Generate a new key (so we have 2 keys)
        manager.rotate_keys().expect("Failed to rotate keys");
        assert_eq!(manager.keys.len(), 2);
        
        // Cleanup with max_age 0 should remove old keys but keep current
        manager.cleanup_old_keys(0);
        assert_eq!(manager.keys.len(), 1);
        assert!(manager.get_key(&original_kid).is_none());
        assert!(manager.get_current_key().is_some());
    }

    #[test]
    fn test_multiple_key_jwks() {
        let mut manager = KeyManager::new().expect("Failed to create key manager");
        manager.rotate_keys().expect("Failed to rotate keys");
        manager.rotate_keys().expect("Failed to rotate keys");
        
        let jwks = manager.get_jwks().expect("Failed to get JWKS");
        assert_eq!(jwks.keys.len(), 3); // Should have 3 keys
        
        // All keys should have unique key IDs
        let mut kids: Vec<String> = jwks.keys.iter().map(|k| k.kid.clone()).collect();
        kids.sort();
        kids.dedup();
        assert_eq!(kids.len(), 3); // All should be unique
    }
}