nym/wasm/smolmix/src/device.rs

// Copyright 2026 - Nym Technologies SA <contact@nymtech.net>
// SPDX-License-Identifier: Apache-2.0

//! smoltcp `Device` implementation backed by in-memory VecDeque buffers.
//!
//! The native smolmix uses `NymAsyncDevice` (Stream/Sink over mpsc channels) fed
//! into `tokio-smoltcp`. On wasm32 we drive smoltcp directly, so we need a sync
//! `Device` impl instead. The bridge pushes incoming IP packets into `rx_queue`
//! and pops outgoing packets from `tx_queue`.

use std::collections::VecDeque;

use smoltcp::phy::{Device, DeviceCapabilities, Medium, RxToken, TxToken};
use smoltcp::time::Instant;

/// smoltcp device backed by in-memory packet queues.
///
/// The bridge task feeds incoming IP packets via [`push_rx`](Self::push_rx) and
/// drains outgoing packets via [`drain_tx`](Self::drain_tx). The reactor calls
/// smoltcp's `Interface::poll()` which invokes the `Device` trait methods.
pub struct WasmDevice {
    rx_queue: VecDeque<Vec<u8>>,
    tx_queue: VecDeque<Vec<u8>>,
    capabilities: DeviceCapabilities,
}

impl WasmDevice {
    pub fn new() -> Self {
        let mut capabilities = DeviceCapabilities::default();
        capabilities.medium = Medium::Ip;
        // Sized so one IP packet fits in one sphinx packet payload (no
        // chunking-layer fragmentation). Budget in bytes from the 2048 B
        // sphinx plaintext: − 344 (SURB-ack) − 32 (x25519 ephemeral key,
        // Repliable msgs) − 7 (frag header) − 1 (padding) − 53 (LP+IPR
        // framing + AEAD) ≈ 1611. 1600 leaves ~11 B headroom for IPR
        // overhead variability.
        capabilities.max_transmission_unit = 1600;
        // Native smolmix also uses Some(1) in the device, but tokio-smoltcp
        // compensates with a burst loop that calls Interface::poll() up to 100
        // times per reactor iteration (each processing 1 packet). Our WASM
        // reactor only calls poll() once per tick, so we set a higher burst
        // size to let smoltcp drain all pending packets in a single poll().
        capabilities.max_burst_size = Some(100);

        Self {
            rx_queue: VecDeque::new(),
            tx_queue: VecDeque::new(),
            capabilities,
        }
    }

    /// Push an incoming IP packet (from the mixnet) into the receive queue.
    pub fn push_rx(&mut self, packet: Vec<u8>) {
        self.rx_queue.push_back(packet);
    }

    /// Drain all outgoing IP packets (generated by smoltcp) from the transmit queue.
    pub fn drain_tx(&mut self) -> impl Iterator<Item = Vec<u8>> + '_ {
        self.tx_queue.drain(..)
    }
}

impl Device for WasmDevice {
    type RxToken<'a> = WasmRxToken;
    type TxToken<'a> = WasmTxToken<'a>;

    fn receive(&mut self, _timestamp: Instant) -> Option<(Self::RxToken<'_>, Self::TxToken<'_>)> {
        let packet = self.rx_queue.pop_front()?;
        Some((
            WasmRxToken { buffer: packet },
            WasmTxToken {
                queue: &mut self.tx_queue,
            },
        ))
    }

    fn transmit(&mut self, _timestamp: Instant) -> Option<Self::TxToken<'_>> {
        Some(WasmTxToken {
            queue: &mut self.tx_queue,
        })
    }

    fn capabilities(&self) -> DeviceCapabilities {
        self.capabilities.clone()
    }
}

/// Receive token: delivers one packet from the rx queue to smoltcp.
pub struct WasmRxToken {
    buffer: Vec<u8>,
}

impl RxToken for WasmRxToken {
    fn consume<R, F>(self, f: F) -> R
    where
        F: FnOnce(&[u8]) -> R,
    {
        f(&self.buffer)
    }
}

/// Transmit token: captures one packet from smoltcp into the tx queue.
pub struct WasmTxToken<'a> {
    queue: &'a mut VecDeque<Vec<u8>>,
}

impl<'a> TxToken for WasmTxToken<'a> {
    fn consume<R, F>(self, len: usize, f: F) -> R
    where
        F: FnOnce(&mut [u8]) -> R,
    {
        let mut buffer = vec![0u8; len];
        let result = f(&mut buffer);
        self.queue.push_back(buffer);
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn push_rx_and_receive() {
        let mut dev = WasmDevice::new();
        dev.push_rx(vec![1, 2, 3]);

        let now = Instant::from_millis(0);
        let (rx, _tx) = dev.receive(now).expect("should have a packet");
        let data = rx.consume(|buf| buf.to_vec());
        assert_eq!(data, vec![1, 2, 3]);
    }

    #[test]
    fn transmit_and_drain() {
        let mut dev = WasmDevice::new();
        let now = Instant::from_millis(0);

        let tx = dev.transmit(now).expect("should get tx token");
        tx.consume(4, |buf| {
            buf.copy_from_slice(&[0xDE, 0xAD, 0xBE, 0xEF]);
        });

        let packets: Vec<_> = dev.drain_tx().collect();
        assert_eq!(packets.len(), 1);
        assert_eq!(packets[0], vec![0xDE, 0xAD, 0xBE, 0xEF]);
    }

    #[test]
    fn empty_receive_returns_none() {
        let mut dev = WasmDevice::new();
        assert!(dev.receive(Instant::from_millis(0)).is_none());
    }

    #[test]
    fn capabilities_are_ip_mode() {
        let dev = WasmDevice::new();
        let caps = dev.capabilities();
        assert_eq!(caps.medium, Medium::Ip);
        assert_eq!(caps.max_transmission_unit, 1980);
    }
}