KBS: Update KBS protocol to 0.2.0 to fix JWE

Fixes #583.

Due to RFC 7516, the JWE AEAD Auth Tag should be expcilitly be included
inside the `tag` part. Before this commit, the tag is actually included
as the suffix of the `ciphertext`.

We fix this by expcilitly extract the tag and include it into the jwe
body.

Also, we fix the AAD calculation logic, s.t. derived from
ProtectedHeader which is also specifiled by RFC7516. This should be
align with the guest-components side.

This change will make the kbs_client not able to connect to the KBS.
Thus we update the KBS protocol version from 0.1.1 to 0.2.0.

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

Cargo.toml

@@ -33,7 +33,7 @@ jwt-simple = { version = "0.12", default-features = false, features = [
     "pure-rust",
 ] }
 kbs_protocol = { git = "https://github.com/confidential-containers/guest-components.git", rev = "e6999a3c0fd877dae9e68ea78b8b483062db32b8", default-features = false }
-kbs-types = "0.7.0"
+kbs-types = { git = "https://github.com/virtee/kbs-types.git", rev = "a704036" }
 kms = { git = "https://github.com/confidential-containers/guest-components.git", rev = "e6999a3c0fd877dae9e68ea78b8b483062db32b8", default-features = false }
 jsonwebtoken = { version = "9", default-features = false }
 log = "0.4.17"

kbs/Cargo.toml

@@ -83,6 +83,7 @@ attestation-service = { path = "../attestation-service", default-features = fals
 ], optional = true }
 
 [dev-dependencies]
+josekit = "0.10.0"
 tempfile.workspace = true
 rstest.workspace = true
 

kbs/docs/kbs_attestation_protocol.md

@@ -191,6 +191,7 @@ payload that follows the [JSON Web Encryption](https://www.rfc-editor.org/rfc/rf
 {
     "protected": "$jose_header",
     "encrypted_key": "$encrypted_key",
+    "aad": "$aad",
     "iv": "$iv",
     "ciphertext": "$ciphertext",
     "tag": "$tag"
@@ -204,6 +205,7 @@ The above JWE JSON fields are defined as follows:
 let jose_header_string = format!(r#"{{"alg": "{}","enc": "{}"}}"#, alg, enc);
 let jose_header = base64_url::encode(&jose_header_string);
 let encrypted_key = base64_url::encode(enc_kbs_symkey);
+let aad = base64_url::encode(additional_authenticated_data);
 let iv = base64_url::encode(initialization_vector);
 let ciphertext = base64_url::encode(response_output);
 
@@ -212,13 +214,13 @@ tag = base64_url::encode(authentication_tag);
 
 ```
 
-- `alg`
+- `protected.alg`
 
 Algorithm used to encrypt the encryption key at `encrypted_key`.
 Since the key is encrypted using the HW-TEE public key, `alg` must be the same
 value as described in the [`Attestation`](#attestation)'s `tee-pubkey` field.
 
-- `enc`
+- `protected.enc`
 
 Encryption algorithm used to encrypt the output of the KBS service API.
 
@@ -227,6 +229,12 @@ Encryption algorithm used to encrypt the output of the KBS service API.
 The output of the KBS service API. It must be encrypted with the KBS-generated
 ephemeral key.
 
+- `aad` (Required if AEAD is used)
+
+An input to an AEAD operation that is integrity protected but not encrypted.
+Due to [JSON Web Encryption](https://www.rfc-editor.org/rfc/rfc7516), AAD field
+should be calculated by `ASCII(BASE64URL(UTF8(JWE Protected Header)))`
+
 - `iv`
 
 The input to a cryptographic primitive is used to provide the initial state.
@@ -239,6 +247,11 @@ The encrypted symmetric key is used to encrypt `ciphertext`.
 This key is encrypted with the HW-TEE's public key, using the algorithm defined
 in `alg`.
 
+- `tag`
+
+The authentication tag is used to authenticate the ciphertext. If the algorithm
+described by `enc` used does not need it, this field is left blank.
+
 ## Key Format
 
 ### Public Key

kbs/src/attestation/backend.rs

@@ -25,8 +25,8 @@ use super::{
 };
 
 static KBS_MAJOR_VERSION: u64 = 0;
-static KBS_MINOR_VERSION: u64 = 1;
-static KBS_PATCH_VERSION: u64 = 1;
+static KBS_MINOR_VERSION: u64 = 2;
+static KBS_PATCH_VERSION: u64 = 0;
 
 lazy_static! {
     static ref VERSION_REQ: VersionReq = {

kbs/src/jwe.rs

@@ -2,18 +2,20 @@
 // Licensed under the Apache License, Version 2.0, see LICENSE for details.
 // SPDX-License-Identifier: Apache-2.0
 
-use aes_gcm::{aead::Aead, Aes256Gcm, KeyInit, Nonce};
+use std::collections::BTreeMap;
+
+use aes_gcm::{aead::AeadMutInPlace, Aes256Gcm, KeyInit, Nonce};
 use anyhow::{anyhow, bail, Context, Result};
 use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
-use kbs_types::{Response, TeePubKey};
+use kbs_types::{ProtectedHeader, Response, TeePubKey};
 use rand::{rngs::OsRng, Rng};
 use rsa::{BigUint, Pkcs1v15Encrypt, RsaPublicKey};
-use serde_json::json;
 
+// TODO: Use RSA-OEAP rather than PKCS1v15
 const RSA_ALGORITHM: &str = "RSA1_5";
 const AES_GCM_256_ALGORITHM: &str = "A256GCM";
 
-pub fn jwe(tee_pub_key: TeePubKey, payload_data: Vec<u8>) -> Result<Response> {
+pub fn jwe(tee_pub_key: TeePubKey, mut payload_data: Vec<u8>) -> Result<Response> {
     let TeePubKey::RSA { alg, k_mod, k_exp } = tee_pub_key else {
         bail!("Only RSA key is support for TEE pub key")
     };
@@ -25,11 +27,19 @@ pub fn jwe(tee_pub_key: TeePubKey, payload_data: Vec<u8>) -> Result<Response> {
     let mut rng = rand::thread_rng();
 
     let aes_sym_key = Aes256Gcm::generate_key(&mut OsRng);
-    let cipher = Aes256Gcm::new(&aes_sym_key);
+    let mut cipher = Aes256Gcm::new(&aes_sym_key);
     let iv = rng.gen::<[u8; 12]>();
     let nonce = Nonce::from_slice(&iv);
-    let encrypted_payload_data = cipher
-        .encrypt(nonce, payload_data.as_slice())
+    let protected = ProtectedHeader {
+        alg: RSA_ALGORITHM.to_string(),
+        enc: AES_GCM_256_ALGORITHM.to_string(),
+        other_fields: BTreeMap::new(),
+    };
+
+    let aad = protected.generate_aad().context("Generate JWE AAD")?;
+
+    let tag = cipher
+        .encrypt_in_place_detached(nonce, &aad, &mut payload_data)
         .map_err(|e| anyhow!("AES encrypt Resource payload failed: {e}"))?;
 
     let k_mod = URL_SAFE_NO_PAD
@@ -43,23 +53,65 @@ pub fn jwe(tee_pub_key: TeePubKey, payload_data: Vec<u8>) -> Result<Response> {
 
     let rsa_pub_key =
         RsaPublicKey::new(n, e).context("Building RSA key from modulus and exponent failed")?;
-    let sym_key: &[u8] = aes_sym_key.as_slice();
-    let wrapped_sym_key = rsa_pub_key
-        .encrypt(&mut rng, Pkcs1v15Encrypt, sym_key)
+    let encrypted_key = rsa_pub_key
+        .encrypt(&mut rng, Pkcs1v15Encrypt, aes_sym_key.as_slice())
         .context("RSA encrypt sym key failed")?;
 
-    let protected_header = json!(
-    {
-       "alg": RSA_ALGORITHM.to_string(),
-       "enc": AES_GCM_256_ALGORITHM.to_string(),
-    });
-
     Ok(Response {
-        protected: serde_json::to_string(&protected_header)
-            .context("serde protected_header failed")?,
-        encrypted_key: URL_SAFE_NO_PAD.encode(wrapped_sym_key),
-        iv: URL_SAFE_NO_PAD.encode(iv),
-        ciphertext: URL_SAFE_NO_PAD.encode(encrypted_payload_data),
-        tag: "".to_string(),
+        protected,
+        encrypted_key,
+        iv: iv.into(),
+        ciphertext: payload_data,
+        aad: None,
+        tag: tag.to_vec(),
     })
 }
+
+#[cfg(test)]
+mod tests {
+    use core::assert_eq;
+
+    use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
+    use josekit::jwe::{alg::rsaes::RsaesJweAlgorithm::Rsa1_5, JweContext, JweHeader};
+    use kbs_types::TeePubKey;
+    use openssl::rsa::Rsa;
+
+    use super::jwe;
+
+    #[test]
+    fn jwe_compability() {
+        let test_data = b"this is a test data";
+
+        // Generate a 4096-bit RSA key pair
+        let rsa_key = Rsa::generate(4096).unwrap();
+        let k_mod = URL_SAFE_NO_PAD.encode(rsa_key.n().to_vec());
+        let k_exp = URL_SAFE_NO_PAD.encode(rsa_key.e().to_vec());
+        let tee_key = TeePubKey::RSA {
+            alg: "RSA1_5".into(),
+            k_mod,
+            k_exp,
+        };
+
+        // Generate a JWE response
+        let response = jwe(tee_key, test_data.to_vec()).unwrap();
+        let response_string = serde_json::to_string(&response).unwrap();
+
+        // Decrypt with josekit crate
+        let decrypter = Rsa1_5
+            .decrypter_from_pem(rsa_key.private_key_to_pem().unwrap())
+            .unwrap();
+        let mut header = JweHeader::new();
+        header.set_token_type("JWT");
+        header.set_content_encryption("A256GCM");
+        let context = JweContext::new();
+        let (decrypted_data, header) = context
+            .deserialize_json(&response_string, &decrypter)
+            .unwrap();
+        assert_eq!(decrypted_data, test_data);
+
+        let mut jwe_header = JweHeader::new();
+        jwe_header.set_claim("alg", Some("RSA1_5".into())).unwrap();
+        jwe_header.set_claim("enc", Some("A256GCM".into())).unwrap();
+        assert_eq!(header, jwe_header);
+    }
+}