// Copyright 2022, The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Project Rialto main source file.

#![no_main]
#![no_std]

mod communication;
mod error;
mod exceptions;
mod fdt;

extern crate alloc;

use crate::communication::VsockStream;
use crate::error::{Error, Result};
use crate::fdt::{read_dice_range_from, read_is_strict_boot, read_vendor_hashtree_root_digest};
use alloc::boxed::Box;
use ciborium_io::Write;
use core::num::NonZeroUsize;
use core::slice;
use diced_open_dice::{bcc_handover_parse, DiceArtifacts};
use log::{debug, error, info};
use service_vm_comm::{ServiceVmRequest, VmType};
use service_vm_fake_chain::service_vm;
use service_vm_requests::{process_request, RequestContext};
use virtio_drivers::{
    device::socket::{VsockAddr, VMADDR_CID_HOST},
    transport::{pci::bus::PciRoot, DeviceType, Transport},
    Hal,
};
use vmbase::{
    configure_heap,
    fdt::pci::PciInfo,
    fdt::SwiotlbInfo,
    generate_image_header,
    layout::crosvm,
    main,
    memory::{
        init_shared_pool, map_rodata, map_rodata_outside_main_memory, resize_available_memory,
        SIZE_128KB,
    },
    power::reboot,
    virtio::{
        pci::{self, PciTransportIterator, VirtIOSocket},
        HalImpl,
    },
};

fn host_addr(fdt: &libfdt::Fdt) -> Result<VsockAddr> {
    Ok(VsockAddr { cid: VMADDR_CID_HOST, port: vm_type(fdt)?.port() })
}

fn vm_type(fdt: &libfdt::Fdt) -> Result<VmType> {
    if read_is_strict_boot(fdt)? {
        Ok(VmType::ProtectedVm)
    } else {
        Ok(VmType::NonProtectedVm)
    }
}

/// # Safety
///
/// Behavior is undefined if any of the following conditions are violated:
/// * The `fdt_addr` must be a valid pointer and points to a valid `Fdt`.
unsafe fn try_main(fdt_addr: usize) -> Result<()> {
    info!("Welcome to Rialto!");

    let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
    map_rodata(fdt_addr, fdt_size)?;
    // SAFETY: The tracker validated the range to be in main memory, mapped, and not overlap.
    let fdt = unsafe { slice::from_raw_parts(fdt_addr as *mut u8, fdt_size.into()) };
    // We do not need to validate the DT since it is already validated in pvmfw.
    let fdt = libfdt::Fdt::from_slice(fdt)?;

    let memory_range = fdt.first_memory_range()?;
    resize_available_memory(&memory_range).inspect_err(|_| {
        error!("Failed to use memory range value from DT: {memory_range:#x?}");
    })?;

    let swiotlb_range = SwiotlbInfo::new_from_fdt(fdt)
        .inspect_err(|_| {
            error!("Rialto failed when access swiotlb");
        })?
        .and_then(|info| info.fixed_range());
    init_shared_pool(swiotlb_range).inspect_err(|_| {
        error!("Failed to initialize shared pool.");
    })?;

    let bcc_handover: Box<dyn DiceArtifacts> = match vm_type(fdt)? {
        VmType::ProtectedVm => {
            let dice_range = read_dice_range_from(fdt)?;
            info!("DICE range: {dice_range:#x?}");
            let dice_size = dice_range.len().try_into().unwrap();
            // SAFETY: The DICE memory region has been generated by pvmfw and doesn't overlap.
            unsafe { map_rodata_outside_main_memory(dice_range.start, dice_size) }.inspect_err(
                |_| {
                    error!("Failed to use DICE range from DT: {dice_range:#x?}");
                },
            )?;
            let dice_start = dice_range.start as *const u8;
            // SAFETY: There's no memory overlap and the region is mapped as read-only data.
            let bcc_handover = unsafe { slice::from_raw_parts(dice_start, dice_range.len()) };
            Box::new(bcc_handover_parse(bcc_handover)?)
        }
        // Currently, a sample DICE data is used for non-protected VMs, as these VMs only run
        // in tests at the moment.
        VmType::NonProtectedVm => Box::new(service_vm::fake_service_vm_dice_artifacts()?),
    };

    let pci_info = PciInfo::from_fdt(fdt)?;
    debug!("PCI: {pci_info:#x?}");
    let mut pci_root = pci::initialize(pci_info).map_err(Error::PciInitializationFailed)?;
    debug!("PCI root: {pci_root:#x?}");
    let socket_device = find_socket_device::<HalImpl>(&mut pci_root)?;
    debug!("Found socket device: guest cid = {:?}", socket_device.guest_cid());
    let vendor_hashtree_root_digest = read_vendor_hashtree_root_digest(fdt)?;
    let request_context =
        RequestContext { dice_artifacts: bcc_handover.as_ref(), vendor_hashtree_root_digest };

    let mut vsock_stream = VsockStream::new(socket_device, host_addr(fdt)?)?;
    while let ServiceVmRequest::Process(req) = vsock_stream.read_request()? {
        info!("Received request: {}", req.name());
        let response = process_request(req, &request_context);
        info!("Sending response: {}", response.name());
        vsock_stream.write_response(&response)?;
        vsock_stream.flush()?;
    }
    vsock_stream.shutdown()?;

    Ok(())
}

fn find_socket_device<T: Hal>(pci_root: &mut PciRoot) -> Result<VirtIOSocket<T>> {
    PciTransportIterator::<T>::new(pci_root)
        .find(|t| DeviceType::Socket == t.device_type())
        .map(VirtIOSocket::<T>::new)
        .transpose()
        .map_err(Error::VirtIOSocketCreationFailed)?
        .ok_or(Error::MissingVirtIOSocketDevice)
}

/// Entry point for Rialto.
pub fn main(fdt_addr: u64, _a1: u64, _a2: u64, _a3: u64) {
    log::set_max_level(log::LevelFilter::Debug);
    // SAFETY: `fdt_addr` is supposed to be a valid pointer and points to
    // a valid `Fdt`.
    if let Err(e) = unsafe { try_main(fdt_addr as usize) } {
        error!("Rialto failed with {e}");
        reboot()
    }
}

generate_image_header!();
main!(main);
configure_heap!(SIZE_128KB * 2);