starry-kernel 0.9.0

use alloc::{borrow::Cow, format, sync::Arc, vec, vec::Vec};
use core::{
    ffi::c_int,
    mem::size_of,
    sync::atomic::{AtomicBool, Ordering},
};

use ax_io::prelude::*;
use ax_net::{InterfaceId, InterfaceInfo, InterfaceKind};
use axpoll::{IoEvents, Pollable};
use axpoll_set::PollSet;
use linux_raw_sys::{
    general::{O_RDWR, S_IFSOCK},
    net::{AF_PACKET, sockaddr},
};

use super::{
    FileLike, Kstat,
    net::{ARPHRD_ETHER, first_visible_ethernet, in_root_net_ns, visible_interface_by_id},
};
use crate::{
    Errno, StarryError, StarryResult,
    file::{IoDst, IoSrc, get_file_like},
    mm::{vm_read_slice, vm_write_slice},
    sync::Mutex,
    task::{
        current_user_task,
        future::{block_on_user, poll_io},
    },
};

const PACKET_HOST: u8 = 0;
const SYNTHETIC_PEER_HWADDR: [u8; 6] = [0x02, 0x00, 0x00, 0x00, 0x00, 0x02];
const ETH_P_IP: u16 = 0x0800;
const ETH_P_ARP: u16 = 0x0806;
const ARPOP_REQUEST: u16 = 1;
const ARPOP_REPLY: u16 = 2;

#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct SockAddrLl {
    pub sll_family: u16,
    pub sll_protocol: u16,
    pub sll_ifindex: i32,
    pub sll_hatype: u16,
    pub sll_pkttype: u8,
    pub sll_halen: u8,
    pub sll_addr: [u8; 8],
}

impl SockAddrLl {
    fn from_interface(info: &InterfaceInfo, protocol: u16) -> StarryResult<Self> {
        if info.kind != InterfaceKind::Ethernet {
            return Err(StarryError::NoSuchDevice);
        }
        let mac = info.mac.ok_or(StarryError::NoSuchDevice)?;
        let mut sll_addr = [0; 8];
        sll_addr[..mac.0.len()].copy_from_slice(&mac.0);
        Ok(Self {
            sll_family: AF_PACKET as u16,
            sll_protocol: protocol,
            sll_ifindex: info.id.to_linux_ifindex(),
            sll_hatype: ARPHRD_ETHER,
            sll_pkttype: PACKET_HOST,
            sll_halen: mac.0.len() as u8,
            sll_addr,
        })
    }

    pub fn read_from_user(
        current: &crate::task::UserTaskRef,
        addr: *const sockaddr,
        addrlen: u32,
    ) -> crate::StarryResult<Self> {
        if addrlen < size_of::<Self>() as u32 {
            return Err(StarryError::InvalidInput);
        }
        let data = read_user_bytes::<{ size_of::<Self>() }>(current, addr as *const u8)?;
        let addr = Self {
            sll_family: u16::from_ne_bytes(data[0..2].try_into().unwrap()),
            sll_protocol: u16::from_ne_bytes(data[2..4].try_into().unwrap()),
            sll_ifindex: i32::from_ne_bytes(data[4..8].try_into().unwrap()),
            sll_hatype: u16::from_ne_bytes(data[8..10].try_into().unwrap()),
            sll_pkttype: data[10],
            sll_halen: data[11],
            sll_addr: data[12..20].try_into().unwrap(),
        };
        if addr.sll_family as u32 != AF_PACKET {
            return Err(StarryError::from(Errno::EAFNOSUPPORT));
        }
        Ok(addr)
    }

    pub fn write_to_user(
        &self,
        current: &crate::task::UserTaskRef,
        addr: *mut sockaddr,
        addrlen: &mut u32,
    ) -> crate::StarryResult<()> {
        let len = (*addrlen as usize).min(size_of::<Self>());
        let data = unsafe { core::slice::from_raw_parts(self as *const Self as *const u8, len) };
        vm_write_slice(current, addr as *mut u8, data)?;
        *addrlen = size_of::<Self>() as u32;
        Ok(())
    }
}

struct PacketSocketState {
    bound: SockAddrLl,
    pending: Option<(Vec<u8>, SockAddrLl)>,
}

pub struct PacketSocket {
    state: Mutex<PacketSocketState>,
    non_blocking: AtomicBool,
    poll_rx: PollSet,
}

impl PacketSocket {
    pub fn new(protocol: u16) -> StarryResult<Self> {
        if !in_root_net_ns() {
            return Err(StarryError::PermissionDenied);
        }
        let info = first_visible_ethernet()?;
        Ok(Self {
            state: Mutex::new(PacketSocketState {
                bound: SockAddrLl::from_interface(&info, protocol)?,
                pending: None,
            }),
            non_blocking: AtomicBool::new(false),
            poll_rx: PollSet::new(),
        })
    }

    pub fn bind_ll(&self, addr: SockAddrLl) -> StarryResult<()> {
        if !in_root_net_ns() {
            return Err(StarryError::NoSuchDevice);
        }
        let info = if addr.sll_ifindex == 0 {
            first_visible_ethernet()?
        } else {
            let id = InterfaceId::from_linux_ifindex(addr.sll_ifindex)
                .ok_or(StarryError::InvalidInput)?;
            visible_interface_by_id(id)?
        };
        // from_interface checks kind, no need to check again
        let mut state = self.state.lock();
        state.bound = SockAddrLl::from_interface(&info, addr.sll_protocol)?;
        Ok(())
    }

    pub fn local_addr(&self) -> SockAddrLl {
        self.state.lock().bound
    }

    pub fn send_packet(&self, src: &mut IoSrc) -> StarryResult<usize> {
        if !in_root_net_ns() {
            return Err(StarryError::NoSuchDevice);
        }
        let len = src.remaining();
        if len == 0 {
            return Ok(0);
        }
        let mut data = vec![0; len];
        let read = src.read(&mut data)?;
        data.truncate(read);

        let bound = self.state.lock().bound;
        if let Some(reply) = build_arp_reply(&data, bound) {
            {
                self.state.lock().pending = Some(reply);
            }
            // Pending packet is stored before waking readers.
            unsafe { self.poll_rx.wake(IoEvents::IN) };
        }
        Ok(read)
    }

    pub fn recv_packet(&self, dst: &mut IoDst) -> crate::StarryResult<(usize, SockAddrLl)> {
        let task = current_user_task();
        block_on_user(
            &task,
            poll_io(self, IoEvents::IN, self.nonblocking(), || {
                let (data, from) = {
                    let mut state = self.state.lock();
                    state.pending.take().ok_or(crate::StarryError::WouldBlock)?
                };
                let written = dst.write(&data)?;
                Ok((written, from))
            }),
        )
        .into_result()?
    }

    pub fn from_fd(fd: c_int) -> StarryResult<Arc<Self>> {
        get_file_like(fd)?
            .downcast_arc()
            .map_err(|_| StarryError::NotASocket)
    }
}

fn build_arp_reply(request: &[u8], bound: SockAddrLl) -> Option<(Vec<u8>, SockAddrLl)> {
    let id = InterfaceId::from_linux_ifindex(bound.sll_ifindex)?;
    let info = visible_interface_by_id(id).ok()?;
    let mac = info.mac?;
    if request.len() < 28
        || u16::from_be_bytes([request[0], request[1]]) != ARPHRD_ETHER
        || u16::from_be_bytes([request[2], request[3]]) != ETH_P_IP
        || request[4] != mac.0.len() as u8
        || request[5] != 4
        || u16::from_be_bytes([request[6], request[7]]) != ARPOP_REQUEST
    {
        return None;
    }

    let request_sender_protocol: [u8; 4] = request[14..18].try_into().ok()?;
    let request_target_protocol: [u8; 4] = request[24..28].try_into().ok()?;
    if !is_modeled_peer_ipv4(&info, request_target_protocol) {
        return None;
    }

    let mut reply = request.to_vec();
    reply[6..8].copy_from_slice(&ARPOP_REPLY.to_be_bytes());
    reply[8..14].copy_from_slice(&SYNTHETIC_PEER_HWADDR);
    reply[14..18].copy_from_slice(&request_target_protocol);
    reply[18..24].copy_from_slice(&request[8..14]);
    reply[24..28].copy_from_slice(&request_sender_protocol);

    let mut from = SockAddrLl::from_interface(&info, ETH_P_ARP.to_be()).ok()?;
    from.sll_addr[..SYNTHETIC_PEER_HWADDR.len()].copy_from_slice(&SYNTHETIC_PEER_HWADDR);

    Some((reply, from))
}

fn is_modeled_peer_ipv4(info: &InterfaceInfo, ip: [u8; 4]) -> bool {
    info.ipv4
        .and_then(|config| config.gateway)
        .is_some_and(|gateway| gateway.octets() == ip)
}

fn read_user_bytes<const N: usize>(
    current: &crate::task::UserTaskRef,
    ptr: *const u8,
) -> crate::StarryResult<[u8; N]> {
    let mut buf = [core::mem::MaybeUninit::<u8>::uninit(); N];
    vm_read_slice(current, ptr, &mut buf)?;
    Ok(buf.map(|b| unsafe { b.assume_init() }))
}

impl FileLike for PacketSocket {
    fn stat(&self) -> StarryResult<Kstat> {
        Ok(Kstat {
            mode: S_IFSOCK | 0o777u32,
            blksize: 4096,
            ..Default::default()
        })
    }

    fn path(&self) -> Cow<'_, str> {
        format!("packet:[{}]", self as *const _ as usize).into()
    }

    fn open_flags(&self) -> u32 {
        O_RDWR
    }

    fn set_nonblocking(&self, nonblocking: bool) -> StarryResult {
        self.non_blocking.store(nonblocking, Ordering::Release);
        Ok(())
    }

    fn nonblocking(&self) -> bool {
        self.non_blocking.load(Ordering::Acquire)
    }

    fn ioctl(
        &self,
        current: &crate::task::UserTaskRef,
        cmd: u32,
        arg: usize,
    ) -> crate::StarryResult<usize> {
        // The SIOCGIF* device ioctls are family-agnostic in Linux sock_ioctl ->
        // dev_ioctl, so AF_PACKET answers them through the same shared helper as
        // the other socket families. Interface visibility (and thus netns
        // scoping) is enforced by device_ioctl's own lookups.
        if let Some(result) = crate::file::net::device_ioctl(current, cmd, arg) {
            return result;
        }
        Err(StarryError::NotATty)
    }
}

impl Pollable for PacketSocket {
    fn poll(&self) -> IoEvents {
        let mut events = IoEvents::OUT;
        events.set(IoEvents::IN, self.state.lock().pending.is_some());
        events
    }

    unsafe fn register_shared(
        &self,
        sink: &mut dyn axpoll::SharedRegistrationSink,
        events: IoEvents,
    ) {
        if events.contains(IoEvents::IN) {
            unsafe { sink.register_shared(&self.poll_rx, IoEvents::IN) };
        }
    }

    unsafe fn register_exclusive(
        &self,
        sink: &mut dyn axpoll::ExclusiveRegistrationSink,
        events: IoEvents,
    ) {
        if events.contains(IoEvents::IN) {
            unsafe { sink.register_exclusive(&self.poll_rx, IoEvents::IN) };
        }
    }
}