From 9d9575843cffb6f6077057becd5681ddf7ec0651 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Thu, 14 Nov 2024 00:14:25 +0800
Subject: [PATCH] feat: update jose-test

---
 __crypto/jose-test/src/jose.rs | 36 +++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 14 deletions(-)

diff --git a/__crypto/jose-test/src/jose.rs b/__crypto/jose-test/src/jose.rs
index 87b0b15..aa08f6c 100644
--- a/__crypto/jose-test/src/jose.rs
+++ b/__crypto/jose-test/src/jose.rs
@@ -62,6 +62,11 @@ const LOCAL_KMS_PREFIX: &str = "LKMS:";
 //
 // https://security.stackexchange.com/questions/80966/what-is-the-point-of-aes-key-wrap-with-json-web-encryption
 
+const JWE_ENC_A256GCM: &str = "A256GCM";
+const JWE_ALG_A256KW: &str = "A256KW";
+const JWE_ALG_RSA_OAEP: &str = "RSA-OAEP";
+const JWE_DOT: &str = ".";
+
 pub fn generate_rsa_key_2(bits: u32) -> XResult<RsaPrivateKey> {
     let mut rng = OsRng::default();
     Ok(RsaPrivateKey::new(&mut rng, bits as usize)?)
@@ -73,9 +78,10 @@ pub fn generate_rsa_key(bits: u32) -> XResult<RsaKeyPair> {
 
 pub fn serialize_jwe_rsa_2(payload: &[u8], rsa_public_key: &RsaPublicKey) -> XResult<String> {
     let header = JweHeader2 {
-        enc: "A256GCM".to_string(),
-        alg: "RSA-OAEP".to_string(),
+        enc: JWE_ENC_A256GCM.to_string(),
+        alg: JWE_ALG_RSA_OAEP.to_string(),
         vendor: "local-mini-kms".to_string(),
+        version: None,
     };
     serialize_jwe_fn(&header, payload, |data_key| -> XResult<Vec<u8>> {
         let mut r = ThreadRng::default();
@@ -85,7 +91,7 @@ pub fn serialize_jwe_rsa_2(payload: &[u8], rsa_public_key: &RsaPublicKey) -> XRe
 
 pub fn serialize_jwe_rsa(payload: &[u8], jwk: &Jwk) -> XResult<String> {
     let mut header = JweHeader::new();
-    header.set_content_encryption("A256GCM");
+    header.set_content_encryption(JWE_ENC_A256GCM);
     header.set_claim("vendor", Some(Value::String("local-mini-kms".to_string())))?;
     let encrypter = RsaesJweAlgorithm::RsaOaep.encrypter_from_jwk(jwk)?;
     Ok(format!("{}{}", LOCAL_KMS_PREFIX, jwe::serialize_compact(payload, &header, &encrypter)?))
@@ -93,7 +99,7 @@ pub fn serialize_jwe_rsa(payload: &[u8], jwk: &Jwk) -> XResult<String> {
 
 pub fn deserialize_jwe_rsa_2(jwe: &str, rsa: &RsaPrivateKey) -> XResult<(Vec<u8>, JweHeader2)> {
     deserialize_jwe_fn(jwe, |alg, key_wrap| -> XResult<Vec<u8>> {
-        if alg != "RSA-OAEP" {
+        if alg != JWE_ALG_RSA_OAEP {
             return simple_error!("Invalid JWE header alg: {}", alg);
         }
         Ok(opt_result!(rsa.decrypt(Oaep::new::<Sha1>(), &key_wrap), "Unwrap key failed: {}"))
@@ -110,13 +116,15 @@ pub struct JweHeader2 {
     pub enc: String,
     pub alg: String,
     pub vendor: String,
+    pub version: Option<String>,
 }
 
 pub fn serialize_jwe_aes_2(payload: &[u8], key: [u8; 32]) -> XResult<String> {
     let header = JweHeader2 {
-        enc: "A256GCM".to_string(),
-        alg: "A256KW".to_string(),
+        enc: JWE_ENC_A256GCM.to_string(),
+        alg: JWE_ALG_A256KW.to_string(),
         vendor: "local-mini-kms".to_string(),
+        version: None,
     };
     serialize_jwe_fn(&header, payload, |data_key| -> XResult<Vec<u8>> {
         let kek = Kek::from(key);
@@ -126,7 +134,7 @@ pub fn serialize_jwe_aes_2(payload: &[u8], key: [u8; 32]) -> XResult<String> {
 
 pub fn serialize_jwe_aes(payload: &[u8], key: &[u8]) -> XResult<String> {
     let mut header = JweHeader::new();
-    header.set_content_encryption("A256GCM");
+    header.set_content_encryption(JWE_ENC_A256GCM);
     header.set_claim("vendor", Some(Value::String("local-mini-kms".to_string())))?;
     // header.set_claim("version", Some(Value::String(get_master_key_checksum(key))))?;
     let encrypter = AeskwJweAlgorithm::A256kw.encrypter_from_bytes(key)?;
@@ -135,7 +143,7 @@ pub fn serialize_jwe_aes(payload: &[u8], key: &[u8]) -> XResult<String> {
 
 pub fn deserialize_jwe_aes_2(jwe: &str, key: [u8; 32]) -> XResult<(Vec<u8>, JweHeader2)> {
     deserialize_jwe_fn(jwe, |alg, key_wrap| -> XResult<Vec<u8>> {
-        if alg != "A256KW" {
+        if alg != JWE_ALG_A256KW {
             return simple_error!("Invalid JWE header alg: {}", alg);
         }
         let kek = Kek::from(key);
@@ -167,13 +175,13 @@ where
 
     let mut jwe = String::new();
     jwe.push_str(&header_b64);
-    jwe.push_str(".");
+    jwe.push_str(JWE_DOT);
     jwe.push_str(&URL_SAFE_NO_PAD.encode(&wrap_key));
-    jwe.push_str(".");
+    jwe.push_str(JWE_DOT);
     jwe.push_str(&URL_SAFE_NO_PAD.encode(&nonce));
-    jwe.push_str(".");
+    jwe.push_str(JWE_DOT);
     jwe.push_str(&URL_SAFE_NO_PAD.encode(&e));
-    jwe.push_str(".");
+    jwe.push_str(JWE_DOT);
     jwe.push_str(&URL_SAFE_NO_PAD.encode(&t));
 
     Ok(jwe)
@@ -183,13 +191,13 @@ fn deserialize_jwe_fn<F>(jwe: &str, key_unwrap_fn: F) -> XResult<(Vec<u8>, JweHe
 where
     F: Fn(&str, &[u8]) -> XResult<Vec<u8>>,
 {
-    let jwe_parts = jwe.split(".").collect::<Vec<&str>>();
+    let jwe_parts = jwe.split(JWE_DOT).collect::<Vec<&str>>();
     if jwe_parts.len() != 5 {
         return simple_error!("Invalid JWE: {}", jwe);
     }
     let header_bytes = opt_result!(decode_url_safe_no_pad(jwe_parts[0]), "Invalid JWE header: {}, JWE: {}", jwe);
     let header: JweHeader2 = opt_result!(serde_json::from_slice(&header_bytes), "Invalid JWE header: {}, JWE: {}", jwe);
-    if header.enc != "A256GCM" {
+    if header.enc != JWE_ENC_A256GCM {
         return simple_error!("Invalid JWE header enc: {}", header.enc);
     }