yggdrasil/driver/net/core/src/lib.rs

#![feature(map_try_insert)]
#![allow(clippy::type_complexity)]
#![no_std]

extern crate alloc;

use core::mem::size_of;

use alloc::sync::Arc;
use bytemuck::Pod;
use ethernet::L2Packet;
use libk::{mem::PageBox, runtime, util::queue::UnboundedMpmcQueue};
use l3::L3Packet;
use yggdrasil_abi::{error::Error, net::protocols::EthernetFrame};

pub mod ethernet;
pub mod l3;
pub mod l4;

pub mod socket;

pub mod config;
pub mod interface;
pub mod util;

pub use interface::register_interface;

pub struct Packet {
    // TODO info about "received" interface
    buffer: PageBox<[u8]>,
    offset: usize,
    iface: u32,
}

pub struct PacketBuilder {
    data: PageBox<[u8]>,
    pos: usize,
    len: usize,
}

impl PacketBuilder {
    pub fn new(l2_offset: usize, l2_size: usize) -> Result<Self, Error> {
        let data = PageBox::new_slice(0, l2_offset + l2_size)?;
        Ok(Self {
            data,
            pos: l2_offset,
            len: l2_offset,
        })
    }

    #[inline]
    pub fn push<T: Pod>(&mut self, value: &T) -> Result<(), Error> {
        self.push_bytes(bytemuck::bytes_of(value))
    }

    pub fn push_bytes(&mut self, bytes: &[u8]) -> Result<(), Error> {
        if self.pos + bytes.len() > self.data.len() {
            return Err(Error::OutOfMemory);
        }
        self.data[self.pos..self.pos + bytes.len()].copy_from_slice(bytes);
        self.pos += bytes.len();
        Ok(())
    }

    pub fn finish(self) -> (PageBox<[u8]>, usize) {
        (self.data, self.len)
    }
}

impl Packet {
    #[inline]
    pub fn new(buffer: PageBox<[u8]>, offset: usize, iface: u32) -> Self {
        Self {
            buffer,
            offset,
            iface,
        }
    }
}

static PACKET_QUEUE: UnboundedMpmcQueue<Packet> = UnboundedMpmcQueue::new();
static ACCEPT_QUEUE: UnboundedMpmcQueue<L3Packet> = UnboundedMpmcQueue::new();

#[inline]
pub fn receive_packet(packet: Packet) -> Result<(), Error> {
    PACKET_QUEUE.push_back(packet);
    Ok(())
}

pub fn start_network_tasks() -> Result<(), Error> {
    runtime::spawn(l2_packet_handler_worker())?;
    for _ in 0..4 {
        runtime::spawn(l3_accept_worker())?;
    }
    runtime::spawn(config::network_config_service())?;

    Ok(())
}

async fn l2_packet_handler_worker() {
    loop {
        let packet = PACKET_QUEUE.pop_front().await;

        let eth_frame: &EthernetFrame = bytemuck::from_bytes(
            &packet.buffer[packet.offset..packet.offset + size_of::<EthernetFrame>()],
        );

        let l2_packet = L2Packet {
            interface_id: packet.iface,

            source_address: eth_frame.source_mac,
            destination_address: eth_frame.destination_mac,

            l2_offset: packet.offset,
            l3_offset: packet.offset + size_of::<EthernetFrame>(),

            data: Arc::new(packet.buffer),
        };

        ethernet::handle(l2_packet);
    }
}

async fn l3_accept_worker() {
    loop {
        let l3_packet = ACCEPT_QUEUE.pop_front().await;

        // log::debug!(
        //     "INPUT {} {}:{:?} -> {}:{:?}: ACCEPT",
        //     l3_packet.protocol,
        //     l3_packet.source_address,
        //     l3_packet.source_port,
        //     l3_packet.destination_address,
        //     l3_packet.destination_port
        // );

        if let Err(error) = l3::handle_accepted(l3_packet).await {
            log::error!("L3 handle error: {:?}", error);
        }
    }
}