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

//! Backing for SNP partitions.

use super::Hcl;
use super::HclVp;
use super::MshvVtl;
use super::NoRunner;
use super::ProcessorRunner;
use super::hcl_pvalidate_pages;
use super::hcl_rmpadjust_pages;
use super::hcl_rmpquery_pages;
use super::mshv_pvalidate;
use super::mshv_rmpadjust;
use super::mshv_rmpquery;
use crate::GuestVtl;
use crate::vmsa::VmsaWrapper;
use hv1_structs::VtlArray;
use hvdef::HV_PAGE_SIZE;
use hvdef::HvRegisterName;
use hvdef::HvRegisterValue;
use memory_range::MemoryRange;
use sidecar_client::SidecarVp;
use std::cell::UnsafeCell;
use std::os::fd::AsRawFd;
use thiserror::Error;
use x86defs::snp::SevAvicPage;
use x86defs::snp::SevRmpAdjust;
use x86defs::snp::SevVmsa;

/// Runner backing for SNP partitions.
pub struct Snp<'a> {
    vmsa: VtlArray<&'a UnsafeCell<SevVmsa>, 2>,
    avic_pages: VtlArray<&'a UnsafeCell<SevAvicPage>, 2>,
}

/// Error returned by failing SNP operations.
#[derive(Debug, Error)]
#[expect(missing_docs)]
pub enum SnpError {
    #[error("operating system error")]
    Os(#[source] nix::Error),
    #[error("isa error {0:?}")]
    Isa(u32),
}

/// Error returned by failing SNP page operations.
#[derive(Debug, Error)]
#[expect(missing_docs)]
pub enum SnpPageError {
    #[error("pvalidate failed")]
    Pvalidate(#[source] SnpError),
    #[error("rmpadjust failed")]
    Rmpadjust(#[source] SnpError),
    #[error("rmpquery failed")]
    Rmpquery(#[source] SnpError),
}

impl MshvVtl {
    /// Execute the pvalidate instruction on the specified memory range.
    ///
    /// The range must not be mapped in the kernel as RAM.
    pub fn pvalidate_pages(
        &self,
        range: MemoryRange,
        validate: bool,
        terminate_on_failure: bool,
    ) -> Result<(), SnpPageError> {
        tracing::debug!(%range, validate, terminate_on_failure, "pvalidate");
        // SAFETY: TODO SNP FUTURE: we are passing parameters as the kernel requires.
        // For defense in depth it could be useful to prevent usermode from changing
        // visibility of a VTL2 kernel page in the kernel.
        let ret = unsafe {
            hcl_pvalidate_pages(
                self.file.as_raw_fd(),
                &mshv_pvalidate {
                    start_pfn: range.start() / HV_PAGE_SIZE,
                    page_count: (range.end() - range.start()) / HV_PAGE_SIZE,
                    validate: validate as u8,
                    terminate_on_failure: terminate_on_failure as u8,
                    ram: 0,
                    padding: [0; 1],
                },
            )
            .map_err(SnpError::Os)
            .map_err(SnpPageError::Pvalidate)?
        };

        if ret != 0 {
            return Err(SnpPageError::Pvalidate(SnpError::Isa(ret as u32)));
        }

        Ok(())
    }

    /// Execute the rmpadjust instruction on the specified memory range.
    ///
    /// The range must not be mapped in the kernel as RAM.
    pub fn rmpadjust_pages(
        &self,
        range: MemoryRange,
        value: SevRmpAdjust,
        terminate_on_failure: bool,
    ) -> Result<(), SnpPageError> {
        // SAFETY: TODO SNP FUTURE: For defense in depth it could be useful to prevent
        // usermode from changing permissions of a VTL2 kernel page in the kernel.
        let ret = unsafe {
            hcl_rmpadjust_pages(
                self.file.as_raw_fd(),
                &mshv_rmpadjust {
                    start_pfn: range.start() / HV_PAGE_SIZE,
                    page_count: (range.end() - range.start()) / HV_PAGE_SIZE,
                    value: value.into(),
                    terminate_on_failure: terminate_on_failure as u8,
                    ram: 0,
                    padding: Default::default(),
                },
            )
            .map_err(SnpError::Os)
            .map_err(SnpPageError::Rmpadjust)?
        };

        if ret != 0 {
            return Err(SnpPageError::Rmpadjust(SnpError::Isa(ret as u32)));
        }

        Ok(())
    }

    /// Gets the current vtl permissions for a page.
    /// Note: only supported on Genoa+
    pub fn rmpquery_page(&self, gpa: u64, vtl: GuestVtl) -> Result<SevRmpAdjust, SnpPageError> {
        let page_count = 1u64;
        let mut flags = [u64::from(SevRmpAdjust::new().with_target_vmpl(match vtl {
            GuestVtl::Vtl0 => 2,
            GuestVtl::Vtl1 => 1,
        })); 1];

        let mut page_size = [0; 1];
        let mut pages_processed = 0u64;

        debug_assert!(flags.len() == page_count as usize);
        debug_assert!(page_size.len() == page_count as usize);

        let query = mshv_rmpquery {
            start_pfn: gpa / HV_PAGE_SIZE,
            page_count,
            terminate_on_failure: 0,
            ram: 0,
            padding: Default::default(),
            flags: flags.as_mut_ptr(),
            page_size: page_size.as_mut_ptr(),
            pages_processed: &mut pages_processed,
        };

        // SAFETY: the input query is the correct type for this ioctl
        unsafe {
            hcl_rmpquery_pages(self.file.as_raw_fd(), &query)
                .map_err(SnpError::Os)
                .map_err(SnpPageError::Rmpquery)?;
        }

        assert!(pages_processed <= page_count);

        Ok(SevRmpAdjust::from(flags[0]))
    }
}

impl<'a> super::private::BackingPrivate<'a> for Snp<'a> {
    fn new(vp: &'a HclVp, sidecar: Option<&SidecarVp<'_>>, _hcl: &Hcl) -> Result<Self, NoRunner> {
        assert!(sidecar.is_none());
        let super::BackingState::Snp {
            vtl0_apic_page,
            vtl1_apic_page,
            vmsa,
        } = &vp.backing
        else {
            return Err(NoRunner::MismatchedIsolation);
        };

        // SAFETY: The mapping is held for the appropriate lifetime, and the
        // APIC page is never accessed as any other type, or by any other location.
        let vtl1_apic_page = unsafe { &*vtl1_apic_page.base().cast() };

        Ok(Self {
            avic_pages: [vtl0_apic_page.as_ref(), vtl1_apic_page].into(),
            vmsa: vmsa.each_ref().map(|mp| mp.as_ref()),
        })
    }

    fn try_set_reg(
        _runner: &mut ProcessorRunner<'a, Self>,
        _vtl: GuestVtl,
        _name: HvRegisterName,
        _value: HvRegisterValue,
    ) -> bool {
        false
    }

    fn must_flush_regs_on(_runner: &ProcessorRunner<'a, Self>, _name: HvRegisterName) -> bool {
        false
    }

    fn try_get_reg(
        _runner: &ProcessorRunner<'a, Self>,
        _vtl: GuestVtl,
        _name: HvRegisterName,
    ) -> Option<HvRegisterValue> {
        None
    }

    fn flush_register_page(_runner: &mut ProcessorRunner<'a, Self>) {}
}

impl<'a> ProcessorRunner<'a, Snp<'a>> {
    /// Gets a reference to the VMSA and backing state of a VTL
    pub fn vmsa(&self, vtl: GuestVtl) -> VmsaWrapper<'_, &SevVmsa> {
        // SAFETY: the VMSA will not be concurrently accessed by the processor
        // while this VP is in VTL2.
        let vmsa = unsafe { &*self.state.vmsa[vtl].get() };

        VmsaWrapper::new(vmsa, &self.hcl.snp_register_bitmap)
    }

    /// Gets a mutable reference to the VMSA and backing state of a VTL.
    pub fn vmsa_mut(&mut self, vtl: GuestVtl) -> VmsaWrapper<'_, &mut SevVmsa> {
        // SAFETY: the VMSA will not be concurrently accessed by the processor
        // while this VP is in VTL2.
        let vmsa = unsafe { &mut *self.state.vmsa[vtl].get() };

        VmsaWrapper::new(vmsa, &self.hcl.snp_register_bitmap)
    }

    /// Returns the VMSAs for [VTL0, VTL1].
    pub fn vmsas_mut(&mut self) -> [VmsaWrapper<'_, &mut SevVmsa>; 2] {
        self.state
            .vmsa
            .each_mut()
            .map(|vmsa| {
                // SAFETY: the VMSA will not be concurrently accessed by the processor
                // while this VP is in VTL2.
                let vmsa = unsafe { &mut *vmsa.get() };

                VmsaWrapper::new(vmsa, &self.hcl.snp_register_bitmap)
            })
            .into_inner()
    }

    /// Gets a PFN of the VTL0 secure AVIC page.
    /// TODO: Maybe there is a better way other than passing `cpu_index` here.
    pub fn secure_avic_vtl0_pfn(&self, cpu_index: u32) -> u64 {
        self.hcl.secure_avic_vtl0_pfn(cpu_index)
    }

    /// Gets a reference to the secure AVIC page for the given VTL.
    pub fn secure_avic_page(&self, vtl: GuestVtl) -> &SevAvicPage {
        // SAFETY: the APIC pages will not be concurrently accessed by the processor
        // while this VP is in VTL2.
        unsafe { &*self.state.avic_pages[vtl].get() }
    }

    /// Gets a mutable reference to the secure AVIC page for the given VTL.
    pub fn secure_avic_page_mut(&mut self, vtl: GuestVtl) -> &mut SevAvicPage {
        // SAFETY: the AVIC pages will not be concurrently accessed by the processor
        // while this VP is in VTL2.
        unsafe { &mut *self.state.avic_pages[vtl].get() }
    }

    /// Gets a mutable reference to the secure AVIC page for the given VTL.
    pub fn secure_avic_page_vmsa_mut(
        &mut self,
        vtl: GuestVtl,
    ) -> (&mut SevAvicPage, VmsaWrapper<'_, &mut SevVmsa>) {
        // SAFETY: the AVIC pages will not be concurrently accessed by the processor
        // while this VP is in VTL2.
        let avic_page = unsafe { &mut *self.state.avic_pages[vtl].get() };
        let vmsa = self.vmsa_mut(vtl);

        (avic_page, vmsa)
    }

    /// Gets a mutable reference to the secure AVIC page and the proxy_irr_exit field
    pub fn secure_avic_page_proxy_irr_exit_vtl0_mut(
        &mut self,
    ) -> (&mut SevAvicPage, &mut [u32; 8]) {
        // SAFETY: the AVIC pages will not be concurrently accessed by the processor
        // while this VP is in VTL2.
        let avic_page = unsafe { &mut *self.state.avic_pages[GuestVtl::Vtl0].get() };
        // SAFETY: The `proxy_irr_exit` field of the run page will not be concurrently updated.
        let proxy_irr_vtl0 = unsafe { &mut (*self.run.get()).proxy_irr_exit };

        (avic_page, proxy_irr_vtl0)
    }
}
