// SPDX-License-Identifier: MIT
//
// Copyright IBM Corp. 2024
use std::ptr;

use zerocopy::{FromZeros, IntoBytes};

use super::{ffi, AttestationUserData, ConfigUid, UvCmd};
use crate::{Error, Result};

/// _Retrieve Attestation Measurement_ UVC
///
/// The Attestation Request has two input and three outputs.
/// ARCB and user-data are inputs for the UV.
/// Measurement, additional data, and the Configuration Unique ID are outputs generated by UV.
///
/// The Attestation Request Control Block (ARCB) is a cryptographically verified
/// and secured request to UV and user-Data is some plaintext data which is
/// going to be included in the Attestation Measurement calculation.
///
/// Measurement is a cryptographic measurement of the callers properties,
/// optional data configured by the ARCB and the user-data. If specified by the
/// ARCB, UV will add some additional Data to the measurement calculation.
/// This additional data is then returned as well.
///
/// If the Retrieve Attestation Measurement UV facility is not present,
/// UV will return invalid command rc.
///
/// # Example
///
/// ```rust,no_run
/// # use s390_pv_core::uv::UvDevice;
/// # use s390_pv_core::uv::AttestationCmd;
/// # fn main() -> s390_pv_core::Result<()> {
/// let arcb = std::fs::read("arcb")?.into();
/// let user_data = vec![0, 1, 2, 3];
/// // Hard-coded example
/// let mut cmd = AttestationCmd::new_request(arcb, Some(user_data), 64, 0)?;
/// let uv = UvDevice::open()?;
/// # uv.send_cmd(&mut cmd)?;
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
pub struct AttestationCmd {
    // all sizes are guaranteed to fit in the exchange format/UV-Call at any time
    // attestation data, these must not changed by this tooling, this is an invariant of this
    // struct, so that the raw pointer stay valid through the lifetime of this struct.
    // no mutable references are ever passed from this struct
    arcb: Box<[u8]>,
    measurement: Vec<u8>,
    additional: Option<Vec<u8>>,
    // raw IOCTL struct
    uvio_attest: ffi::uvio_attest,
}

impl AttestationCmd {
    /// Maximum size for Additional-data
    pub const ADDITIONAL_MAX_SIZE: u32 = ffi::UVIO_ATT_ADDITIONAL_MAX_LEN as u32;
    /// Maximum size for Attestation Request Control Block
    pub const ARCB_MAX_SIZE: u32 = ffi::UVIO_ATT_ARCB_MAX_LEN as u32;
    /// Maximum size for Configuration UID
    pub const CUID_SIZE: u32 = ffi::UVIO_ATT_UID_LEN as u32;
    /// Maximum size for Measurement-data
    pub const MEASUREMENT_MAX_SIZE: u32 = ffi::UVIO_ATT_MEASUREMENT_MAX_LEN as u32;
    /// Maximum size for user-data
    pub const USER_MAX_SIZE: u32 = ffi::UVIO_ATT_USER_DATA_LEN as u32;

    fn verify_size(size: u32, min_size: u32, max_size: u32, field: &'static str) -> Result<()> {
        if size < min_size {
            return Err(Error::AttDataSizeSmall { field, min_size });
        }
        if size > max_size {
            return Err(Error::AttDataSizeLarge { field, max_size });
        }

        Ok(())
    }

    fn verify_slice(val: &[u8], max_size: u32, field: &'static str) -> Result<()> {
        if val.len() > max_size as usize {
            Err(Error::AttDataSizeLarge { field, max_size })
        } else {
            Ok(())
        }
    }

    /// Creates a new [`AttestationCmd`]
    ///
    /// * `arcb` - binary attestation request
    /// * `user_data` - optional, up to 256 bytes of arbitrary data
    /// * `exp_measurement` - expected size the Attestation measurement requires
    /// * `exp_additional` - expected size of the additional data.
    ///
    /// Creates a new Retrieve Attestation Measurement UVC.
    pub fn new_request(
        arcb: Box<[u8]>,
        user_data: Option<Vec<u8>>,
        exp_measurement: u32,
        exp_additional: u32,
    ) -> Result<Self> {
        Self::verify_size(
            exp_measurement,
            1,
            Self::MEASUREMENT_MAX_SIZE,
            "Expected measurement size",
        )?;
        Self::verify_size(
            exp_additional,
            0,
            Self::ADDITIONAL_MAX_SIZE,
            "Expected additional data size",
        )?;
        Self::verify_slice(&arcb, Self::ARCB_MAX_SIZE, "Attestation request")?;
        if let Some(ref data) = user_data {
            Self::verify_slice(data, Self::USER_MAX_SIZE, "User data")?;
        }

        let (user_len, user_data) = match user_data {
            // enforced by tests before invariants
            Some(user) => (Some(user.len() as u16), {
                let mut user_data = AttestationUserData::new_zeroed();
                user_data[0..user.len()].clone_from_slice(&user);
                Some(user_data)
            }),
            None => (None, None),
        };

        let mut additional = match exp_additional {
            0 => None,
            size => Some(vec![0u8; size as usize]),
        };
        let mut measurement = vec![0u8; exp_measurement as usize];
        let uvio_attest = unsafe {
            ffi::uvio_attest::new(
                &arcb,
                &mut measurement,
                additional.as_deref_mut(),
                user_data,
                user_len,
            )
        };
        Ok(Self {
            arcb,
            measurement,
            additional,
            uvio_attest,
        })
    }

    /// Provides the additional data  calculated by UV after a successful UVC
    ///
    /// Truncates the additional data to the correct length in place.
    /// If called before a successful attestation the data in this buffer is undefined.
    pub fn additional(&mut self) -> Option<&[u8]> {
        // truncate the add size to the UV reported size
        if let Some(ref mut a) = &mut self.additional {
            a.truncate(self.uvio_attest.add_data_len as usize)
        }
        self.additional.as_deref()
    }

    /// Copies the additional data calculated by UV after a successful UVC into a Vec
    ///
    /// Truncates the additional data to the correct length.
    /// If called before a successful attestation the data in this buffer is undefined.
    pub fn additional_owned(&self) -> Option<Vec<u8>> {
        let mut additional = self.additional.clone()?;
        additional.truncate(self.uvio_attest.add_data_len as usize);
        Some(additional)
    }

    /// Provides the Configuration Unique Identifier received from UV after a successful UVC
    ///
    /// If called before a successful attestation the data in this buffer is undefined.
    pub fn cuid(&self) -> &ConfigUid {
        &self.uvio_attest.config_uid
    }

    /// Provides the attestation measurement calculated by UV after a successful UVC
    ///
    /// If called before a successful attestation the data in this buffer is undefined.
    pub fn measurement(&self) -> &[u8] {
        &self.measurement
    }

    /// Returns a reference to the request of this [`AttestationCmd`].
    pub fn arcb(&self) -> &[u8] {
        self.arcb.as_ref()
    }
}

impl UvCmd for AttestationCmd {
    const UV_IOCTL_NR: u8 = ffi::UVIO_IOCTL_ATT_NR;

    fn rc_fmt(&self, rc: u16, _rrc: u16) -> Option<&'static str> {
        match rc {
            // should not happen, uvdevice local value
            0x0101 => Some("Invalid continuation token specified"),
            0x010a => Some("Unsupported plaintext attestation flag set"),
            0x010b => Some("Unsupported measurement algorithm specified."),
            0x010c => Some("Unable to decrypt attestation request control block. Probably no valid host-key was provided"),
            0x0106 => Some("Unsupported attestation request version"),
            // should not happen, protected by AttestationCmd constructors
            0x0102 => Some("User data length is greater than 256"),
            // should not happen, uvdevice ensures this
            0x0103 => Some("Access exception recognized when accessing the attestation request control block"),
            // should not happen, uvdevice ensures this
            0x0104 => Some("Access exception recognized when accessing the measurement data area"),
            // should not happen, uvdevice ensures this
            0x0105 => Some("Access exception recognized when accessing the additional data area"),
            // should not happen, ensured by Attestation Request builder
            0x0107 => Some("Invalid attestation request length for the specified attestation request version"),
            // 0 case should not happen, ensured by Attestation Request builder
            0x0108 => Some("Number of key slots is either equal to 0 or greater than the maximum number supported by the specified attestation request version"),
            // should not happen, ensured by Attestation Request builder
            0x0109 => Some("Size of encrypted area does not match measurement length plus any optional items"),
            // should not happen, ensured by Attestation Request builder
            0x010d => Some("Measurement data length is not large enough to store measurement"),
            // should not happen, ensured by Attestation Request builder
            0x010e => Some("Additional data length not large enough to hold all requested additional data"),
            _ => None,
        }
    }

    fn data(&mut self) -> Option<&mut [u8]> {
        Some(self.uvio_attest.as_mut_bytes())
    }
}

fn opt_to_mut_ptr_u64(opt: &mut Option<&mut [u8]>) -> u64 {
    (match opt {
        Some(v) => v.as_mut_ptr(),
        None => ptr::null_mut(),
    }) as u64
}

impl ffi::uvio_attest {
    /// Create a new attestation IOCTL control block
    ///
    /// Happily converts slice lengths into u32/u16 without verifying.
    /// Therefore marked as unsafe.
    ///
    /// # SAFETY
    /// It is safe to call this function iff:
    ///  - `arcb.len() < u32::MAX`
    ///  - `additional.len() < u32::MAX`
    ///  - `user_len() <= 256`
    ///  - pointer fits into an u64
    unsafe fn new(
        arcb: &[u8],
        measurement: &mut [u8],
        mut additional: Option<&mut [u8]>,
        user: Option<AttestationUserData>,
        user_len: Option<u16>,
    ) -> Self {
        Self {
            arcb_addr: arcb.as_ptr() as u64,
            meas_addr: measurement.as_ptr() as u64,
            add_data_addr: opt_to_mut_ptr_u64(&mut additional),
            user_data: user.unwrap_or([0; 256]),
            config_uid: [0; 16],
            arcb_len: arcb.len() as u32,
            meas_len: measurement.len() as u32,
            add_data_len: additional.unwrap_or_default().len() as u32,
            user_data_len: user_len.unwrap_or_default(),
            reserved136: 0,
        }
    }
}
