initial version of the tool to create TEE Attestation Payload

Signed-off-by: Wojciech Ozga <woz@zurich.ibm.com>

tools/attestation_payload/Cargo.toml

@@ -0,0 +1,18 @@
+[package]
+name = "ace-tap"
+version = "0.1.0"
+authors = ["Wojciech Ozga <woz@zurich.ibm.com>"]
+description = "Tool to create TEE attestation payload for a confidential VM"
+edition = "2021"
+
+[dependencies]
+clap = { version="4", features = ["derive"] } # for command line argument parsing
+
+byteorder = "1.5" # to serialize numbers using little/big endianness
+sha2 = "0.10" # to calculate integrity measurements of kernel, initramfs etc.
+hmac = "0.12" # to calculate HMAC over encrypted payload
+rsa = "0.9" # to create lockboxes: encrypt symetric key using public keys of target TEEs
+rand = "0.8" # to generate symmetric key used to encrypted payload
+aes-gcm = "0.10.3" # for symmetric encryption of payload
+
+thiserror = "1.0"  # provides macros that help removing boilerplate code in the rust error handling

tools/attestation_payload/README.md

@@ -0,0 +1,3 @@
+# Usage
+
+Generate symmetric key

tools/attestation_payload/src/append.rs

@@ -0,0 +1,24 @@
+// SPDX-FileCopyrightText: 2024 IBM Corporation
+// SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
+// SPDX-License-Identifier: Apache-2.0
+use crate::error::Error;
+
+pub fn append_tap(input_file: String, tap_file: String, output_file: Option<String>) -> Result<(), Error> {
+    use std::fs::OpenOptions;
+    use byteorder::{LittleEndian, WriteBytesExt};
+
+    let output_file_name = match output_file {
+        Some(f) if input_file != f => {
+            std::fs::copy(input_file, f.clone())?;
+            f
+        }
+        Some(f) => f,
+        None => input_file,
+    };
+    let mut file1 = OpenOptions::new().write(true).append(true).open(output_file_name)?;
+    let mut file2 = OpenOptions::new().read(true).open(tap_file)?;
+    let tap_size_in_bytes = std::io::copy(&mut file2, &mut file1)?;
+    file1.write_u32::<LittleEndian>(crate::constants::ACE_TAP_MAGIC).unwrap();
+    file1.write_u16::<LittleEndian>(tap_size_in_bytes.try_into().unwrap()).unwrap();
+    Ok(())
+}

tools/attestation_payload/src/constants.rs

@@ -0,0 +1,30 @@
+// SPDX-FileCopyrightText: 2024 IBM Corporation
+// SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
+// SPDX-License-Identifier: Apache-2.0
+
+pub const ACE_TAP_MAGIC: u32 = 0xACEACE00;
+
+#[repr(u8)]
+pub enum TapDigestEntryType {
+    Kernel,
+    KernelCommandLine,
+    Initramfs
+}
+
+#[repr(u8)]
+pub enum TapDigestAlgorithm {
+    Sha512,
+}
+
+impl TapDigestAlgorithm {
+    pub fn digest_size(&self) -> u16 {
+        match self {
+            &Self::Sha512 => 512/8,
+        }
+    }
+}
+
+#[repr(u8)]
+pub enum TapLockboxAlgorithm {
+    Rsa_2048_Sha256_OASP = 0,
+}

tools/attestation_payload/src/error.rs

@@ -0,0 +1,24 @@
+// SPDX-FileCopyrightText: 2024 IBM Corporation
+// SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
+// SPDX-License-Identifier: Apache-2.0
+
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("IO Error")]
+    IoError(#[from] std::io::Error),
+    #[error("PKCS1 Error")]
+    Pkcs1Error(#[from] rsa::pkcs1::Error),
+    #[error("RSA Error")]
+    RsaError(#[from] rsa::Error),
+    #[error("Invalid parameter")]
+    InvalidParameter(String),
+
+    #[error("Int casting Error")]
+    TryFromIntErr(#[from] std::num::TryFromIntError),
+
+    #[error("Symmetric key size does not match the HMAC function")]
+    InvalidSizeOfSymmetricKey(),
+
+    #[error("Cannot open file {0}")]
+    CannotOpenFile(String)
+}

tools/attestation_payload/src/generate.rs

@@ -0,0 +1,123 @@
+// SPDX-FileCopyrightText: 2024 IBM Corporation
+// SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
+// SPDX-License-Identifier: Apache-2.0
+use crate::error::Error;
+use byteorder::{LittleEndian, WriteBytesExt};
+use std::io::{Write};
+use std::fs::OpenOptions;
+use sha2::{Digest};
+use rand::RngCore;
+// use hex_literal::hex;
+
+use crate::constants::TapDigestAlgorithm;
+use crate::constants::TapDigestEntryType;
+use crate::constants::TapLockboxAlgorithm;
+
+pub fn generate_tap(kernel_file: String, kernel_commandline: String, initramfs_file: String, confidential_vm_secrets: Vec<(usize, String)>, tee_public_keys_files: Vec<String>, output_file: String) -> Result<(), Error> {
+    // Sanity check:
+    // - number of secrets must be lower than 256
+    if confidential_vm_secrets.len() >= 256 {
+        return Err(Error::InvalidParameter(format!("Confidential VM can receive maximum 256 secrets")));
+    }
+    // - number of lockboxes must be lower than 1024
+    if tee_public_keys_files.len() >= 1024 {
+        return Err(Error::InvalidParameter(format!("Confidential VM TAP support max 1024 lockboxes")));
+    }
+
+    // symmetric key will be used to encrypt payload and calculate hmac
+    let mut symmetric_key = [0u8; 32];
+    rand::thread_rng().fill_bytes(&mut symmetric_key);
+
+    // generate lockboxes
+    let mut lockboxes = vec![];
+    // number of lock boxes
+    lockboxes.write_u16::<LittleEndian>(tee_public_keys_files.len().try_into()?)?;
+    for public_key_file in tee_public_keys_files.iter() {
+        // lock box entry: entry size; algorithm; ciphertext
+        let mut ciphertext = encrypt_rsa_2048_sha256_oasp(&symmetric_key, public_key_file.to_string())?;
+        lockboxes.write_u16::<LittleEndian>((8 + ciphertext.len()) as u16)?;
+        lockboxes.write_u8(TapLockboxAlgorithm::Rsa_2048_Sha256_OASP as u8)?;
+        lockboxes.append(&mut ciphertext);
+    }
+
+    let mut payload = vec![];
+    // number of digests
+    payload.write_u8(3)?;
+    // first digest: entry size; entry type; algorithm; digest 
+    payload.write_u16::<LittleEndian>(2+TapDigestAlgorithm::Sha512.digest_size())?;
+    payload.write_u8(TapDigestEntryType::Kernel as u8)?;
+    payload.write_u8(TapDigestAlgorithm::Sha512 as u8)?;
+    payload.append(&mut measure_file_integrity(kernel_file)?);
+    // second digest
+    payload.write_u16::<LittleEndian>(2+TapDigestAlgorithm::Sha512.digest_size())?;
+    payload.write_u8(TapDigestEntryType::KernelCommandLine as u8)?;
+    payload.write_u8(TapDigestAlgorithm::Sha512 as u8)?;
+    payload.append(&mut measure_string_integrity(kernel_commandline)?);
+    // third digest
+    payload.write_u16::<LittleEndian>(2+TapDigestAlgorithm::Sha512.digest_size())?;
+    payload.write_u8(TapDigestEntryType::Initramfs as u8)?;
+    payload.write_u8(TapDigestAlgorithm::Sha512 as u8)?;
+    payload.append(&mut measure_file_integrity(initramfs_file)?);
+    // number of confidential VM's secrets
+    payload.write_u8(confidential_vm_secrets.len().try_into()?)?;
+    for (key, value) in confidential_vm_secrets.iter() {
+        // each confidential VM's secret entry consist of: entry size; key; value
+        payload.write_u16::<LittleEndian>((8 + value.as_bytes().len()) as u16)?;
+        payload.write_u64::<LittleEndian>(*key as u64)?;
+        payload.append(&mut value.as_bytes().to_vec());
+    }
+    // payload's HMAC
+    let mut payload_hmac = hmac_sha512(&symmetric_key, &payload)?;
+
+    // encrypt payload and hmac
+    let mut serialized_payload_and_hmac = vec![];
+    serialized_payload_and_hmac.write_u16::<LittleEndian>(payload.len().try_into()?)?;
+    serialized_payload_and_hmac.append(&mut payload);
+    serialized_payload_and_hmac.write_u16::<LittleEndian>(payload_hmac.len().try_into()?)?;
+    serialized_payload_and_hmac.append(&mut payload_hmac);
+
+    use aes_gcm::{Aes256Gcm, KeyInit, AeadInPlace, Key, AeadCore};
+    use rand::rngs::OsRng;
+    let key: Key<Aes256Gcm> = symmetric_key.into();
+    let cipher = Aes256Gcm::new(&key);
+    let nonce = Aes256Gcm::generate_nonce(&mut OsRng);
+    cipher.encrypt_in_place(&nonce, b"", &mut serialized_payload_and_hmac).unwrap();
+
+    // write the entire TAP to the output file
+    let mut output = OpenOptions::new().create_new(true).write(true).append(false).open(output_file.clone()).map_err(|_| Error::CannotOpenFile(output_file))?;
+    output.write(&lockboxes)?;
+    output.write(&serialized_payload_and_hmac)?;
+
+    Ok(())
+}
+
+fn encrypt_rsa_2048_sha256_oasp(value: &[u8], public_key_file: String) -> Result<Vec<u8>, Error> {
+    use rsa::pkcs1::DecodeRsaPublicKey;
+    let public_key_pem: Vec<u8> = std::fs::read(public_key_file.clone()).map_err(|_| Error::CannotOpenFile(public_key_file))?;
+    let public_key = rsa::RsaPublicKey::from_pkcs1_pem(&String::from_utf8_lossy(&public_key_pem))?;
+    let padding = rsa::Oaep::new::<sha2::Sha256>();
+    let encrypted_data = public_key.encrypt(&mut rand::thread_rng(), padding, value)?;
+    Ok(encrypted_data)
+}
+
+fn hmac_sha512(hmac_key: &[u8], payload: &[u8]) -> Result<Vec<u8>, Error> {
+    use hmac::{Mac};
+    let mut hmac = hmac::Hmac::<sha2::Sha512>::new_from_slice(hmac_key).map_err(|_| Error::InvalidSizeOfSymmetricKey())?;
+    hmac.update(payload);
+    Ok(hmac.finalize().into_bytes().to_vec())
+}
+
+fn measure_file_integrity(file: String) -> Result<Vec<u8>, Error> {
+    let mut hasher = sha2::Sha512::new();
+    let mut file = std::fs::File::open(file.clone()).map_err(|_| Error::CannotOpenFile(file.clone()))?;
+    std::io::copy(&mut file, &mut hasher)?;
+    let digest = hasher.finalize();
+    println!("Measured file {:?}, digest={:?}", file, digest);
+    Ok(digest.to_vec())
+}
+
+fn measure_string_integrity(value: String) -> Result<Vec<u8>, Error> {    
+    let mut hasher = sha2::Sha512::new();
+    hasher.update(value.as_bytes());
+    Ok(hasher.finalize().to_vec())
+}

tools/attestation_payload/src/main.rs

@@ -0,0 +1,65 @@
+// SPDX-FileCopyrightText: 2024 IBM Corporation
+// SPDX-FileContributor: Wojciech Ozga <woz@zurich.ibm.com>, IBM Research - Zurich
+// SPDX-License-Identifier: Apache-2.0
+use clap::{Parser, Subcommand};
+use crate::error::Error;
+
+mod error;
+mod constants;
+mod generate;
+mod append;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    #[command(subcommand)]
+    cmd: Commands
+}
+
+#[derive(Subcommand, Debug, Clone)]
+enum Commands {
+    Append {
+        #[arg(short, long)]
+        input_file: String,
+        #[arg(short, long)]
+        tap_file: String,
+        #[arg(short, long)]
+        output_file: Option<String>,
+    },
+    Generate {
+        #[arg(short='k', long)]
+        kernel_file: String,
+        #[arg(short='b', long)]
+        kernel_commandline: String,
+        #[arg(short, long)]
+        initramfs_file: String,
+        #[arg(value_parser = parse_key_val::<usize, String>)]
+        confidential_vm_secrets: Vec<(usize, String)>,
+        #[clap(short, long, value_delimiter = ' ', num_args = 1..)]
+        tee_public_keys_files: Vec<String>,
+        #[arg(short, long)]
+        output_file: String,
+    }
+}
+
+/// Parse a single key-value pair
+fn parse_key_val<T, U>(s: &str) -> Result<(T, U), Box<dyn std::error::Error + Send + Sync + 'static>>
+where
+    T: std::str::FromStr,
+    T::Err: std::error::Error + Send + Sync + 'static,
+    U: std::str::FromStr,
+    U::Err: std::error::Error + Send + Sync + 'static,
+{
+    let pos = s
+        .find('=')
+        .ok_or_else(|| format!("invalid KEY=value: no `=` found in `{s}`"))?;
+    Ok((s[..pos].parse()?, s[pos + 1..].parse()?))
+}
+
+fn main() -> Result<(), Error> {
+    Ok(match Args::parse().cmd {
+        Commands::Append {input_file, tap_file, output_file} => append::append_tap(input_file, tap_file, output_file),
+        Commands::Generate {kernel_file, kernel_commandline, initramfs_file, confidential_vm_secrets, tee_public_keys_files, output_file} => generate::generate_tap(kernel_file, kernel_commandline, initramfs_file, confidential_vm_secrets, tee_public_keys_files, output_file),
+    }?)
+}
+