From ef021c7a88439d526284af93d8a9626b17436907 Mon Sep 17 00:00:00 2001
From: shay7sev <shay7sev@gmail.com>
Date: Wed, 11 Feb 2026 13:56:19 +0800
Subject: [PATCH] fix(key): fix key

---
 .gitignore        |  1 +
 Cargo.lock        | 36 +------------------------------
 src/utils/keys.rs | 55 +++++++++++++++++++++++++++++++++++------------
 3 files changed, 43 insertions(+), 49 deletions(-)

diff --git a/.gitignore b/.gitignore
index db4cec5..55a0c45 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,3 +5,4 @@
 !.env.example
 /log
 *.log
+/data
diff --git a/Cargo.lock b/Cargo.lock
index 172f3ac..7843d48 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -87,10 +87,9 @@ checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0"
 
 [[package]]
 name = "auth-kit"
-version = "0.1.0"
+version = "0.1.1"
 dependencies = [
  "axum",
- "axum-extra",
  "base64",
  "common-telemetry",
  "dashmap",
@@ -184,28 +183,6 @@ dependencies = [
  "tracing",
 ]
 
-[[package]]
-name = "axum-extra"
-version = "0.12.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "fef252edff26ddba56bbcdf2ee3307b8129acb86f5749b68990c168a6fcc9c76"
-dependencies = [
- "axum",
- "axum-core",
- "bytes",
- "cookie",
- "futures-core",
- "futures-util",
- "http",
- "http-body",
- "http-body-util",
- "mime",
- "pin-project-lite",
- "tower-layer",
- "tower-service",
- "tracing",
-]
-
 [[package]]
 name = "backon"
 version = "1.6.0"
@@ -413,17 +390,6 @@ dependencies = [
  "unicode-segmentation",
 ]
 
-[[package]]
-name = "cookie"
-version = "0.18.1"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "4ddef33a339a91ea89fb53151bd0a4689cfce27055c291dfa69945475d22c747"
-dependencies = [
- "percent-encoding",
- "time",
- "version_check",
-]
-
 [[package]]
 name = "core-foundation"
 version = "0.9.4"
diff --git a/src/utils/keys.rs b/src/utils/keys.rs
index 683ecb9..ea86fe4 100644
--- a/src/utils/keys.rs
+++ b/src/utils/keys.rs
@@ -4,7 +4,7 @@ use rsa::pkcs8::{DecodePublicKey, EncodePrivateKey, EncodePublicKey};
 use rsa::rand_core::OsRng;
 use rsa::traits::PublicKeyParts;
 use rsa::{RsaPrivateKey, RsaPublicKey, pkcs1::LineEnding};
-use std::sync::OnceLock;
+use std::{fs, path::PathBuf, sync::OnceLock};
 
 pub struct KeyPair {
     pub encoding_key: jsonwebtoken::EncodingKey,
@@ -30,19 +30,46 @@ pub fn get_keys() -> &'static KeyPair {
                 (priv_pem, pub_pem, public_key)
             }
             _ => {
-                let bits = 2048;
-                let private_key =
-                    RsaPrivateKey::new(&mut OsRng, bits).expect("failed to generate a key");
-                let public_key = RsaPublicKey::from(&private_key);
-                let private_pem = private_key
-                    .to_pkcs8_pem(LineEnding::LF)
-                    .expect("failed to encode private key")
-                    .to_string();
-                let public_pem = public_key
-                    .to_public_key_pem(LineEnding::LF)
-                    .expect("failed to encode public key")
-                    .to_string();
-                (private_pem, public_pem, public_key)
+                let key_dir = std::env::var("JWT_KEY_DIR").unwrap_or_else(|_| "./data".to_string());
+                let private_path = std::env::var("JWT_PRIVATE_KEY_PATH").unwrap_or_else(|_| {
+                    format!("{}/jwt_private_key.pem", key_dir.trim_end_matches('/'))
+                });
+                let public_path = std::env::var("JWT_PUBLIC_KEY_PATH").unwrap_or_else(|_| {
+                    format!("{}/jwt_public_key.pem", key_dir.trim_end_matches('/'))
+                });
+
+                let from_files = || -> Option<(String, String, RsaPublicKey)> {
+                    let priv_pem = fs::read_to_string(&private_path).ok()?;
+                    let pub_pem = fs::read_to_string(&public_path).ok()?;
+                    let public_key = RsaPublicKey::from_pkcs1_pem(&pub_pem)
+                        .or_else(|_| RsaPublicKey::from_public_key_pem(&pub_pem))
+                        .ok()?;
+                    Some((priv_pem, pub_pem, public_key))
+                };
+
+                if let Some((priv_pem, pub_pem, public_key)) = from_files() {
+                    (priv_pem, pub_pem, public_key)
+                } else {
+                    let bits = 2048;
+                    let private_key =
+                        RsaPrivateKey::new(&mut OsRng, bits).expect("failed to generate a key");
+                    let public_key = RsaPublicKey::from(&private_key);
+                    let private_pem = private_key
+                        .to_pkcs8_pem(LineEnding::LF)
+                        .expect("failed to encode private key")
+                        .to_string();
+                    let public_pem = public_key
+                        .to_public_key_pem(LineEnding::LF)
+                        .expect("failed to encode public key")
+                        .to_string();
+
+                    let key_dir_path = PathBuf::from(&key_dir);
+                    let _ = fs::create_dir_all(&key_dir_path);
+                    let _ = fs::write(&private_path, &private_pem);
+                    let _ = fs::write(&public_path, &public_pem);
+
+                    (private_pem, public_pem, public_key)
+                }
             }
         };