// Copyright © 2022 Intel Corporation

//

// SPDX-License-Identifier: Apache-2.0

#![no_main]

use std::os::unix::io::{AsRawFd, FromRawFd};

use std::sync::Arc;

use libfuzzer_sys::{fuzz_target, Corpus};

use seccompiler::SeccompAction;

use virtio_devices::{VirtioDevice, VirtioInterrupt, VirtioInterruptType};

use virtio_queue::{Queue, QueueT};

use vm_memory::bitmap::AtomicBitmap;

use vm_memory::{Bytes, GuestAddress, GuestMemoryAtomic};

use vmm_sys_util::eventfd::{EventFd, EFD_NONBLOCK};

type GuestMemoryMmap = vm_memory::GuestMemoryMmap<AtomicBitmap>;

const QUEUE_DATA_SIZE: usize = 4;

const MEM_SIZE: usize = 512 * 1024;

const BALLOON_SIZE: u64 = 512 * 1024;

// Number of queues

const QUEUE_NUM: usize = 3;

// Max entries in the queue.

const QUEUE_SIZE: u16 = 64;

// Descriptor table alignment

const DESC_TABLE_ALIGN_SIZE: u64 = 16;

// Available ring alignment

const AVAIL_RING_ALIGN_SIZE: u64 = 2;

// Used ring alignment

const USED_RING_ALIGN_SIZE: u64 = 4;

// Descriptor table size

const DESC_TABLE_SIZE: u64 = 16_u64 * QUEUE_SIZE as u64;

// Available ring size

const AVAIL_RING_SIZE: u64 = 6_u64 + 2 * QUEUE_SIZE as u64;

// Used ring size

const USED_RING_SIZE: u64 = 6_u64 + 8 * QUEUE_SIZE as u64;

fuzz_target!(|bytes: &[u8]| -> Corpus {

    if bytes.len() < QUEUE_DATA_SIZE * QUEUE_NUM

        || bytes.len() > (QUEUE_DATA_SIZE * QUEUE_NUM + MEM_SIZE)

    {

        return Corpus::Reject;

    }

    let mut balloon = virtio_devices::Balloon::new(

        "fuzzer_balloon".to_owned(),

        BALLOON_SIZE,

        true,

        true,

        SeccompAction::Allow,

        EventFd::new(EFD_NONBLOCK).unwrap(),

        None,

    )

    .unwrap();

    let queue_data = &bytes[..QUEUE_DATA_SIZE * QUEUE_NUM];

    let mem_bytes = &bytes[QUEUE_DATA_SIZE * QUEUE_NUM..];

    // Setup the guest memory with the input bytes

    let mem = GuestMemoryMmap::from_ranges(&[(GuestAddress(0), MEM_SIZE)]).unwrap();

    if mem.write_slice(mem_bytes, GuestAddress(0 as u64)).is_err() {

        return Corpus::Reject;

    }

    let guest_memory = GuestMemoryAtomic::new(mem);

    // Setup the virt queues with the input bytes

    let mut queues = setup_virt_queues(

        &[

            &queue_data[..QUEUE_DATA_SIZE].try_into().unwrap(),

            &queue_data[QUEUE_DATA_SIZE..QUEUE_DATA_SIZE * 2]

                .try_into()

                .unwrap(),

            &queue_data[QUEUE_DATA_SIZE * 2..QUEUE_DATA_SIZE * 3]

                .try_into()

                .unwrap(),

        ],

        0,

    );

    let inflate_q = queues.remove(0);

    let inflate_evt = EventFd::new(0).unwrap();

    let inflate_queue_evt = unsafe { EventFd::from_raw_fd(libc::dup(inflate_evt.as_raw_fd())) };

    let deflate_q = queues.remove(0);

    let deflate_evt = EventFd::new(0).unwrap();

    let deflate_queue_evt = unsafe { EventFd::from_raw_fd(libc::dup(deflate_evt.as_raw_fd())) };

    let reporting_q = queues.remove(0);

    let reporting_evt = EventFd::new(0).unwrap();

    let reporting_queue_evt = unsafe { EventFd::from_raw_fd(libc::dup(reporting_evt.as_raw_fd())) };

    // Kick the 'queue' events before activate the balloon device

    inflate_queue_evt.write(1).unwrap();

    deflate_queue_evt.write(1).unwrap();

    reporting_queue_evt.write(1).unwrap();

    balloon

        .activate(

            guest_memory,

            Arc::new(NoopVirtioInterrupt {}),

            vec![

                (0, inflate_q, inflate_evt),

                (1, deflate_q, deflate_evt),

                (2, reporting_q, reporting_evt),

            ],

        )

        .ok();

    // Wait for the events to finish and balloon device worker thread to return

    balloon.wait_for_epoll_threads();

    Corpus::Keep

});

pub struct NoopVirtioInterrupt {}

impl VirtioInterrupt for NoopVirtioInterrupt {

    fn trigger(&self, _int_type: VirtioInterruptType) -> std::result::Result<(), std::io::Error> {

        Ok(())

    }

}

macro_rules! align {

    ($n:expr, $align:expr) => {{

        $n.div_ceil($align) * $align

    }};

}

fn setup_virt_queues(bytes: &[&[u8; QUEUE_DATA_SIZE]], base_addr: u64) -> Vec<Queue> {

    let mut queues = Vec::new();

    let mut base_addr = base_addr;

    for b in bytes {

        let mut q = Queue::new(QUEUE_SIZE).unwrap();

        let desc_table_addr = align!(base_addr, DESC_TABLE_ALIGN_SIZE);

        let avail_ring_addr = align!(desc_table_addr + DESC_TABLE_SIZE, AVAIL_RING_ALIGN_SIZE);

        let used_ring_addr = align!(avail_ring_addr + AVAIL_RING_SIZE, USED_RING_ALIGN_SIZE);

        q.try_set_desc_table_address(GuestAddress(desc_table_addr))

            .unwrap();

        q.try_set_avail_ring_address(GuestAddress(avail_ring_addr))

            .unwrap();

        q.try_set_used_ring_address(GuestAddress(used_ring_addr))

            .unwrap();

        q.set_next_avail(b[0] as u16); // 'u8' is enough given the 'QUEUE_SIZE' is small

        q.set_next_used(b[1] as u16);

        q.set_event_idx(b[2] % 2 != 0);

        q.set_size(b[3] as u16 % QUEUE_SIZE);

        q.set_ready(true);

        queues.push(q);

        base_addr = used_ring_addr + USED_RING_SIZE;

    }

    queues

}