lancer / kernel / src / ring / process.rs

use core::sync::atomic::Ordering;

use crate::cap::cnode;
use crate::cap::object::ObjectTag;
use crate::cap::ops;
use crate::cap::pool::POOL;
use crate::cap::table::Rights;
use crate::error::KernelError;
use crate::ipc::message;
use crate::ipc::{endpoint, notification};
use crate::mem::addr;
use crate::proc::context::IpcMessage;
use crate::proc::{BlockedReason, PROCESSES};
use crate::ring::{
    CompletionEntry, MAX_CQ_ENTRIES, MAX_SQ_ENTRIES, RingHeader, RingIndex, RingOpcode,
    SubmissionEntry,
};
use crate::types::Pid;

const MAX_RING_BATCH: u32 = 16;

fn _assert_from_bytes<T: zerocopy::FromBytes>() {}
#[allow(dead_code)]
fn _assert_ring_types_are_zerocopy() {
    _assert_from_bytes::<SubmissionEntry>();
    _assert_from_bytes::<CompletionEntry>();
}

struct UserSnapshot {
    sq_tail: RingIndex,
    cq_head: RingIndex,
}

unsafe fn snapshot_user_fields(ring_base: *const u8) -> UserSnapshot {
    let header = unsafe { &*(ring_base as *const RingHeader) };
    UserSnapshot {
        sq_tail: RingIndex::new(header.sq_tail.load(Ordering::Acquire)),
        cq_head: RingIndex::new(header.cq_head.load(Ordering::Acquire)),
    }
}

unsafe fn commit_header(ring_base: *mut u8, new_sq_head: RingIndex, new_cq_tail: RingIndex) {
    let header = unsafe { &*(ring_base as *const RingHeader) };
    header.sq_head.store(new_sq_head.raw(), Ordering::Release);
    header.cq_tail.store(new_cq_tail.raw(), Ordering::Release);
}

unsafe fn read_sq_entry(ring_base: *const u8, index: u32) -> SubmissionEntry {
    let sq_base = unsafe { ring_base.add(super::ring_sq_offset()) };
    let entry_ptr =
        unsafe { sq_base.add((index as usize) * core::mem::size_of::<SubmissionEntry>()) };
    unsafe { core::ptr::read_volatile(entry_ptr as *const SubmissionEntry) }
}

unsafe fn write_cq_entry(ring_base: *mut u8, index: u32, entry: CompletionEntry) {
    let cq_base = unsafe { ring_base.add(super::ring_cq_offset()) };
    let entry_ptr =
        unsafe { cq_base.add((index as usize) * core::mem::size_of::<CompletionEntry>()) };
    unsafe { core::ptr::write_volatile(entry_ptr as *mut CompletionEntry, entry) };
}

struct CqeResult {
    value: i64,
    extra: u64,
}

impl CqeResult {
    const fn ok() -> Self {
        Self { value: 0, extra: 0 }
    }

    const fn success(value: i64) -> Self {
        Self { value, extra: 0 }
    }

    const fn with_extra(value: i64, extra: u64) -> Self {
        Self { value, extra }
    }
}

fn ring_cap_create(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;
    let tag = ObjectTag::try_from(sqe.arg0)?;

    let ptable = PROCESSES.lock();
    let (cnode_id, cnode_gen, depth, gv, gb) = cnode::cnode_coords(pid, &ptable)?;
    let mut pool = POOL.lock_after(&ptable);
    ops::create_via_cnode(&mut pool, cnode_id, cnode_gen, address, depth, gv, gb, tag)
        .map(|phys| CqeResult::success(phys.raw() as i64))
}

fn ring_cap_derive(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let src_addr = sqe.cap_slot as u64;
    let dest_addr = sqe.arg0;
    let rights_mask = Rights::from_bits(sqe.arg1 as u16);

    let ptable = PROCESSES.lock();
    let (cnode_id, cnode_gen, depth, gv, gb) = cnode::cnode_coords(pid, &ptable)?;
    let mut pool = POOL.lock_after(&ptable);
    ops::derive_via_cnode(
        &mut pool,
        cnode_id,
        cnode_gen,
        src_addr,
        dest_addr,
        depth,
        gv,
        gb,
        rights_mask,
    )
    .map(|()| CqeResult::ok())
}

fn ring_ipc_send(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;
    let msg = IpcMessage::from_regs([sqe.arg0, sqe.arg1, sqe.arg2, 0, 0, 0]);

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid,
            address,
            ObjectTag::Endpoint,
            Rights::WRITE,
            &ptable,
            &pool,
        )?
    };

    ptable.exec_mut(pid).unwrap().ipc_message = msg;

    let recv_pid = {
        let mut pool = POOL.lock_after(&ptable);
        let ep = pool.write_as::<lancer_core::object_layout::EndpointObject>(
            cap.phys(),
            cap.generation(),
        )?;
        let mut receivers = endpoint::load_receivers(ep);
        let dequeued = endpoint::dequeue_genuine_receiver(&mut receivers, &mut ptable);
        if let Some(ref br) = dequeued {
            endpoint::store_holder(ep, Some(br.pid()));
        }
        endpoint::store_receivers(ep, &receivers);
        dequeued
    };

    match recv_pid {
        Some(blocked_recv) => {
            let recv_pid_val = blocked_recv.pid();
            let sender_prio = ptable[pid].effective_priority();

            ptable.unblock_and_enqueue(recv_pid_val, blocked_recv)?;
            let recv_exec = ptable.exec_mut(recv_pid_val).unwrap();
            recv_exec.ipc_message = msg;
            recv_exec.ipc_badge = pid.raw() as u64;
            message::inject_into_context(&mut recv_exec.saved_context, &msg);
            recv_exec.saved_context.rax = pid.raw() as u64;
            recv_exec.seal_context();

            crate::sched::boost_effective(&mut ptable, recv_pid_val, sender_prio);

            Ok(CqeResult::ok())
        }
        None => Err(KernelError::WouldBlock),
    }
}

fn ring_ipc_recv(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid,
            address,
            ObjectTag::Endpoint,
            Rights::READ,
            &ptable,
            &pool,
        )?
    };

    let sender_pid = {
        let mut pool = POOL.lock_after(&ptable);
        let ep = pool.write_as::<lancer_core::object_layout::EndpointObject>(
            cap.phys(),
            cap.generation(),
        )?;
        let mut senders = endpoint::load_senders(ep);
        let dequeued = endpoint::dequeue(&mut senders, &mut ptable);
        if dequeued.is_some() {
            endpoint::store_holder(ep, Some(pid));
        }
        endpoint::store_senders(ep, &senders);
        dequeued
    };

    match sender_pid {
        Some(blocked_sender) => {
            let sender_pid = blocked_sender.pid();
            let sender_msg = ptable.exec(sender_pid).unwrap().ipc_message;
            let recv_exec = ptable.exec_mut(pid).unwrap();
            recv_exec.ipc_message = sender_msg;
            recv_exec.ipc_badge = sender_pid.raw() as u64;

            match ptable[sender_pid].blocked_reason() {
                Some(BlockedReason::Calling(_, _)) => {
                    let mut pool = POOL.lock_after(&ptable);
                    match pool.write_as::<lancer_core::object_layout::EndpointObject>(
                        cap.phys(),
                        cap.generation(),
                    ) {
                        Ok(ep) => {
                            let mut receivers = endpoint::load_receivers(ep);
                            match endpoint::enqueue(&mut receivers, blocked_sender, &mut ptable) {
                                Ok(()) => {
                                    endpoint::store_receivers(ep, &receivers);
                                    ptable.exec_mut(pid).unwrap().reply_target = Some(sender_pid);
                                }
                                Err(_) => {
                                    endpoint::store_receivers(ep, &receivers);
                                    let sender_exec = ptable.exec_mut(sender_pid).unwrap();
                                    sender_exec.saved_context.rax =
                                        crate::error::KernelError::ResourceExhausted.to_errno()
                                            as u64;
                                    sender_exec.seal_context();
                                    {
                                        let proof = ptable[sender_pid].blocked_proof();
                                        let r = ptable.unblock_and_enqueue(sender_pid, proof);
                                        debug_assert!(r.is_ok());
                                    }
                                }
                            }
                        }
                        Err(e) => {
                            let sender_exec = ptable.exec_mut(sender_pid).unwrap();
                            sender_exec.saved_context.rax = e.to_errno() as u64;
                            sender_exec.seal_context();
                            {
                                let r = ptable.unblock_and_enqueue(sender_pid, blocked_sender);
                                debug_assert!(r.is_ok());
                            }
                        }
                    }
                }
                _ => {
                    ptable.unblock_and_enqueue(sender_pid, blocked_sender)?;
                }
            }

            let sender_prio = ptable[sender_pid].effective_priority();
            crate::sched::reset_effective(&mut ptable, pid);
            crate::sched::boost_effective(&mut ptable, pid, sender_prio);

            Ok(CqeResult::with_extra(
                sender_pid.raw() as i64,
                sender_msg.regs[0],
            ))
        }
        None => Err(KernelError::WouldBlock),
    }
}

fn ring_notify_signal(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;
    let bits = sqe.arg0;

    let mut ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid,
            address,
            ObjectTag::Notification,
            Rights::WRITE,
            &ptable,
            &pool,
        )?
    };

    notification::do_signal(&cap, bits, &mut ptable).map(|_| CqeResult::ok())
}

fn ring_notify_poll(sqe: &SubmissionEntry, pid: Pid) -> Result<CqeResult, KernelError> {
    let address = sqe.cap_slot as u64;

    let ptable = PROCESSES.lock();
    let cap = {
        let pool = POOL.lock_after(&ptable);
        cnode::resolve_caller_validate(
            pid,
            address,
            ObjectTag::Notification,
            Rights::READ,
            &ptable,
            &pool,
        )?
    };

    drop(ptable);
    notification::do_poll(&cap).map(|val| CqeResult::with_extra(0, val))
}

fn process_submission(sqe: &SubmissionEntry, pid: Pid) -> CompletionEntry {
    let result = match RingOpcode::from_u8(sqe.opcode) {
        Some(RingOpcode::Nop) => Ok(CqeResult::ok()),
        Some(RingOpcode::CapCreate) => ring_cap_create(sqe, pid),
        Some(RingOpcode::CapDerive) => ring_cap_derive(sqe, pid),
        Some(RingOpcode::IpcSend) => ring_ipc_send(sqe, pid),
        Some(RingOpcode::IpcRecv) => ring_ipc_recv(sqe, pid),
        Some(RingOpcode::NotifySignal) => ring_notify_signal(sqe, pid),
        Some(RingOpcode::NotifyPoll) => ring_notify_poll(sqe, pid),
        None => Err(KernelError::InvalidParameter),
    };

    match result {
        Ok(cqe) => CompletionEntry {
            result: cqe.value,
            user_data: sqe.user_data as u64,
            extra: cqe.extra,
        },
        Err(e) => CompletionEntry {
            result: e.to_errno(),
            user_data: sqe.user_data as u64,
            extra: 0,
        },
    }
}

pub fn ring_enter(
    ring_phys_base: x86_64::PhysAddr,
    pid: Pid,
    _min_complete: u32,
) -> Result<i64, KernelError> {
    if !ring_phys_base.as_u64().is_multiple_of(4096) {
        return Err(KernelError::InvalidAddress);
    }

    let max_phys = (crate::mem::phys::BitmapFrameAllocator::total_frames() as u64) * 4096;
    let ring_end = ring_phys_base
        .as_u64()
        .checked_add(super::ring_total_size() as u64)
        .ok_or(KernelError::InvalidAddress)?;
    if ring_end > max_phys {
        return Err(KernelError::InvalidAddress);
    }

    let hhdm = addr::hhdm_offset();
    let ring_virt = (ring_phys_base.as_u64() + hhdm) as *mut u8;

    let mut ptable = PROCESSES.lock();
    let exec = ptable.exec_mut(pid).ok_or(KernelError::InvalidObject)?;

    let sq_head = exec.ring_sq_head;
    let cq_tail = exec.ring_cq_tail;
    let user_snap = unsafe { snapshot_user_fields(ring_virt) };

    let pending = user_snap.sq_tail.distance(sq_head);
    if pending > MAX_SQ_ENTRIES {
        return Err(KernelError::InvalidParameter);
    }

    let cq_used = cq_tail.distance(user_snap.cq_head);
    let cq_available = MAX_CQ_ENTRIES.saturating_sub(cq_used);
    let to_process = pending.min(cq_available).min(MAX_RING_BATCH);

    if to_process == 0 {
        return Ok(0);
    }
    drop(ptable);

    let (completed, cq_write) = (0..to_process).fold((0u32, cq_tail), |(done, cq_w), i| {
        let sq_slot = sq_head.advance(i).slot(MAX_SQ_ENTRIES);
        let sqe = unsafe { read_sq_entry(ring_virt, sq_slot) };
        let cqe = process_submission(&sqe, pid);

        let cq_slot = cq_w.slot(MAX_CQ_ENTRIES);
        unsafe { write_cq_entry(ring_virt, cq_slot, cqe) };
        (done + 1, cq_w.advance(1))
    });

    let new_sq_head = sq_head.advance(to_process);
    unsafe { commit_header(ring_virt, new_sq_head, cq_write) };

    let mut ptable = PROCESSES.lock();
    if let Some(exec) = ptable.exec_mut(pid) {
        exec.ring_sq_head = new_sq_head;
        exec.ring_cq_tail = cq_write;
    }

    Ok(completed as i64)
}