routex-settlement 0.3.0

use std::future::Future;

use anyhow::anyhow;
use base64::prelude::*;
use chacha_box_ietf::{PublicKey, SecretKey, unseal};
use clerk_report::PublishedVersionEntry;
use http::HeaderValue;
use log::info;
use routex_api::keys::{Response, SettlementBoxMessage};

#[derive(Debug)]
/// Client-side key settlement handling
pub struct KeySettlement<C> {
    secret_key: SecretKey,
    server_key: Option<(PublicKey, HeaderValue)>,
    #[allow(clippy::struct_field_names)]
    core: C,
    system_version: Option<PublishedVersionEntry>,
}

pub trait KeySettlementCore {
    type Data: ?Sized;

    fn request(
        &self,
        public_key: [u8; PublicKey::size()],
        data: &Self::Data,
    ) -> impl Future<Output = anyhow::Result<Response>>;

    fn new_session(&mut self, _id: [u8; 32], _public_key: &PublicKey, _secret_key: &SecretKey) {}
}

impl<C: KeySettlementCore> KeySettlement<C> {
    pub fn new(core: C) -> Self {
        Self {
            secret_key: generate_key(),
            server_key: None,
            core,
            system_version: None,
        }
    }

    /// Settle a new key with the service
    ///
    /// # Errors
    ///
    /// Returns errors if the request fails, the attestation verification fails or the response does not meet the expectations.
    ///
    /// # Panics
    ///
    /// Panics if Base64 values are not valid HTTP header values.
    pub async fn settle(&mut self, data: &C::Data) -> anyhow::Result<()> {
        // Request a remote attestation from the TEE which authenticates the TEE's public key and our client public key
        let response = self
            .core
            .request(*self.secret_key.public_key().as_bytes(), data)
            .await?;

        if let Err(err) = verify_attestation(&response) {
            return Err(anyhow!(
                "Invalid attestation report ({err:?}): {:?}",
                response.attestation_report
            ));
        }

        let (server_public_key, session_id) = unseal(&self.secret_key, &response.chacha_box)
            .map_err(|err| {
                anyhow!("Could not unseal chacha box containing routex's public key: {err:?}")
            })
            .and_then(|box_bytes| {
                serde_json::from_slice::<SettlementBoxMessage>(&box_bytes).map_err(Into::into)
            })
            .and_then(|contents| {
                PublicKey::from_slice(&contents.public_key)
                    .map_err(|err| anyhow!("Could not deserialize routex's public key: {err:?}"))
                    .map(|public_key| (public_key, contents.session_id))
            })?;

        self.core
            .new_session(session_id, &server_public_key, &self.secret_key);

        self.server_key = Some((
            server_public_key,
            HeaderValue::from_str(&BASE64_STANDARD.encode(session_id))
                .expect("Value should be valid"),
        ));

        self.system_version = Some(response.system_version);

        Ok(())
    }

    /// System version for the currently established session
    pub fn system_version(&self) -> Option<&PublishedVersionEntry> {
        self.system_version.as_ref()
    }

    /// Seal data for the service
    ///
    /// Settles a key if none is settled.
    ///
    /// # Errors
    ///
    /// Forwards errors from [`settle`](Self::settle)
    ///
    /// # Panics
    ///
    /// Panics if [`chacha_box_ietf::seal`] panics.
    pub async fn seal(&mut self, data: &[u8], user_data: &C::Data) -> anyhow::Result<Vec<u8>> {
        self.try_f(|s| s.try_seal(data), user_data).await
    }

    /// Seal data for the service
    ///
    /// Returns [`None`] if no key is settled.
    ///
    /// # Panics
    ///
    /// Panics if [`chacha_box_ietf::seal`] panics.
    pub fn try_seal(&self, data: &[u8]) -> Option<Vec<u8>> {
        self.server_key
            .as_ref()
            .map(|(key, _)| chacha_box_ietf::seal(key, data).expect("Encrypt should work"))
    }

    /// Unseal data from the service
    ///
    /// # Errors
    ///
    /// Forwards errors from [`chacha_box_ietf::unseal`].
    pub fn unseal(&self, data: &[u8]) -> Result<Vec<u8>, chacha_box_ietf::Error> {
        chacha_box_ietf::unseal(&self.secret_key, data)
    }

    /// Return a settled session ID
    ///
    /// Settles a key if none is settled.
    ///
    /// # Errors
    ///
    /// Forwards errors from [`settle`](Self::settle)
    pub async fn session_id(&mut self, user_data: &C::Data) -> anyhow::Result<&HeaderValue> {
        self.try_f(|s: &mut KeySettlement<C>| s.try_session_id(), user_data)
            .await
    }

    /// Return a settled session ID
    ///
    /// Returns [`None`] if no key is settled.
    pub fn try_session_id(&self) -> Option<&HeaderValue> {
        self.server_key.as_ref().map(|(_, session_id)| session_id)
    }

    async fn try_f<'a, T>(
        &'a mut self,
        f: impl FnOnce(&'a mut Self) -> Option<T>,
        user_data: &C::Data,
    ) -> anyhow::Result<T> {
        if self.server_key.is_none() {
            info!("No key set, running a key settlement.");
            self.settle(user_data).await?;
        }

        Ok(f(self).expect("Key should be set"))
    }
}

#[cfg(feature = "unattested")]
fn generate_key() -> SecretKey {
    routex_keys_fixtures::fixed_client_key().into()
}

#[cfg(not(feature = "unattested"))]
fn generate_key() -> SecretKey {
    SecretKey::generate()
}

/// Verify that
/// - a YAXI-provisioned TEE created and signed the attestation report
/// - the attestation report authenticates the chacha box which seals the TEE's public key.
///   Therefore, after verification of the attestation report, unsealing the
///   box with our secret key authenticates the public keys of both parties.
/// - the attestation report authenticates the chacha box, so transitively also our public
///   key and the TEE's public key when we are able to unseal the box
///
/// # Errors
///
/// Returns an error when any step of the verification fails.
pub fn verify_attestation(response: &Response) -> std::result::Result<(), anyhow::Error> {
    use clerk_report::verification::{Requirements, RootStore, verify_report};

    let root_store = RootStore::default();
    let report = verify_report(
        &response.attestation_report,
        std::io::Cursor::new(response.vcek.as_bytes()),
        &root_store,
        &Requirements::default(),
    )
    .map_err(|err| anyhow!("Verification resulted in error: {err:?}"))?;

    // Attestation is signed by AMD, now verify that the chacha box is part of attestation
    verify_chacha_box(response, &report)?;

    // Verify that the data in `system_version` was signed by a well-known YAXI key
    response
        .system_version
        .verify_signature()
        .map_err(|err| anyhow!("Could not verify system version's signature: {err:?}"))?;

    // The reported measurement has to match the measurement specified in `system_version`. As
    // the expected measurement was signed by a YAXI key (see the step above), this guarantees
    // that the TEE is YAXI-provisioned
    if report.measurement == *response.system_version.launch_measurement {
        Ok(())
    } else {
        Err(anyhow!(
            "Reported measurement {:?} doesn't match expected measurement {:?}",
            &report.measurement,
            response.system_version.launch_measurement
        ))
    }
}

fn verify_chacha_box(
    response: &Response,
    report: &clerk_report::AttestationReport,
) -> std::result::Result<(), anyhow::Error> {
    use sha2::{Digest, Sha256};

    if &report.report_data[..32] == Sha256::digest(&response.chacha_box).as_slice() {
        Ok(())
    } else {
        Err(anyhow!(
            "Data in attestation report doesn't match chacha box"
        ))
    }
}

#[cfg(test)]
mod tests {
    use base64::prelude::*;
    use chacha_box_ietf::{PublicKey, SecretKey};
    use rand::{TryRng, rngs::SysRng};
    use routex_api::keys::Response;
    use routex_keys_fixtures::fixed_test_key_response;

    use super::{KeySettlement, KeySettlementCore};

    struct FixedResponse(routex_api::keys::Response);

    impl FixedResponse {
        fn new(response: routex_api::keys::Response) -> Self {
            Self(response)
        }
    }

    impl KeySettlementCore for FixedResponse {
        type Data = ();

        fn request(
            &self,
            _public_key: [u8; 32],
            _data: &Self::Data,
        ) -> impl std::future::Future<Output = anyhow::Result<routex_api::keys::Response>> {
            std::future::ready(Ok(self.0.clone()))
        }
    }

    struct ShouldNotSettle;

    impl KeySettlementCore for ShouldNotSettle {
        type Data = ();

        fn request(
            &self,
            _public_key: [u8; 32],
            _data: &Self::Data,
        ) -> impl std::future::Future<Output = anyhow::Result<routex_api::keys::Response>> {
            panic!("Unexpectedly called key settlement function");
            // Otherwise, rustc complains about () not being a Future
            #[allow(unreachable_code)]
            std::future::ready(Ok(fixed_test_key_response()))
        }
    }

    fn fixed_settlement() -> KeySettlement<FixedResponse> {
        settlement_with_response(fixed_test_key_response())
    }

    fn settlement_with_response(
        response: routex_api::keys::Response,
    ) -> KeySettlement<FixedResponse> {
        KeySettlement {
            secret_key: routex_keys_fixtures::fixed_client_key().into(),
            server_key: None,
            core: FixedResponse::new(response),
            system_version: None,
        }
    }

    fn assert_err_starts_with<T: std::fmt::Debug, E: std::fmt::Debug + std::fmt::Display>(
        value: Result<T, E>,
        expectation: &str,
    ) {
        if let Err(err) = value {
            let err_str = err.to_string();
            if !err_str.starts_with(expectation) {
                assert_eq!(expectation, err_str);
            }
        } else {
            panic!("Expected Err, got {value:?}");
        }
    }

    #[tokio::test]
    async fn test_seal_unseal_roundtrip() {
        let key = chacha_box_ietf::SecretKey::generate();
        let mut settlement = KeySettlement {
            secret_key: key.clone(),
            server_key: Some((key.public_key(), "session-id".try_into().unwrap())),
            core: ShouldNotSettle,
            system_version: None,
        };
        let mut data = [0u8; 42];
        SysRng.try_fill_bytes(&mut data).unwrap();

        let secret_box = settlement.seal(&data, &()).await.unwrap();
        let unsealed_data = settlement.unseal(&secret_box).unwrap();

        assert_ne!(&data[..], secret_box);
        assert_eq!(&data[..], unsealed_data);
    }

    #[tokio::test]
    async fn test_settle() {
        let mut settlement = fixed_settlement();
        settlement.settle(&()).await.unwrap();
    }

    #[tokio::test]
    async fn test_settle_invalid_vcek() {
        let mut settlement = settlement_with_response({
            let mut response = fixed_test_key_response();
            response.vcek = "invalid".into();
            response
        });
        let result = settlement.settle(&()).await;
        assert_err_starts_with(
            result,
            "Invalid attestation report (Verification resulted in error: ChainBroken)",
        );
    }

    #[tokio::test]
    async fn test_settle_invalid_attestation_report_signature() {
        let mut settlement = settlement_with_response({
            let mut response = fixed_test_key_response();
            response.attestation_report[0] = 42;
            response
        });
        let result = settlement.settle(&()).await;
        assert_err_starts_with(
            result,
            "Invalid attestation report (Verification resulted in error: ReportSignatureMismatch",
        );
    }

    #[tokio::test]
    async fn test_settle_invalid_system_version_signature() {
        let mut settlement = settlement_with_response({
            let mut response = fixed_test_key_response();
            response.system_version.signature.value[0] = 42;
            response
        });
        let result = settlement.settle(&()).await;
        assert_err_starts_with(
            result,
            "Invalid attestation report (Could not verify system version's signature: SignatureError",
        );
    }

    #[tokio::test]
    async fn test_session_id() {
        let mut settlement = fixed_settlement();
        let session_id = settlement.session_id(&()).await.unwrap();
        assert_eq!(
            session_id.to_str().unwrap(),
            BASE64_STANDARD.encode(routex_keys_fixtures::fixed_session_id())
        );
    }

    #[tokio::test]
    async fn test_session_id_settled_key() {
        let key = chacha_box_ietf::SecretKey::generate();
        let expected_session_id: http::HeaderValue = "session-id".try_into().unwrap();
        let mut settlement = KeySettlement {
            secret_key: key.clone(),
            server_key: Some((key.public_key(), expected_session_id.clone())),
            core: ShouldNotSettle,
            system_version: None,
        };

        let session_id = settlement.session_id(&()).await.unwrap();

        assert_eq!(session_id, &expected_session_id);
    }

    #[test]
    fn test_seal_settled_key() {
        let secret_key = SecretKey::from(routex_keys_fixtures::fixed_client_key());
        let server_key = Some((secret_key.public_key(), "session-id".try_into().unwrap()));

        let settlement = KeySettlement {
            secret_key,
            server_key,
            core: ShouldNotSettle,
            system_version: None,
        };
        let sealed = settlement.try_seal(&[]);
        assert!(sealed.is_some());
    }

    #[test]
    fn test_seal_not_settled_yet() {
        let settlement = KeySettlement::new(ShouldNotSettle);
        assert_eq!(None, settlement.try_seal(&[]));
    }

    #[tokio::test]
    async fn test_new_session_callback() {
        struct Core {
            session: bool,
        }

        impl KeySettlementCore for Core {
            type Data = ();

            fn request(
                &self,
                _public_key: [u8; PublicKey::size()],
                _data: &Self::Data,
            ) -> impl Future<Output = anyhow::Result<Response>> {
                std::future::ready(Ok(fixed_test_key_response()))
            }

            fn new_session(
                &mut self,
                _session_id: [u8; 32],
                _public_key: &PublicKey,
                _secret_key: &SecretKey,
            ) {
                self.session = true;
            }
        }

        let mut settlement = KeySettlement {
            secret_key: routex_keys_fixtures::fixed_client_key().into(),
            server_key: None,
            core: Core { session: false },
            system_version: None,
        };

        settlement.settle(&()).await.unwrap();

        assert!(settlement.core.session);
    }
}