// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

//! The crate includes the abstraction layer of Linux SEV-SNP Guest APIs and
//! definitions of data structures according to SEV-SNP specification.

use bitfield_struct::bitfield;
use zerocopy::AsBytes;
use zerocopy::FromBytes;
use zerocopy::FromZeroes;

/// Ioctl type defined by Linux.
pub const SNP_GUEST_REQ_IOC_TYPE: u8 = b'S';

/// Value for the `msg_version` member in [`SNP_GUEST_REQ_MSG_VERSION`].
/// Use 1 for now.
pub const SNP_GUEST_REQ_MSG_VERSION: u32 = 1;

/// Ioctl struct defined by Linux.
#[repr(C)]
pub struct SnpGuestRequestIoctl {
    /// Message version number (must be non-zero).
    pub msg_version: u32,
    /// Request struct address.
    pub req_data: u64,
    /// Response struct address.
    pub resp_data: u64,
    /// VMM error code.
    pub exitinfo: VmmErrorCode,
}

/// VMM error code.
#[repr(C)]
#[derive(FromZeroes)]
pub struct VmmErrorCode {
    /// Firmware error
    pub fw_error: u32,
    /// VMM error
    pub vmm_error: u32,
}

/// Request structure for the `SNP_GET_REPORT` ioctl.
/// See `MSG_REPORT_REQ` in Table 21, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
#[repr(C)]
#[derive(AsBytes, FromBytes, FromZeroes)]
pub struct SnpReportReq {
    /// Guest-provided data to be included in the attestation report.
    pub user_data: [u8; 64],
    /// The VMPL to put in the attestation report. Must be greater than
    /// or equal to the current VMPL and, at most, three.
    pub vmpl: u32,
    /// Reserved
    // TODO SNP: Support VLEK feature if needed
    pub rsvd: [u8; 28],
}

const SNP_REPORT_RESP_DATA_SIZE: usize =
    size_of::<u32>() + size_of::<u32>() + 24 + size_of::<SnpReport>();

/// The size of the response data defined by the Linux kernel.
const LINUX_SNP_REPORT_RESP_DATA_SIZE: usize = 4000;

/// Response structure for the `SNP_GET_REPORT` ioctl.
/// See `MSG_REPORT_RSP` in Table 24, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
#[repr(C)]
#[derive(AsBytes, FromBytes, FromZeroes)]
pub struct SnpReportResp {
    /// The status of key derivation operation.
    /// 0h: Success.
    /// 16h: Invalid parameters.
    /// 27h: Invalid key selection.
    pub status: u32,
    /// Size in bytes of the report.
    pub report_size: u32,
    /// Reserved
    pub _reserved0: [u8; 24],
    /// The attestation report generated by the firmware.
    pub report: SnpReport,
    /// Reserved
    pub _reserved1: [u8; LINUX_SNP_REPORT_RESP_DATA_SIZE - SNP_REPORT_RESP_DATA_SIZE],
}

static_assertions::const_assert_eq!(LINUX_SNP_REPORT_RESP_DATA_SIZE, size_of::<SnpReportResp>());

/// Size of the [`SnpReport`].
pub const SNP_REPORT_SIZE: usize = 0x4a0;

/// Size of `report_data` member in [`SnpReport`].
pub const SNP_REPORT_DATA_SIZE: usize = 64;

/// Report structure.
/// See `ATTESTATION_REPORT` in Table 22, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
#[repr(C)]
#[derive(AsBytes, FromBytes, FromZeroes)]
pub struct SnpReport {
    /// Version number of this attestation report.
    /// Set to 2h for this specification.
    pub version: u32,
    /// The guest SVN.
    pub guest_svn: u32,
    /// The guest policy.
    pub policy: u64,
    /// The family ID provided at launch.
    pub family: u128,
    /// The image ID provided at launch.
    pub image_id: u128,
    /// The request VMPL for the attestation
    /// report.
    pub vmpl: u32,
    /// The signature algorithm used to sign
    /// this report.
    pub signature_algo: u32,
    /// CurrentTcb.
    pub current_tcb: u64,
    /// Information about the platform.
    pub platform_info: u64,
    /// Flags
    pub flags: u32,
    /// Reserved
    pub _reserved0: u32,
    /// Guest-provided data.
    pub report_data: [u8; SNP_REPORT_DATA_SIZE],
    /// The measurement calculated at
    /// launch.
    pub measurement: [u8; 48],
    /// Data provided by the hypervisor at
    /// launch.
    pub host_data: [u8; 32],
    /// SHA-384 digest of the ID public key
    /// that signed the ID block provided in
    /// SNP_LAUNCH_FINISH.
    pub id_key_digest: [u8; 48],
    /// SHA-384 digest of the Author public
    /// key that certified the ID key, if
    /// provided in SNP_LAUNCH_FINISH.
    pub author_key_digest: [u8; 48],
    /// Report ID of this guest.
    pub report_id: [u8; 32],
    /// Report ID of this guest’s migration
    /// agent
    pub report_id_ma: [u8; 32],
    /// Reported TCB version used to derive
    /// the VCEK that signed this report.
    pub reported_tcb: u64,
    /// Reserved
    pub _reserved1: [u8; 24],
    /// If MaskChipId is set to 0, Identifier
    /// unique to the chip as output by
    /// GET_ID. Otherwise, set to 0h.
    pub chip_id: [u8; 64],
    /// CommittedTcb.
    pub committed_tcb: u64,
    /// The build number of CurrentVersion.
    pub current_build: u8,
    /// The minor number of CurrentVersion.
    pub current_minor: u8,
    /// The major number of CurrentVersion.
    pub current_major: u8,
    /// Reserved
    pub _reserved2: u8,
    /// The build number of CommittedVersion.
    pub committed_build: u8,
    /// The minor version of CommittedVersion.
    pub committed_minor: u8,
    /// The major version of CommittedVersion.
    pub committed_major: u8,
    /// Reserved
    pub _reserved3: u8,
    /// The CurrentTcb at the time the guest
    /// was launched or imported.
    pub launch_tcb: u64,
    /// Reserved
    pub _reserved4: [u8; 168],
    /// Signature of bytes inclusive of this report.
    pub signature: [u8; 512],
}

static_assertions::const_assert_eq!(SNP_REPORT_SIZE, size_of::<SnpReport>());

/// Request structure for the `SNP_GET_DERIVED_KEY` ioctl.
/// See `MSG_KEY_REQ` in Table 18, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
#[repr(C)]
#[derive(AsBytes, FromBytes, FromZeroes)]
pub struct SnpDerivedKeyReq {
    /// Selects the root key from which to derive the key.
    /// 0 indicates VCEK
    /// 1 indicates VMRK
    // TODO: Support VLEK feature if needed
    pub root_key_select: u32,
    /// Reserved
    pub rsvd: u32,
    /// Bitmask indicating which data will be mixed into the
    /// derived key.
    pub guest_field_select: u64,
    /// The VMPL to mix into the derived key. Must be greater
    /// than or equal to the current VMPL.
    pub vmpl: u32,
    /// The guest SVN to mix into the key. Must not exceed the
    /// guest SVN provided at launch in the ID block.
    pub guest_svn: u32,
    /// The TCB version to mix into the derived key. Must not
    /// exceed CommittedTcb.
    pub tcb_version: u64,
}

/// Indicate which guest-selectable fields will be mixed into the key.
/// See `GUEST_FIELD_SELECT` in Table 19, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
#[bitfield(u64)]
pub struct GuestFieldSelect {
    /// Indicate that the guest policy will be mixed into the key.
    pub guest_policy: bool,
    /// Indicate that the image ID of the guest will be mixed into the key.
    pub image_id: bool,
    /// Indicate the family ID of the guest will be mixed into the key.
    pub family_id: bool,
    /// Indicate the measurement of the guest during launch will be mixed into the key.
    pub measurement: bool,
    /// Indicate that the guest-provided SVN will be mixed into the key.
    pub guest_svn: bool,
    /// Indicate that the guest-provided TCB_VERSION will be mixed into the key.
    pub tcb_version: bool,
    /// Reserved
    #[bits(58)]
    pub _reserved: u64,
}

/// See `DERIVED_KEY` in Table 20, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
pub const SNP_DERIVED_KEY_SIZE: usize = 32;

/// Response structure for the `SNP_GET_DERIVED_KEY` ioctl.
/// See `MSG_KEY_RSP` in Table 20, "SEV Secure Nested Paging Firmware ABI specification", Revision 1.55.
#[repr(C)]
#[derive(AsBytes, FromBytes, FromZeroes)]
pub struct SnpDerivedKeyResp {
    /// The status of key derivation operation.
    /// 0h: Success.
    /// 16h: Invalid parameters.
    /// 27h: Invalid key selection.
    pub status: u32,
    /// Reserved
    pub _reserved: [u8; 28],
    /// The requested derived key.
    pub derived_key: [u8; SNP_DERIVED_KEY_SIZE],
}

static_assertions::const_assert_eq!(
    // The size of the response data defined by the Linux kernel.
    64,
    size_of::<SnpDerivedKeyResp>()
);
