From 75f9c877382c0f0380e667ab8b6b253c98cdd0cd Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Thu, 24 Jul 2025 23:51:36 +0800
Subject: [PATCH] feat: sm4 jwe

---
 .gitignore |   1 +
 Cargo.toml |  13 +++++
 README.md  |   3 ++
 src/lib.rs | 144 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 161 insertions(+)
 create mode 100644 Cargo.toml
 create mode 100644 src/lib.rs

diff --git a/.gitignore b/.gitignore
index 0eabc4c..ad19215 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+.idea/
 # ---> Rust
 # Generated by Cargo
 # will have compiled files and executables
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..b137072
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,13 @@
+[package]
+name = "sm4-jwe"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+base64 = "0.22.1"
+rand = "0.9.2"
+rust_util = "0.6.47"
+serde = { version = "1.0.219", features = ["derive"] }
+serde_json = "1.0.141"
+sm4-gcm = "0.1.2"
+zeroize = "1.8.1"
diff --git a/README.md b/README.md
index 0a6ae91..2cafc1e 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,5 @@
 # sm4-jwe
 
+> Specification: https://openwebstandard.org/rfc5
+
+
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..5799c7e
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,144 @@
+use base64::engine::general_purpose::URL_SAFE_NO_PAD;
+use base64::Engine;
+use rust_util::{opt_result, simple_error, SimpleError, XResult};
+use serde::{Deserialize, Serialize};
+use sm4_gcm::{sm4_gcm_encrypt, Sm4GcmStreamEncryptor, Sm4Key};
+use zeroize::Zeroize;
+
+const SM4GCM: &str = "SM4GCM";
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct JweHeader {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    kid: Option<String>,
+    enc: String,
+    alg: String,
+}
+
+#[derive(Debug, Eq, PartialEq, Copy, Clone)]
+pub enum JweAlg {
+    Dir,
+    Sm2Pke,
+    Sm4Skw,
+}
+
+impl JweAlg {
+    pub fn get_name(&self) -> &'static str {
+        match self {
+            JweAlg::Dir => "dir",
+            JweAlg::Sm2Pke => "SM2PKE",
+            JweAlg::Sm4Skw => "SM4SKW",
+        }
+    }
+}
+
+pub struct Sm4Jwk {
+    key_id: Option<String>,
+    sm4key: Sm4Key,
+}
+
+impl TryFrom<Vec<u8>> for Sm4Jwk {
+    type Error = SimpleError;
+
+    fn try_from(mut value: Vec<u8>) -> Result<Self, Self::Error> {
+        let sm4key = opt_result!(Sm4Key::from_slice(&value), "Invalid SM4 key: {}");
+        value.zeroize();
+        Ok(Sm4Jwk {
+            key_id: None,
+            sm4key,
+        })
+    }
+}
+
+impl TryFrom<[u8; 16]> for Sm4Jwk {
+    type Error = SimpleError;
+
+    fn try_from(mut value: [u8; 16]) -> Result<Self, Self::Error> {
+        let sm4key = opt_result!(Sm4Key::from_slice(&value), "Invalid SM4 key: {}");
+        value.zeroize();
+        Ok(Sm4Jwk {
+            key_id: None,
+            sm4key,
+        })
+    }
+}
+
+#[test]
+fn test() {
+    let k = Sm4Jwk::try_from(b"0123456789012345".to_vec())
+        .unwrap()
+        .key_id("k001");
+    println!("{}", k.encrypt(JweAlg::Dir, b"hello world").unwrap());
+    println!("{}", k.encrypt(JweAlg::Sm4Skw, b"hello world").unwrap());
+}
+
+impl Sm4Jwk {
+    pub fn key_id(mut self, key_id: &str) -> Self {
+        self.key_id = Some(key_id.to_string());
+        self
+    }
+
+    pub fn encrypt(&self, alg: JweAlg, message: &[u8]) -> XResult<String> {
+        if alg == JweAlg::Sm2Pke {
+            return simple_error!("SM2PKE is not supported");
+        }
+
+        let encrypted_temp_key;
+        let temp_key: [u8; 16];
+        if alg == JweAlg::Sm4Skw {
+            temp_key = rand::random();
+            encrypted_temp_key = encode_url_safe_no_pad(&encrypt_sm4ske(&self.sm4key, &temp_key)?);
+        } else {
+            temp_key = [0_u8; 16];
+            encrypted_temp_key = "".to_string();
+        }
+
+        let jwe_header = JweHeader {
+            kid: self.key_id.clone(),
+            enc: SM4GCM.to_string(),
+            alg: alg.get_name().to_string(),
+        };
+
+        let header_str = opt_result!(
+            serde_json::to_string(&jwe_header),
+            "serialize header failed: {}"
+        );
+        let header_base64 = encode_url_safe_no_pad(header_str.as_bytes());
+
+        let nonce: [u8; 12] = rand::random();
+
+        let mut encryptor;
+        if alg == JweAlg::Sm4Skw {
+            let temp_sm4_key = Sm4Key::from_slice(&temp_key)?;
+            encryptor = Sm4GcmStreamEncryptor::new(&temp_sm4_key, &nonce);
+        } else {
+            encryptor = Sm4GcmStreamEncryptor::new(&self.sm4key, &nonce);
+        }
+        encryptor.init_adata(header_base64.as_bytes());
+
+        let mut ciphertext = encryptor.update(message);
+        let (enc2, tag) = encryptor.finalize();
+        ciphertext.extend_from_slice(&enc2);
+
+        Ok(format!(
+            "{}.{}.{}.{}.{}",
+            header_base64,
+            encrypted_temp_key,
+            encode_url_safe_no_pad(&nonce),
+            encode_url_safe_no_pad(&ciphertext),
+            encode_url_safe_no_pad(&tag)
+        ))
+    }
+}
+
+fn encrypt_sm4ske(sm4key: &Sm4Key, key: &[u8]) -> XResult<Vec<u8>> {
+    let nonce: [u8; 12] = rand::random();
+    let encrypted_key = sm4_gcm_encrypt(sm4key, &nonce, key);
+    let mut ske = nonce.to_vec();
+    ske.extend_from_slice(&encrypted_key);
+    Ok(ske)
+}
+
+fn encode_url_safe_no_pad(bytes: &[u8]) -> String {
+    URL_SAFE_NO_PAD.encode(bytes)
+}