From 25061526d713ed153c717c7006a370fed512425b Mon Sep 17 00:00:00 2001
From: Yohei Kitamura <ykitamura@mdsol.com>
Date: Mon, 23 Dec 2024 11:21:06 +0900
Subject: [PATCH] Normalize provided private key string to the correct PEM
 format

---
 CHANGELOG.md      |  7 ++++++
 Cargo.toml        |  4 ++--
 src/lib.rs        |  2 ++
 src/pem_format.rs | 54 +++++++++++++++++++++++++++++++++++++++++++++++
 src/signer.rs     |  5 ++++-
 5 files changed, 69 insertions(+), 3 deletions(-)
 create mode 100644 src/pem_format.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 48e3ea6..4a34937 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,10 @@
+# 0.6.0
+- Normalize provided private key string to the correct PEM format.
+- Update thiserror to version 2.
+
+# 0.5.0
+- Convert from anyhow to thiserror.
+
 # 0.4.1
 - Add the Debug and Clone traits to `Signer` and `Verifier`.
 
diff --git a/Cargo.toml b/Cargo.toml
index 49fa11a..a64b6b1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "mauth-core"
-version = "0.5.0"
+version = "0.6.0"
 edition = "2021"
 authors = ["Medidata Solutions <support@mdsol.com>"]
 description = "Generate and verify Medidata MAuth protocol signatures"
@@ -13,7 +13,7 @@ keywords = ["security", "authentication"]
 categories = ["authentication"]
 
 [dependencies]
-thiserror = "1"
+thiserror = "2"
 base64 = "0.22"
 hex = "0.4"
 lazy-regex = "3"
diff --git a/src/lib.rs b/src/lib.rs
index 2f2fa85..e47c283 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -8,3 +8,5 @@ pub(crate) mod signable;
 pub mod signer;
 /// Signature verification for incoming requests
 pub mod verifier;
+
+mod pem_format;
diff --git a/src/pem_format.rs b/src/pem_format.rs
new file mode 100644
index 0000000..21a0ce0
--- /dev/null
+++ b/src/pem_format.rs
@@ -0,0 +1,54 @@
+const RSA_PRIVATE_KEY_HEADER: &str = "-----BEGIN RSA PRIVATE KEY-----";
+const RSA_PRIVATE_KEY_FOOTER: &str = "-----END RSA PRIVATE KEY-----";
+
+pub fn normalize_rsa_private_key(key: impl Into<String>) -> String {
+    let key = key.into();
+
+    match key.contains('\n') {
+        true => key,
+        false => {
+            let body = key
+                .trim()
+                .trim_start_matches(RSA_PRIVATE_KEY_HEADER)
+                .trim_end_matches(RSA_PRIVATE_KEY_FOOTER)
+                .trim();
+
+            let body = match body.contains(' ') {
+                true => body.replace(' ', "\n"),
+                false => body
+                    .chars()
+                    .collect::<Vec<char>>()
+                    .chunks(64)
+                    .map(|chunk| chunk.iter().collect::<String>())
+                    .collect::<Vec<String>>()
+                    .join("\n"),
+            };
+            format!(
+                "{}\n{}\n{}",
+                RSA_PRIVATE_KEY_HEADER, body, RSA_PRIVATE_KEY_FOOTER
+            )
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn normalize_rsa_private_key_test() {
+        let rsa_private_key = "-----BEGIN RSA PRIVATE KEY-----\nMIIBOgIBAAJBAKj34GkxFhD90vcNLYLInFEX6Ppy1tPf9Cnzj4p4WGeKLs1Pt8Qu\n/5OiPgoTdSy7bcF9IGpSE8ZgGKzgYQVZeN97YE00\n-----END RSA PRIVATE KEY-----";
+        let rsa_private_key_with_space = rsa_private_key.replace('\n', " ");
+        let rsa_private_key_without_newline = rsa_private_key.replace('\n', "");
+
+        assert_eq!(normalize_rsa_private_key(rsa_private_key), rsa_private_key);
+        assert_eq!(
+            normalize_rsa_private_key(rsa_private_key_with_space),
+            rsa_private_key
+        );
+        assert_eq!(
+            normalize_rsa_private_key(rsa_private_key_without_newline),
+            rsa_private_key
+        );
+    }
+}
diff --git a/src/signer.rs b/src/signer.rs
index 3d105da..76a8197 100644
--- a/src/signer.rs
+++ b/src/signer.rs
@@ -1,3 +1,4 @@
+use crate::pem_format;
 use crate::{error::Error, signable::Signable};
 use base64::{engine::general_purpose, Engine as _};
 use rsa::pkcs1::DecodeRsaPrivateKey;
@@ -27,7 +28,9 @@ impl Signer {
     /// assert!(signer.is_ok());
     /// ```
     pub fn new(app_uuid: impl Into<String>, private_key_data: String) -> Result<Self, Error> {
-        let private_key = RsaPrivateKey::from_pkcs1_pem(&private_key_data)?;
+        let private_key = RsaPrivateKey::from_pkcs1_pem(&pem_format::normalize_rsa_private_key(
+            private_key_data,
+        ))?;
         let signing_key = rsa::pkcs1v15::SigningKey::<Sha512>::new(private_key.to_owned());
 
         Ok(Self {