yggdrasil/kernel/driver/net/core/src/interface.rs

use core::{
    mem::{size_of, MaybeUninit},
    net::IpAddr,
    sync::atomic::{AtomicU32, AtomicUsize, Ordering},
};

use alloc::{boxed::Box, collections::BTreeMap, format, sync::Arc};
use libk::dma::DmaBuffer;
// TODO: link state management?
use libk_util::{
    sync::spin_rwlock::{IrqSafeRwLock, IrqSafeRwLockReadGuard},
    OneTimeInit,
};
use yggdrasil_abi::{
    error::Error,
    net::{link::LinkState, protocols::EthernetFrame, MacAddress},
};

use crate::{
    l3::{arp::ArpTable, Route},
    TxPacketBuilder,
};

pub trait NetworkDevice: Sync + Send {
    fn allocate_transmit_buffer(&self, len: usize) -> Result<DmaBuffer<[MaybeUninit<u8>]>, Error>;
    fn transmit_buffer(&self, buffer: DmaBuffer<[u8]>) -> Result<(), Error>;

    fn packet_prefix_size(&self) -> usize;

    fn read_hardware_address(&self) -> MacAddress;
    fn link_state(&self) -> LinkState {
        LinkState::Other { up: true }
    }
}

pub struct NetworkInterface {
    pub(crate) name: Box<str>,
    pub(crate) device: Arc<dyn NetworkDevice>,
    pub(crate) mac: MacAddress,

    pub(crate) address: IrqSafeRwLock<Option<IpAddr>>,
    pub(crate) id: u32,
}

#[derive(PartialEq, Eq)]
pub enum NetworkInterfaceType {
    Ethernet,
    Loopback,
}

static INTERFACES: IrqSafeRwLock<BTreeMap<u32, Arc<NetworkInterface>>> =
    IrqSafeRwLock::new(BTreeMap::new());
static LAST_INTERFACE_ID: AtomicU32 = AtomicU32::new(1);
static LOOPBACK: OneTimeInit<Arc<NetworkInterface>> = OneTimeInit::new();

impl NetworkInterface {
    pub fn id(&self) -> u32 {
        self.id
    }

    pub fn loopback() -> &'static Arc<Self> {
        LOOPBACK.get()
    }

    pub fn get(id: u32) -> Result<Arc<Self>, Error> {
        INTERFACES
            .read()
            .get(&id)
            .cloned()
            .ok_or(Error::DoesNotExist)
    }

    pub fn query_by_name(name: &str) -> Result<Arc<Self>, Error> {
        INTERFACES
            .read()
            .iter()
            .find_map(|(_, iface)| {
                if iface.name.as_ref() == name {
                    Some(iface.clone())
                } else {
                    None
                }
            })
            .ok_or(Error::DoesNotExist)
    }

    pub fn list_ref() -> IrqSafeRwLockReadGuard<'static, BTreeMap<u32, Arc<NetworkInterface>>> {
        INTERFACES.read()
    }

    pub fn set_address(&self, address: IpAddr) {
        // Flush routes associated with the interface
        {
            let mut routes = Route::list_mut();
            routes.retain(|route| route.interface != self.id);
        }

        let mut addr = self.address.write();
        // Flush owned ARP entries related to the old address
        if let Some(address) = *addr {
            ArpTable::flush_address(self.id, address);
        }

        log::info!("Assigned address: {:?} -> {:?}", self.name, address);
        addr.replace(address);

        ArpTable::insert_address(self.id, self.mac, address, true);
    }

    pub fn link_state(&self) -> LinkState {
        self.device.link_state()
    }

    pub fn send_l2(&self, l2_frame: &EthernetFrame, l2_data: &[u8]) -> Result<(), Error> {
        let mut builder = TxPacketBuilder::new(self, size_of::<EthernetFrame>() + l2_data.len())?;
        builder.push_bytes(bytemuck::bytes_of(l2_frame))?;
        builder.push_bytes(l2_data)?;
        builder.transmit(None)
    }
}

pub fn register_interface(
    ty: NetworkInterfaceType,
    dev: Arc<dyn NetworkDevice>,
) -> Arc<NetworkInterface> {
    let name = match ty {
        NetworkInterfaceType::Ethernet => {
            static LAST_ETHERNET_ID: AtomicUsize = AtomicUsize::new(0);
            let eth_id = LAST_ETHERNET_ID.fetch_add(1, Ordering::SeqCst);
            format!("eth{eth_id}").into_boxed_str()
        }
        NetworkInterfaceType::Loopback => "lo".into(),
    };

    let mac = dev.read_hardware_address();
    let id = LAST_INTERFACE_ID.fetch_add(1, Ordering::SeqCst);

    log::info!("Register network interface {} (#{}): {}", name, id, mac);

    let iface = NetworkInterface {
        name,
        device: dev,
        mac,
        address: IrqSafeRwLock::new(None),
        id,
    };

    let interface = Arc::new(iface);

    INTERFACES.write().insert(id, interface.clone());

    if ty == NetworkInterfaceType::Loopback {
        LOOPBACK.init(interface.clone());
    }

    interface
}