refactor: replace single shared with key rsa keys

9 files changed, 434 insertions, 31 deletions

diff --git a/Cargo.lock b/Cargo.lock
index a676533..62de408 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -15,18 +15,39 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6"
 
 [[package]]
+name = "base64ct"
+version = "1.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "55248b47b0caf0546f7988906588779981c43bb1bc9d0c44087278f80cdb44ba"
+
+[[package]]
 name = "bitflags"
 version = "2.9.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "1b8e56985ec62d17e9c1001dc89c88ecd7dc08e47eba5ec7c29c7b5eeecde967"
 
 [[package]]
+name = "block-buffer"
+version = "0.10.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71"
+dependencies = [
+ "generic-array",
+]
+
+[[package]]
 name = "bumpalo"
 version = "3.18.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "793db76d6187cd04dff33004d8e6c9cc4e05cd330500379d2394209271b4aeee"
 
 [[package]]
+name = "byteorder"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b"
+
+[[package]]
 name = "cc"
 version = "1.2.26"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -42,6 +63,42 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
 [[package]]
+name = "const-oid"
+version = "0.9.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8"
+
+[[package]]
+name = "cpufeatures"
+version = "0.2.17"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280"
+dependencies = [
+ "libc",
+]
+
+[[package]]
+name = "crypto-common"
+version = "0.1.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3"
+dependencies = [
+ "generic-array",
+ "typenum",
+]
+
+[[package]]
+name = "der"
+version = "0.7.10"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb"
+dependencies = [
+ "const-oid",
+ "pem-rfc7468",
+ "zeroize",
+]
+
+[[package]]
 name = "deranged"
 version = "0.4.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -51,6 +108,17 @@ dependencies = [
 ]
 
 [[package]]
+name = "digest"
+version = "0.10.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292"
+dependencies = [
+ "block-buffer",
+ "const-oid",
+ "crypto-common",
+]
+
+[[package]]
 name = "displaydoc"
 version = "0.2.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -71,6 +139,16 @@ dependencies = [
 ]
 
 [[package]]
+name = "generic-array"
+version = "0.14.7"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a"
+dependencies = [
+ "typenum",
+ "version_check",
+]
+
+[[package]]
 name = "getrandom"
 version = "0.2.16"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -234,12 +312,27 @@ dependencies = [
 ]
 
 [[package]]
+name = "lazy_static"
+version = "1.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe"
+dependencies = [
+ "spin",
+]
+
+[[package]]
 name = "libc"
 version = "0.2.172"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d750af042f7ef4f724306de029d18836c26c1765a54a6a3f094cbd23a7267ffa"
 
 [[package]]
+name = "libm"
+version = "0.2.15"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f9fbbcab51052fe104eb5e5d351cf728d30a5be1fe14d9be8a3b097481fb97de"
+
+[[package]]
 name = "litemap"
 version = "0.8.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -268,6 +361,23 @@ dependencies = [
 ]
 
 [[package]]
+name = "num-bigint-dig"
+version = "0.8.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc84195820f291c7697304f3cbdadd1cb7199c0efc917ff5eafd71225c136151"
+dependencies = [
+ "byteorder",
+ "lazy_static",
+ "libm",
+ "num-integer",
+ "num-iter",
+ "num-traits",
+ "rand",
+ "smallvec",
+ "zeroize",
+]
+
+[[package]]
 name = "num-conv"
 version = "0.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -283,12 +393,24 @@ dependencies = [
 ]
 
 [[package]]
+name = "num-iter"
+version = "0.1.45"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1429034a0490724d0075ebb2bc9e875d6503c3cf69e235a8941aa757d83ef5bf"
+dependencies = [
+ "autocfg",
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
 name = "num-traits"
 version = "0.2.19"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
 dependencies = [
  "autocfg",
+ "libm",
 ]
 
 [[package]]
@@ -308,12 +430,42 @@ dependencies = [
 ]
 
 [[package]]
+name = "pem-rfc7468"
+version = "0.7.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412"
+dependencies = [
+ "base64ct",
+]
+
+[[package]]
 name = "percent-encoding"
 version = "2.3.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e3148f5046208a5d56bcfc03053e3ca6334e51da8dfb19b6cdc8b306fae3283e"
 
 [[package]]
+name = "pkcs1"
+version = "0.7.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f"
+dependencies = [
+ "der",
+ "pkcs8",
+ "spki",
+]
+
+[[package]]
+name = "pkcs8"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7"
+dependencies = [
+ "der",
+ "spki",
+]
+
+[[package]]
 name = "potential_utf"
 version = "0.1.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -406,6 +558,26 @@ dependencies = [
 ]
 
 [[package]]
+name = "rsa"
+version = "0.9.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "78928ac1ed176a5ca1d17e578a1825f3d81ca54cf41053a592584b020cfd691b"
+dependencies = [
+ "const-oid",
+ "digest",
+ "num-bigint-dig",
+ "num-integer",
+ "num-traits",
+ "pkcs1",
+ "pkcs8",
+ "rand_core",
+ "signature",
+ "spki",
+ "subtle",
+ "zeroize",
+]
+
+[[package]]
 name = "rustversion"
 version = "1.0.21"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -450,12 +622,33 @@ dependencies = [
 ]
 
 [[package]]
+name = "sha2"
+version = "0.10.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
+[[package]]
 name = "shlex"
 version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
 
 [[package]]
+name = "signature"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de"
+dependencies = [
+ "digest",
+ "rand_core",
+]
+
+[[package]]
 name = "simple_asn1"
 version = "0.6.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -474,6 +667,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03"
 
 [[package]]
+name = "spin"
+version = "0.9.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67"
+
+[[package]]
+name = "spki"
+version = "0.7.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d"
+dependencies = [
+ "base64ct",
+ "der",
+]
+
+[[package]]
 name = "stable_deref_trait"
 version = "1.2.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -486,14 +695,22 @@ dependencies = [
  "base64",
  "jsonwebtoken",
  "rand",
+ "rsa",
  "serde",
  "serde_json",
+ "sha2",
  "url",
  "urlencoding",
  "uuid",
 ]
 
 [[package]]
+name = "subtle"
+version = "2.6.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292"
+
+[[package]]
 name = "syn"
 version = "2.0.101"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -577,6 +794,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "typenum"
+version = "1.18.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1dccffe3ce07af9386bfd29e80c0ab1a8205a2fc34e4bcd40364df902cfa8f3f"
+
+[[package]]
 name = "unicode-ident"
 version = "1.0.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -623,6 +846,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "version_check"
+version = "0.9.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a"
+
+[[package]]
 name = "wasi"
 version = "0.11.0+wasi-snapshot-preview1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -849,6 +1078,12 @@ dependencies = [
 ]
 
 [[package]]
+name = "zeroize"
+version = "1.8.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ced3678a2879b30306d323f4542626697a464a97c0a07c9aebf7ebca65cd4dde"
+
+[[package]]
 name = "zerotrie"
 version = "0.2.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
diff --git a/Cargo.toml b/Cargo.toml
index e943299..331cac5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -12,3 +12,5 @@ url = "2.0"
 base64 = "0.22"
 rand = "0.8"
 urlencoding = "2.1"
+rsa = "0.9"
+sha2 = "0.10"
diff --git a/spec/support/server.rb b/spec/support/server.rb
index 8de0d9d..d96239a 100644
--- a/spec/support/server.rb
+++ b/spec/support/server.rb
@@ -22,7 +22,6 @@ RSpec.configure do |config|
         Socket.tcp(ip, port.to_i) { true }
         break
       rescue Errno::ECONNREFUSED => error
-        puts error.inspect
         sleep 1
       end
     end
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
+}