kbs_protocol: Update protocol to v0.2.0

This patch updates KBS protocol to v0.2.0. The change mainly includes
1. Replace RSA-PKCS1v15 to ECDH-ES-A256KW. The former algorithm is not
declared as deprecated in

https://www.ietf.org/archive/id/draft-madden-jose-deprecate-none-rsa15-00.html#section-1.2

Also, some fixups to make the KBS protocol's Response fully compatible
with JWE standard are made, including explicitly parse `tag` in the
flattened JSON serialization.

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>

Cargo.toml

@@ -34,8 +34,12 @@ ctr = "0.9.2"
 env_logger = "0.11.5"
 hex = "0.4.3"
 hmac = "0.12.1"
-jwt-simple = { version = "0.12", default-features = false, features = ["pure-rust"] }
-kbs-types = "0.7.0"
+jwt-simple = { version = "0.12", default-features = false, features = [
+    "pure-rust",
+] }
+# TODO: change this rev to official repo
+# Once https://github.com/virtee/kbs-types/pull/53 gets merged
+kbs-types = { git = "https://github.com/Xynnn007/kbs-types.git", rev = "f56c565" }
 lazy_static = "1.5.0"
 log = "0.4.22"
 nix = "0.29"

attestation-agent/attestation-agent/src/token/kbs.rs

@@ -40,7 +40,7 @@ impl GetToken for KbsTokenGetter {
         let (token, tee_keypair) = client.get_token().await?;
         let message = Message {
             token: token.content,
-            tee_keypair: tee_keypair.to_pkcs1_pem()?.to_string(),
+            tee_keypair: tee_keypair.to_pem()?.to_string(),
         };
 
         let res = serde_json::to_vec(&message)?;

attestation-agent/kbs_protocol/src/builder.rs

@@ -91,7 +91,7 @@ impl<T> KbsClientBuilder<T> {
         }
 
         let tee_key = match self.tee_key {
-            Some(key) => TeeKeyPair::from_pkcs1_pem(&key[..]).context("read tee public key")?,
+            Some(key) => TeeKeyPair::from_pem(&key[..]).context("read tee public key")?,
             None => TeeKeyPair::new()?,
         };
 

attestation-agent/kbs_protocol/src/client/mod.rs

@@ -48,7 +48,7 @@ pub struct KbsClient<T> {
     pub(crate) token: Option<Token>,
 }
 
-pub const KBS_PROTOCOL_VERSION: &str = "0.1.1";
+pub const KBS_PROTOCOL_VERSION: &str = "0.2.0";
 
 pub const KBS_GET_RESOURCE_MAX_ATTEMPT: u64 = 3;
 

attestation-agent/kbs_protocol/src/keypair.rs

@@ -3,80 +3,160 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-use anyhow::{Context, Result};
+use anyhow::{anyhow, bail, Context, Result};
 
 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
 use crypto::{
+    ec::{Curve, EcKeyPair, KeyWrapAlgorithm},
     rsa::{PaddingMode, RSAKeyPair},
-    WrapType,
 };
-use kbs_types::{Response, TeePubKey};
-use serde::Deserialize;
+use kbs_types::{ProtectedHeader, Response, TeePubKey};
+use log::warn;
 use zeroize::Zeroizing;
 
 #[derive(Clone, Debug)]
 pub struct TeeKeyPair {
-    keypair: RSAKeyPair,
+    key: TeeKey,
+}
+
+#[derive(Clone, Debug)]
+pub enum TeeKey {
+    Rsa(Box<RSAKeyPair>),
+    Ec(Box<EcKeyPair>),
 }
 
 impl TeeKeyPair {
+    /// Create a new Tee key pair. We by default to use EC key pair.
     pub fn new() -> Result<Self> {
-        Ok(Self {
-            keypair: RSAKeyPair::new()?,
-        })
+        let key = TeeKey::Ec(Box::default());
+        Ok(Self { key })
     }
 
     /// Export TEE public key as specific structure.
     pub fn export_pubkey(&self) -> Result<TeePubKey> {
-        let k_mod = URL_SAFE_NO_PAD.encode(self.keypair.n());
-        let k_exp = URL_SAFE_NO_PAD.encode(self.keypair.e());
+        match &self.key {
+            TeeKey::Rsa(key) => {
+                let k_mod = URL_SAFE_NO_PAD.encode(key.n());
+                let k_exp = URL_SAFE_NO_PAD.encode(key.e());
 
-        Ok(TeePubKey::RSA {
-            alg: PaddingMode::PKCS1v15.as_ref().to_string(),
-            k_mod,
-            k_exp,
-        })
+                Ok(TeePubKey::RSA {
+                    alg: PaddingMode::OAEP.as_ref().to_string(),
+                    k_mod,
+                    k_exp,
+                })
+            }
+            TeeKey::Ec(key) => {
+                let x = URL_SAFE_NO_PAD.encode(key.x()?);
+                let y = URL_SAFE_NO_PAD.encode(key.y()?);
+
+                Ok(TeePubKey::EC {
+                    crv: Curve::P256.as_ref().to_string(),
+                    alg: KeyWrapAlgorithm::EcdhEsA256Kw.as_ref().to_string(),
+                    x,
+                    y,
+                })
+            }
+        }
     }
 
     #[inline]
-    pub fn decrypt(&self, mode: PaddingMode, cipher_text: Vec<u8>) -> Result<Vec<u8>> {
-        self.keypair.decrypt(mode, cipher_text)
+    pub fn unwrap_cek(&self, header: &ProtectedHeader, cipher_text: Vec<u8>) -> Result<Vec<u8>> {
+        #[allow(deprecated)]
+        if &header.alg[..] == PaddingMode::PKCS1v15.as_ref() {
+            warn!("Use deprecated Rsa PKCSv1.5 algorithm!");
+            let TeeKey::Rsa(key) = &self.key else {
+                bail!("Unmatched key. Must be RSA key");
+            };
+
+            let cek = key.decrypt(PaddingMode::OAEP, cipher_text)?;
+            Ok(cek)
+        } else if &header.alg[..] == PaddingMode::OAEP.as_ref() {
+            let TeeKey::Rsa(key) = &self.key else {
+                bail!("Unmatched key. Must be RSA key");
+            };
+
+            let cek = key.decrypt(PaddingMode::OAEP, cipher_text)?;
+            Ok(cek)
+        } else if &header.alg[..] == KeyWrapAlgorithm::EcdhEsA256Kw.as_ref() {
+            let epk = header
+                .other_fields
+                .get("epk")
+                .ok_or(anyhow!("Invalid JWE ProtectedHeader. Without `epk`"))?;
+            let crv = epk
+                .get("crv")
+                .ok_or(anyhow!("Invalid JWE ProtectedHeader. Without `crv`"))?
+                .as_str()
+                .ok_or(anyhow!(
+                    "Invalid JWE ProtectedHeader. `crv` is not a string"
+                ))?;
+
+            let x = epk
+                .get("x")
+                .ok_or(anyhow!("Invalid JWE ProtectedHeader. Without `x`"))?
+                .as_str()
+                .ok_or(anyhow!("Invalid JWE ProtectedHeader. `x` is not a string"))?;
+
+            let x = URL_SAFE_NO_PAD.decode(x)?;
+
+            let y = epk
+                .get("y")
+                .ok_or(anyhow!("Invalid JWE ProtectedHeader. Without `y`"))?
+                .as_str()
+                .ok_or(anyhow!("Invalid JWE ProtectedHeader. `y` is not a string"))?;
+
+            let y = URL_SAFE_NO_PAD.decode(y)?;
+
+            let TeeKey::Ec(key) = &self.key else {
+                bail!("Unmatched key. Must be RSA key");
+            };
+
+            if crv != key.curve().as_ref() {
+                bail!("Unmatched curve: {}", crv);
+            }
+
+            let cek = key.unwrap_key(cipher_text, x, y, KeyWrapAlgorithm::EcdhEsA256Kw)?;
+            Ok(cek)
+        } else {
+            bail!("Unsupported algorithm: {}", header.alg)
+        }
     }
 
     #[inline]
-    pub fn from_pkcs1_pem(pem: &str) -> Result<Self> {
-        let keypair = RSAKeyPair::from_pkcs1_pem(pem)?;
-        Ok(Self { keypair })
+    pub fn from_pem(pem: &str) -> Result<Self> {
+        if let Ok(keypair) = RSAKeyPair::from_pkcs1_pem(pem) {
+            return Ok(Self {
+                key: TeeKey::Rsa(Box::new(keypair)),
+            });
+        }
+
+        let keypair = EcKeyPair::from_pkcs8_pem(pem)
+            .context("private key is not RSA (PKCS#1) nor EC P256 (PKCS#8)")?;
+        Ok(Self {
+            key: TeeKey::Ec(Box::new(keypair)),
+        })
     }
 
     #[inline]
-    pub fn to_pkcs1_pem(&self) -> Result<Zeroizing<String>> {
-        self.keypair.to_pkcs1_pem()
+    pub fn to_pem(&self) -> Result<Zeroizing<String>> {
+        match &self.key {
+            TeeKey::Rsa(keypair) => keypair.to_pkcs1_pem(),
+            TeeKey::Ec(keypair) => keypair.to_pkcs8_pem(),
+        }
     }
 
     pub fn decrypt_response(&self, response: Response) -> Result<Vec<u8>> {
-        // deserialize the jose header and check that the key type matches
-        let protected: ProtectedHeader = serde_json::from_str(&response.protected)?;
-        let padding_mode = PaddingMode::try_from(&protected.alg[..])
-            .context("Unsupported padding mode for wrapped key")?;
-
         // unwrap the wrapped key
-        let wrapped_symkey: Vec<u8> = URL_SAFE_NO_PAD.decode(&response.encrypted_key)?;
-        let symkey = self.decrypt(padding_mode, wrapped_symkey)?;
-
-        let iv = URL_SAFE_NO_PAD.decode(&response.iv)?;
-        let ciphertext = URL_SAFE_NO_PAD.decode(&response.ciphertext)?;
+        let cek = self.unwrap_cek(&response.protected, response.encrypted_key)?;
 
-        let plaintext = crypto::decrypt(Zeroizing::new(symkey), ciphertext, iv, protected.enc)?;
+        let aad = response.protected.generate_aad()?;
+        let plaintext = crypto::aes256gcm_decrypt(
+            Zeroizing::new(cek),
+            response.ciphertext,
+            response.iv,
+            aad,
+            response.tag,
+        )?;
 
         Ok(plaintext)
     }
 }
-
-#[derive(Deserialize)]
-struct ProtectedHeader {
-    /// enryption algorithm for encrypted key
-    alg: String,
-    /// encryption algorithm for payload
-    enc: WrapType,
-}

attestation-agent/kbs_protocol/src/token_provider/aa/mod.rs

@@ -61,7 +61,7 @@ impl TokenProvider for AATokenProvider {
         })?;
         let token = Token::new(message.token)
             .map_err(|e| Error::AATokenProvider(format!("deserialize token failed: {e:?}")))?;
-        let tee_keypair = TeeKeyPair::from_pkcs1_pem(&message.tee_keypair).map_err(|e| {
+        let tee_keypair = TeeKeyPair::from_pem(&message.tee_keypair).map_err(|e| {
             Error::AATokenProvider(format!("deserialize tee keypair failed: {e:?}"))
         })?;
         Ok((token, tee_keypair))