Example stack

ergot_stack.rs

use {
    bbqio::v0_6 as bbqio6,
    bbqueue::{
        BBQueue,
        traits::{notifier::maitake::MaiNotSpsc, storage::Inline},
    },
    ergot::{
        exports::mutex::raw_impls::local::LocalRawMutex,
        interface_manager::{
            Interface, interface_impls::embedded_io::IoInterface,
            profiles::direct_edge::CENTRAL_NODE_ID,
        },
        toolkits::embedded_io_async_v0_6::{self as kit, tx_worker},
    },
};

use edge_interface_plus::EdgeInterfacePlus;

use cobs_acc::{CobsAccumulator, FeedResult};
use embedded_io_async_06::Read;

use ergot::{
    Header, NetStack,
    interface_manager::{
        InterfaceSendError, InterfaceState, Profile, SetStateError,
        profiles::direct_edge::EDGE_NODE_ID,
    },
    net_stack::NetStackHandle,
    wire_frames::de_frame,
};
use static_cell::ConstStaticCell;

pub const OUT_QUEUE_SIZE: usize = 4096;
pub const MAX_PACKET_SIZE: usize = 1024;

pub const RX_PIPE_SIZE: usize = 1200;
pub const TX_PIPE_SIZE: usize = 1201;

// The type of our netstack
pub type Stack<const N: usize> = NetStack<LocalRawMutex, MyRouter<IoInterface<&'static Queue>, N>>;
// The type of our outgoing queue
pub type Queue = kit::Queue<OUT_QUEUE_SIZE, crate::Coord>;

type MyPipe<const SIZE: usize> = BBQueue<Inline<SIZE>, crate::Coord, MaiNotSpsc>;

pub type RxPipe = MyPipe<{ RX_PIPE_SIZE }>;
pub type TxPipe = MyPipe<{ TX_PIPE_SIZE }>;

pub fn new_stack<const N: usize>(outqs: [&'static Queue; N]) -> Stack<N> {
    use ergot::interface_manager::utils::cobs_stream;
    NetStack::new_with_profile(MyRouter::new(
        outqs.map(|q| cobs_stream::Sink::new_from_handle(q, MAX_PACKET_SIZE as u16)),
    ))
}

pub struct PinnedErgotData<const N: usize> {
    pub outq: [Queue; N],
    pub recv_buf: [ConstStaticCell<[u8; MAX_PACKET_SIZE]>; N],
    pub scratch_buf: [ConstStaticCell<[u8; 64]>; N],
}

impl<const N: usize> PinnedErgotData<N> {
    pub const fn new() -> Self {
        Self {
            outq: [const { Queue::new() }; N],
            recv_buf: [const { ConstStaticCell::new([0u8; MAX_PACKET_SIZE]) }; N],
            scratch_buf: [const { ConstStaticCell::new([0u8; 64]) }; N],
        }
    }
}

pub async fn run_link<N, R>(
    mut rcvr: MyRxWorker<N, R>,
    recv_buf: &'static mut [u8],
    scratch_buf: &'static mut [u8],
    outq: &'static Queue,
    tx: bbqio6::PWrap<&'static TxPipe>,
)
where
    N: NetStackHandle,
    R: Read,
    <<N as NetStackHandle>::Profile as Profile>::InterfaceIdent: From<IfaceIdent>,
{
    embassy_futures::join::join(
        run_rx(&mut rcvr, recv_buf, scratch_buf),
        run_tx(outq, tx),
    )
    .await;
}

/// Worker task for incoming data
async fn run_rx<N, R>(rcvr: &mut MyRxWorker<N, R>, recv_buf: &'static mut [u8], scratch_buf: &'static mut [u8])
where
    N: NetStackHandle,
    R: Read,
    <<N as NetStackHandle>::Profile as Profile>::InterfaceIdent: From<IfaceIdent>,
{
    loop {
        _ = rcvr.run(recv_buf, scratch_buf).await;
    }
}

/// Worker task for outgoing data
async fn run_tx(outq: &'static Queue, mut tx: bbqio6::PWrap<&'static TxPipe>) {
    loop {
        _ = tx_worker(&mut tx, outq.stream_consumer()).await;
    }
}

pub struct MyRxWorker<N, R>
where
    N: NetStackHandle,
    R: Read,
    <<N as NetStackHandle>::Profile as Profile>::InterfaceIdent: From<IfaceIdent>,
{
    nsh: N,
    rx: R,
    net_id: u16,
    ident: IfaceIdent,
}

impl<N, R> MyRxWorker<N, R>
where
    N: NetStackHandle,
    R: Read,
    <<N as NetStackHandle>::Profile as Profile>::InterfaceIdent: From<IfaceIdent>,
{
    pub fn new(nsh: N, rx: R, net_id: u16, ident: IfaceIdent) -> Self {
        Self {
            nsh,
            rx,
            net_id,
            ident,
        }
    }

    pub async fn run(
        &mut self,
        recv_buf: &mut [u8],
        scratch_buf: &mut [u8],
    ) -> Result<(), R::Error> {
        _ = self.nsh.stack().manage_profile(|im| {
            im.set_interface_state(
                self.ident.into(),
                InterfaceState::Active {
                    net_id: self.net_id,
                    node_id: CENTRAL_NODE_ID,
                },
            )
        });
        let res = self.run_inner(recv_buf, scratch_buf).await;
        _ = self
            .nsh
            .stack()
            .manage_profile(|im| im.set_interface_state(self.ident.into(), InterfaceState::Down));
        res
    }

    async fn run_inner(
        &mut self,
        recv_buf: &mut [u8],
        scratch_buf: &mut [u8],
    ) -> Result<(), R::Error> {
        let mut acc = CobsAccumulator::new(recv_buf);
        let Self {
            nsh,
            rx,
            net_id,
            ident,
        } = self;
        'outer: loop {
            let used = rx.read(scratch_buf).await?;
            let mut window = &mut scratch_buf[..used];
            loop {
                match acc.feed_raw(window) {
                    FeedResult::Consumed => continue 'outer,
                    FeedResult::OverFull(remaining) => {
                        window = remaining;
                    }
                    FeedResult::DecodeError(remaining) => {
                        window = remaining;
                    }
                    FeedResult::Success { data, remaining }
                    | FeedResult::SuccessInput { data, remaining } => {
                        process_frame(*net_id, data, nsh, (*ident).into());
                        window = remaining;
                    }
                }
            }
        }
    }
}

impl<N, R> Drop for MyRxWorker<N, R>
where
    N: NetStackHandle,
    R: Read,
    <<N as NetStackHandle>::Profile as Profile>::InterfaceIdent: From<IfaceIdent>,
{
    fn drop(&mut self) {
        self.nsh.stack().manage_profile(|im| {
            _ = im.set_interface_state(self.ident.into(), InterfaceState::Down);
        })
    }
}

const BASE_NET_ID: u16 = 1;

/// Interface identifier: index into the fixed-size interface array
pub type IfaceIdent = u8;

struct Node<I: Interface> {
    edge: EdgeInterfacePlus<I>,
    net_id: u16,
    ident: IfaceIdent,
}

/// A no-heap router backed by a fixed-size array of `N` directly-connected interfaces.
/// Net IDs are assigned contiguously starting at `BASE_NET_ID` (slot 0 → net 1, slot 1 → net 2, …).
pub struct MyRouter<I: Interface, const N: usize> {
    ifaces: [Node<I>; N],
}

fn net_id_to_idx(net_id: u16) -> Option<usize> {
    net_id.checked_sub(BASE_NET_ID).map(|i| i as usize)
}

impl<I: Interface, const N: usize> MyRouter<I, N> {
    pub fn new(sinks: [<I as Interface>::Sink; N]) -> Self {
        let mut n = 0u8;
        Self {
            ifaces: sinks.map(|sink| {
                let net_id = BASE_NET_ID + n as u16;
                let node = Node {
                    edge: EdgeInterfacePlus::new_controller(sink, InterfaceState::Down),
                    net_id,
                    ident: n,
                };
                n += 1;
                node
            }),
        }
    }

    fn find(
        &mut self,
        hdr: &Header,
        source: Option<IfaceIdent>,
    ) -> Result<&mut EdgeInterfacePlus<I>, InterfaceSendError> {
        if hdr.dst.port_id == 0 && hdr.any_all.is_none() {
            return Err(InterfaceSendError::AnyPortMissingKey);
        }

        let idx = net_id_to_idx(hdr.dst.network_id)
            .filter(|&i| i < N)
            .ok_or(InterfaceSendError::NoRouteToDest)?;

        // Is this actually for us (central node on this network)?
        if hdr.dst.node_id == CENTRAL_NODE_ID {
            return Err(InterfaceSendError::DestinationLocal);
        }

        // Routing loop: destination is the same interface the packet came in on
        if let Some(src) = source
            && idx as u8 == src
        {
            return Err(InterfaceSendError::RoutingLoop);
        }

        Ok(&mut self.ifaces[idx].edge)
    }
}

impl<I: Interface, const N: usize> Profile for MyRouter<I, N> {
    type InterfaceIdent = IfaceIdent;

    fn send<T: serde::Serialize>(
        &mut self,
        hdr: &Header,
        data: &T,
    ) -> Result<(), InterfaceSendError> {
        let mut hdr = hdr.clone();
        if hdr.decrement_ttl().is_err() {
            return Err(InterfaceSendError::NoRouteToDest);
        }

        if hdr.dst.port_id == 255 {
            if hdr.any_all.is_none() {
                return Err(InterfaceSendError::AnyPortMissingKey);
            }
            let mut any_good = false;
            for node in self.ifaces.iter_mut() {
                if hdr.dst.network_id == node.net_id {
                    continue;
                }
                let mut hdr = hdr.clone();
                hdr.dst.network_id = node.net_id;
                any_good |= node.edge.send(&hdr, data).is_ok();
            }
            if any_good {
                Ok(())
            } else {
                Err(InterfaceSendError::NoRouteToDest)
            }
        } else {
            let intfc = self.find(&hdr, None)?;
            intfc.send(&hdr, data)
        }
    }

    fn send_err(
        &mut self,
        hdr: &Header,
        err: ergot::ProtocolError,
        source: Option<Self::InterfaceIdent>,
    ) -> Result<(), InterfaceSendError> {
        let mut hdr = hdr.clone();
        if hdr.decrement_ttl().is_err() {
            return Err(InterfaceSendError::NoRouteToDest);
        }
        let intfc = self.find(&hdr, source)?;
        intfc.send_err(&hdr, err)
    }

    fn send_raw(
        &mut self,
        hdr: &ergot::HeaderSeq,
        data: &[u8],
        source: Self::InterfaceIdent,
    ) -> Result<(), InterfaceSendError> {
        let mut hdr = hdr.clone();
        if hdr.decrement_ttl().is_err() {
            return Err(InterfaceSendError::NoRouteToDest);
        }

        if hdr.dst.port_id == 255 {
            if hdr.any_all.is_none() {
                return Err(InterfaceSendError::AnyPortMissingKey);
            }
            if N == 0 {
                return Err(InterfaceSendError::NoRouteToDest);
            }
            let mut any_good = false;
            for node in self.ifaces.iter_mut() {
                if source == node.ident {
                    continue;
                }
                let mut hdr = hdr.clone();
                hdr.dst.network_id = node.net_id;
                any_good |= node.edge.send_raw(&hdr, data).is_ok();
            }
            if any_good {
                Ok(())
            } else {
                Err(InterfaceSendError::NoRouteToDest)
            }
        } else {
            let nshdr: Header = hdr.clone().into();
            let intfc = self.find(&nshdr, Some(source))?;
            intfc.send_raw(&hdr, data)
        }
    }

    fn interface_state(&mut self, ident: Self::InterfaceIdent) -> Option<InterfaceState> {
        self.ifaces
            .get_mut(ident as usize)?
            .edge
            .interface_state(())
    }

    fn set_interface_state(
        &mut self,
        ident: Self::InterfaceIdent,
        state: InterfaceState,
    ) -> Result<(), SetStateError> {
        let Some(node) = self.ifaces.get_mut(ident as usize) else {
            return Err(SetStateError::InterfaceNotFound);
        };
        node.edge.set_interface_state((), state)
    }
}

pub fn process_frame<N>(
    net_id: u16,
    data: &[u8],
    nsh: &N,
    ident: <<N as NetStackHandle>::Profile as Profile>::InterfaceIdent,
) where
    N: NetStackHandle,
{
    // Successfully received a packet, now we need to
    // do something with it.
    if let Some(mut frame) = de_frame(data) {
        // trace!("{} got frame from {:?}", frame.hdr, ident);
        // If the message comes in and has a src net_id of zero,
        // we should rewrite it so it isn't later understood as a
        // local packet.
        if frame.hdr.src.network_id == 0 {
            match frame.hdr.src.node_id {
                0 => {
                    /*log::warn!(
                        "{}: device is sending us frames without a node id, ignoring",
                        frame.hdr
                    );*/
                    return;
                }
                CENTRAL_NODE_ID => {
                    // log::warn!("{}: device is sending us frames as us, ignoring", frame.hdr);
                    return;
                }
                EDGE_NODE_ID => {}
                _ => {
                    /*log::warn!(
                        "{}: device is sending us frames with a bad node id, ignoring",
                        frame.hdr
                    );*/
                    return;
                }
            }

            frame.hdr.src.network_id = net_id;
        }
        // TODO: if the destination IS self.net_id, we could rewrite the
        // dest net_id as zero to avoid a pass through the interface manager.
        //
        // If the dest is 0, should we rewrite the dest as self.net_id? This
        // is the opposite as above, but I dunno how that will work with responses
        let hdr = frame.hdr.clone();
        let nshdr: Header = hdr.clone().into();

        let res = match frame.body {
            Ok(body) => nsh.stack().send_raw(&hdr, body, ident),
            Err(e) => nsh.stack().send_err(&nshdr, e, Some(ident)),
        };
        match res {
            Ok(()) => {}
            Err(_) => {
                // TODO: match on error, potentially try to send NAK?
                // warn!("{} recv->send error: {:?}", frame.hdr, e);
            }
        }
    } else {
        // warn!("Decode error! Ignoring frame on net_id {}", net_id);
    }
}

mod edge_interface_plus {
    use serde::Serialize;

    use ergot::{
        Header, HeaderSeq, ProtocolError,
        interface_manager::{
            Interface, InterfaceSendError, InterfaceSink, InterfaceState, SetStateError,
            profiles::direct_edge::{CENTRAL_NODE_ID, EDGE_NODE_ID},
        },
    };

    pub struct EdgeInterfacePlus<I: Interface> {
        sink: I::Sink,
        seq_no: u16,
        state: InterfaceState,
        own_node_id: u8,
        other_node_id: u8,
    }

    impl<I: Interface> EdgeInterfacePlus<I> {
        pub const fn new_controller(sink: I::Sink, state: InterfaceState) -> Self {
            Self {
                sink,
                seq_no: 0,
                state,
                own_node_id: CENTRAL_NODE_ID,
                other_node_id: EDGE_NODE_ID,
            }
        }
    }

    impl<I: Interface> EdgeInterfacePlus<I> {
        fn common_send<'b>(
            &'b mut self,
            hdr: &Header,
        ) -> Result<(&'b mut I::Sink, HeaderSeq), InterfaceSendError> {
            let net_id = match &self.state {
                InterfaceState::Down | InterfaceState::Inactive => {
                    return Err(InterfaceSendError::NoRouteToDest);
                }
                InterfaceState::ActiveLocal { .. } => {
                    // TODO: maybe also handle this?
                    return Err(InterfaceSendError::NoRouteToDest);
                }
                InterfaceState::Active { net_id, node_id: _ } => *net_id,
            };

            // trace!("{} common_send", hdr);

            // TODO: when this WAS a real Profile, we did a lot of these things, but
            // now they should be done by the router. For now, we just have asserts,
            // eventually we should relax this to debug_asserts?
            assert!(net_id != 0);
            let for_us = hdr.dst.network_id == net_id && hdr.dst.node_id == self.own_node_id;
            assert!(!for_us);

            let mut hdr = hdr.clone();

            // If the source is local, rewrite the source using this interface's
            // information so responses can find their way back here
            if hdr.src.net_node_any() {
                // todo: if we know the destination is EXACTLY this network,
                // we could leave the network_id local to allow for shorter
                // addresses
                hdr.src.network_id = net_id;
                hdr.src.node_id = self.own_node_id;
            }

            // If this is a broadcast message, update the destination, ignoring
            // whatever was there before
            if hdr.dst.port_id == 255 {
                hdr.dst.network_id = net_id;
                hdr.dst.node_id = self.other_node_id;
            }

            // If this message has no seq_no, assign it one
            let header = hdr.to_headerseq_or_with_seq(|| {
                let seq_no = self.seq_no;
                self.seq_no = self.seq_no.wrapping_add(1);
                seq_no
            });
            if [0, 255].contains(&hdr.dst.port_id) && hdr.any_all.is_none() {
                return Err(InterfaceSendError::AnyPortMissingKey);
            }

            Ok((&mut self.sink, header))
        }
    }

    /// NOTE: this LOOKS like a profile impl, because it was, but it's actually not, because
    /// this version of DirectEdge only serves DirectRouter
    impl<I: Interface> EdgeInterfacePlus<I> {
        pub(super) fn send<T: Serialize>(
            &mut self,
            hdr: &Header,
            data: &T,
        ) -> Result<(), InterfaceSendError> {
            let (intfc, header) = self.common_send(hdr)?;

            let res = intfc.send_ty(&header, data);

            match res {
                Ok(()) => Ok(()),
                Err(()) => Err(InterfaceSendError::InterfaceFull),
            }
        }

        pub(super) fn send_err(
            &mut self,
            hdr: &Header,
            err: ProtocolError,
        ) -> Result<(), InterfaceSendError> {
            let (intfc, header) = self.common_send(hdr)?;

            let res = intfc.send_err(&header, err);

            match res {
                Ok(()) => Ok(()),
                Err(()) => Err(InterfaceSendError::InterfaceFull),
            }
        }

        pub(super) fn send_raw(
            &mut self,
            hdr: &HeaderSeq,
            data: &[u8],
        ) -> Result<(), InterfaceSendError> {
            let nshdr: Header = hdr.clone().into();
            let (intfc, header) = self.common_send(&nshdr)?;
            let res = intfc.send_raw(&header, data);

            match res {
                Ok(()) => Ok(()),
                Err(()) => Err(InterfaceSendError::InterfaceFull),
            }
        }

        pub(super) fn interface_state(&mut self, _ident: ()) -> Option<InterfaceState> {
            Some(self.state)
        }

        pub(super) fn set_interface_state(
            &mut self,
            _ident: (),
            state: InterfaceState,
        ) -> Result<(), SetStateError> {
            match state {
                InterfaceState::Down => {
                    self.state = InterfaceState::Down;
                }
                InterfaceState::Inactive => {
                    self.state = InterfaceState::Inactive;
                }
                InterfaceState::ActiveLocal { node_id } => {
                    if node_id != self.own_node_id {
                        return Err(SetStateError::InvalidNodeId);
                    }
                    self.state = InterfaceState::ActiveLocal { node_id };
                }
                InterfaceState::Active { net_id, node_id } => {
                    if node_id != self.own_node_id {
                        return Err(SetStateError::InvalidNodeId);
                    }
                    self.state = InterfaceState::Active { net_id, node_id };
                }
            }
            Ok(())
        }
    }
}