From 66569618dd8ad84195aff493a5e791261d0b8497 Mon Sep 17 00:00:00 2001
From: Xynnn007 <xynnn@linux.alibaba.com>
Date: Mon, 25 Nov 2024 11:15:43 +0800
Subject: [PATCH] AA/kbs_protocol: Update protocol version to 0.2.0 to fix JWE

Per RFC7516, the AEAD's auth tag should be included inside the JWE body.
We fix this to align with trustee side

https://github.com/confidential-containers/trustee/pull/597

Signed-off-by: Xynnn007 <xynnn@linux.alibaba.com>
---
 attestation-agent/deps/crypto/src/lib.rs         |  2 ++
 .../deps/crypto/src/native/aes256gcm.rs          |  9 ++++++---
 .../deps/crypto/src/rust/aes256gcm.rs            | 16 ++++++++++------
 attestation-agent/kbs_protocol/src/client/mod.rs |  2 +-
 attestation-agent/kbs_protocol/src/keypair.rs    |  5 +++--
 5 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/attestation-agent/deps/crypto/src/lib.rs b/attestation-agent/deps/crypto/src/lib.rs
index fc78e7e02..a8606da1a 100644
--- a/attestation-agent/deps/crypto/src/lib.rs
+++ b/attestation-agent/deps/crypto/src/lib.rs
@@ -34,3 +34,5 @@ pub use asymmetric::*;
 
 mod algorithms;
 pub use algorithms::*;
+
+const AEAD_AAD: &[u8] = b"CoCo";
diff --git a/attestation-agent/deps/crypto/src/native/aes256gcm.rs b/attestation-agent/deps/crypto/src/native/aes256gcm.rs
index 73a883a4e..fc6e94410 100644
--- a/attestation-agent/deps/crypto/src/native/aes256gcm.rs
+++ b/attestation-agent/deps/crypto/src/native/aes256gcm.rs
@@ -8,6 +8,8 @@
 use anyhow::*;
 use openssl::symm::Cipher;
 
+use crate::AEAD_AAD;
+
 const TAG_LENGTH: usize = 16;
 
 pub fn decrypt(encrypted_data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>> {
@@ -17,15 +19,16 @@ pub fn decrypt(encrypted_data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>>
     }
 
     let (data, tag) = encrypted_data.split_at(encrypted_data.len() - TAG_LENGTH);
-    openssl::symm::decrypt_aead(cipher, key, Some(iv), &[], data, tag)
+    openssl::symm::decrypt_aead(cipher, key, Some(iv), AEAD_AAD, data, tag)
         .map_err(|e| anyhow!(e.to_string()))
 }
 
 pub fn encrypt(data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>> {
     let cipher = Cipher::aes_256_gcm();
     let mut tag = [0u8; TAG_LENGTH];
-    let mut ciphertext = openssl::symm::encrypt_aead(cipher, key, Some(iv), &[], data, &mut tag)
-        .map_err(|e| anyhow!(e.to_string()))?;
+    let mut ciphertext =
+        openssl::symm::encrypt_aead(cipher, key, Some(iv), AEAD_AAD, data, &mut tag)
+            .map_err(|e| anyhow!(e.to_string()))?;
     ciphertext.extend_from_slice(&tag);
     Ok(ciphertext)
 }
diff --git a/attestation-agent/deps/crypto/src/rust/aes256gcm.rs b/attestation-agent/deps/crypto/src/rust/aes256gcm.rs
index 00a4a0d55..8fb3be087 100644
--- a/attestation-agent/deps/crypto/src/rust/aes256gcm.rs
+++ b/attestation-agent/deps/crypto/src/rust/aes256gcm.rs
@@ -5,26 +5,30 @@
 
 //! This mod implements aes-256-gcm encryption & decryption.
 
-use aes_gcm::{aead::Aead, Aes256Gcm, Key, KeyInit, Nonce};
+use aes_gcm::{AeadInPlace, Aes256Gcm, Key, KeyInit, Nonce};
 use anyhow::*;
 
+use crate::AEAD_AAD;
+
 pub fn decrypt(encrypted_data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>> {
     let decrypting_key = Key::<Aes256Gcm>::from_slice(key);
     let cipher = Aes256Gcm::new(decrypting_key);
     let nonce = Nonce::from_slice(iv);
-    let plain_text = cipher
-        .decrypt(nonce, encrypted_data)
+    let mut plaintext = encrypted_data.to_vec();
+    cipher
+        .decrypt_in_place(nonce, AEAD_AAD, &mut plaintext)
         .map_err(|e| anyhow!("aes-256-gcm decrypt failed: {:?}", e))?;
 
-    Ok(plain_text)
+    Ok(plaintext)
 }
 
 pub fn encrypt(data: &[u8], key: &[u8], iv: &[u8]) -> Result<Vec<u8>> {
     let encrypting_key = Key::<Aes256Gcm>::from_slice(key);
     let cipher = Aes256Gcm::new(encrypting_key);
     let nonce = Nonce::from_slice(iv);
-    let ciphertext = cipher
-        .encrypt(nonce, data)
+    let mut ciphertext = data.to_vec();
+    cipher
+        .encrypt_in_place(nonce, AEAD_AAD, &mut ciphertext)
         .map_err(|e| anyhow!("aes-256-gcm encrypt failed: {:?}", e))?;
 
     Ok(ciphertext)
diff --git a/attestation-agent/kbs_protocol/src/client/mod.rs b/attestation-agent/kbs_protocol/src/client/mod.rs
index 5febb6911..fdb7f23db 100644
--- a/attestation-agent/kbs_protocol/src/client/mod.rs
+++ b/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;
 
diff --git a/attestation-agent/kbs_protocol/src/keypair.rs b/attestation-agent/kbs_protocol/src/keypair.rs
index 897724e3e..732a5b462 100644
--- a/attestation-agent/kbs_protocol/src/keypair.rs
+++ b/attestation-agent/kbs_protocol/src/keypair.rs
@@ -65,8 +65,9 @@ impl TeeKeyPair {
         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 mut ciphertext = URL_SAFE_NO_PAD.decode(&response.ciphertext)?;
+        let mut tag = URL_SAFE_NO_PAD.decode(&response.tag)?;
+        ciphertext.append(&mut tag);
         let plaintext = crypto::decrypt(Zeroizing::new(symkey), ciphertext, iv, protected.enc)?;
 
         Ok(plaintext)