refactor: replace single shared with key rsa keys

6 files changed, 197 insertions, 30 deletions

diff --git a/src/config.rs b/src/config.rs
index 3976d71..a13658b 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -2,21 +2,16 @@
 pub struct Config {
     pub bind_addr: String,
     pub issuer_url: String,
-    pub jwt_secret: String,
 }
 
 impl Config {
     pub fn from_env() -> Self {
         let bind_addr = std::env::var("BIND_ADDR").unwrap_or_else(|_| "127.0.0.1:7878".to_string());
         let issuer_url = format!("http://{}", bind_addr);
-        let jwt_secret = std::env::var("JWT_SECRET").unwrap_or_else(|_| {
-            "your-256-bit-secret-key-here-make-it-very-long-and-secure".to_string()
-        });
 
         Self {
             bind_addr,
             issuer_url,
-            jwt_secret,
         }
     }
-}
\ No newline at end of file
+}
diff --git a/src/http/mod.rs b/src/http/mod.rs
index a133f09..4523d3b 100644
--- a/src/http/mod.rs
+++ b/src/http/mod.rs
@@ -12,15 +12,15 @@ pub struct Server {
 }
 
 impl Server {
-    pub fn new(addr: String) -> Server {
+    pub fn new(addr: String) -> Result<Server, Box<dyn std::error::Error>> {
         let mut config = Config::from_env();
         config.bind_addr = addr;
         config.issuer_url = format!("http://{}", config.bind_addr);
 
-        Server {
-            oauth_server: OAuthServer::new(&config),
+        Ok(Server {
+            oauth_server: OAuthServer::new(&config)?,
             config,
-        }
+        })
     }
 
     pub fn start(&self) {
diff --git a/src/keys.rs b/src/keys.rs
new file mode 100644
index 0000000..6c25681
--- /dev/null
+++ b/src/keys.rs
@@ -0,0 +1,154 @@
+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,
+    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
+        });
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 1563317..1bff27d 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,5 +1,6 @@
 pub mod config;
 pub mod http;
+pub mod keys;
 pub mod oauth;
 
 pub use config::Config;
diff --git a/src/main.rs b/src/main.rs
index 47bd1ff..25d56ea 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -6,7 +6,7 @@ fn main() {
         Err(_) => String::from("127.0.0.1:7878"),
     };
 
-    let server = Server::new(bind_addr);
+    let server = Server::new(bind_addr).expect("Failed to create server");
     server.start();
 }
 
@@ -18,14 +18,13 @@ mod tests {
     #[test]
     fn test_oauth_server_creation() {
         let server = sts::http::Server::new("127.0.0.1:0".to_string());
-        // If we get here without panicking, the server was created successfully
-        assert!(true);
+        assert!(server.is_ok());
     }
 
     #[test]
     fn test_authorization_code_generation() {
         let config = sts::Config::from_env();
-        let oauth_server = sts::OAuthServer::new(&config);
+        let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server");
         let mut params = HashMap::new();
         params.insert("client_id".to_string(), "test_client".to_string());
         params.insert(
@@ -46,7 +45,7 @@ mod tests {
     #[test]
     fn test_missing_client_id() {
         let config = sts::Config::from_env();
-        let oauth_server = sts::OAuthServer::new(&config);
+        let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server");
         let mut params = HashMap::new();
         params.insert(
             "redirect_uri".to_string(),
@@ -62,7 +61,7 @@ mod tests {
     #[test]
     fn test_unsupported_response_type() {
         let config = sts::Config::from_env();
-        let oauth_server = sts::OAuthServer::new(&config);
+        let oauth_server = sts::OAuthServer::new(&config).expect("Failed to create OAuth server");
         let mut params = HashMap::new();
         params.insert("client_id".to_string(), "test_client".to_string());
         params.insert(
diff --git a/src/oauth/server.rs b/src/oauth/server.rs
index fdaddf6..888b0c2 100644
--- a/src/oauth/server.rs
+++ b/src/oauth/server.rs
@@ -1,6 +1,7 @@
 use crate::config::Config;
+use crate::keys::KeyManager;
 use crate::oauth::types::{AuthCode, Claims, ErrorResponse, TokenResponse};
-use jsonwebtoken::{DecodingKey, EncodingKey, Header, encode};
+use jsonwebtoken::{Algorithm, Header, encode};
 use std::collections::HashMap;
 use std::time::{SystemTime, UNIX_EPOCH};
 use url::Url;
@@ -8,26 +9,27 @@ use uuid::Uuid;
 
 pub struct OAuthServer {
     config: Config,
-    encoding_key: EncodingKey,
-    decoding_key: DecodingKey,
+    key_manager: std::sync::Mutex<KeyManager>,
     auth_codes: std::sync::Mutex<HashMap<String, AuthCode>>,
 }
 
 impl OAuthServer {
-    pub fn new(config: &Config) -> Self {
-        Self {
-            encoding_key: EncodingKey::from_secret(config.jwt_secret.as_ref()),
-            decoding_key: DecodingKey::from_secret(config.jwt_secret.as_ref()),
+    pub fn new(config: &Config) -> Result<Self, Box<dyn std::error::Error>> {
+        let key_manager = KeyManager::new()?;
+
+        Ok(Self {
+            key_manager: std::sync::Mutex::new(key_manager),
             auth_codes: std::sync::Mutex::new(HashMap::new()),
             config: config.clone(),
-        }
+        })
     }
 
     pub fn get_jwks(&self) -> String {
-        serde_json::json!({
-            "keys": []
-        })
-        .to_string()
+        let key_manager = self.key_manager.lock().unwrap();
+        match key_manager.get_jwks() {
+            Ok(jwks) => serde_json::to_string(&jwks).unwrap_or_else(|_| "{}".to_string()),
+            Err(_) => serde_json::json!({"keys": []}).to_string(),
+        }
     }
 
     pub fn handle_authorize(&self, params: &HashMap<String, String>) -> Result<String, String> {
@@ -143,6 +145,19 @@ impl OAuthServer {
         client_id: &str,
         scope: &Option<String>,
     ) -> Result<String, String> {
+        let mut key_manager = self.key_manager.lock().unwrap();
+
+        // Check if we need to rotate keys
+        if key_manager.should_rotate() {
+            if let Err(_) = key_manager.rotate_keys() {
+                return Err(self.error_response("server_error", "Key rotation failed"));
+            }
+        }
+
+        let current_key = key_manager
+            .get_current_key()
+            .ok_or_else(|| self.error_response("server_error", "No signing key available"))?;
+
         let now = SystemTime::now()
             .duration_since(UNIX_EPOCH)
             .unwrap()
@@ -157,7 +172,10 @@ impl OAuthServer {
             scope: scope.clone(),
         };
 
-        encode(&Header::default(), &claims, &self.encoding_key)
+        let mut header = Header::new(Algorithm::RS256);
+        header.kid = Some(current_key.kid.clone());
+
+        encode(&header, &claims, &current_key.encoding_key)
             .map_err(|_| self.error_response("server_error", "Failed to generate token"))
     }
 
@@ -168,4 +186,4 @@ impl OAuthServer {
         };
         serde_json::to_string(&error_resp).unwrap_or_else(|_| "{}".to_string())
     }
-}
\ No newline at end of file
+}