1 files changed, 32 insertions, 14 deletions

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
+}