netdesk/native/core/src/transport.rs

//! Pluggable transport for the session protocol. The E2E AES-GCM framing sits
//! ABOVE this layer on both the viewer and host, so a transport only ever moves
//! opaque frames (binary media/file chunks, or text control JSON). The current
//! WSS relay ([`WsRelay`]) and the P2P WebRTC DataChannel (added later) both
//! implement [`Transport`], so `connect_with` / `run_host_with` carry either and
//! the relay→P2P upgrade / downgrade can swap transports without touching the
//! protocol, crypto, or framing.

use crate::error::{NdError, Result};
use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt};
use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio::net::TcpStream;
use tokio::sync::mpsc;
use tokio_tungstenite::tungstenite::Message;
use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream};

use bytes::Bytes;
use webrtc::data_channel::data_channel_message::DataChannelMessage;
use webrtc::data_channel::RTCDataChannel;

/// One message crossing a transport. Opaque here; the caller encrypts/decrypts
/// and dispatches by channel byte above this layer.
#[derive(Debug, Clone)]
pub enum TransportMsg {
    Binary(Vec<u8>),
    Text(String),
}

/// Which logical stream a binary frame belongs to. The caller knows this from
/// the (pre-encryption) channel byte; the transport cannot see it because the
/// byte is inside the E2E ciphertext. It decides per-lane delivery: `Media` is
/// drop-tolerant (video — drop a fresh frame rather than queue it late), `Ctl`
/// must be reliable (file bytes, control), `Move` is drop-tolerant pointer
/// motion. Maps 1:1 to the web/Python `media` / `ctl` / `move` DataChannels;
/// the relay and the native single channel use it only for back-pressure.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Lane {
    Media,
    Ctl,
    Move,
}

/// A bidirectional message transport. Used from a single task (the viewer's
/// socket task / the host's `run_host` loop): `recv` is awaited in one
/// `tokio::select!` arm and `send_*` in other arms' bodies, so `&mut self` for
/// all three is sound (select drops the non-chosen futures first).
#[async_trait]
pub trait Transport: Send + 'static {
    async fn send_binary(&mut self, data: Vec<u8>) -> Result<()>;
    async fn send_text(&mut self, text: String) -> Result<()>;
    /// Next inbound message, or `None` when the transport has closed.
    async fn recv(&mut self) -> Option<Result<TransportMsg>>;

    /// Send a binary frame tagged with its [`Lane`]. The default ignores the
    /// lane (relay, native single channel today), so existing transports need
    /// no change; the P2P multiplexer overrides it to pick the channel /
    /// back-pressure policy. Callers that know the channel (video vs file)
    /// should prefer this over [`Transport::send_binary`].
    async fn send_binary_lane(&mut self, data: Vec<u8>, lane: Lane) -> Result<()> {
        let _ = lane;
        self.send_binary(data).await
    }
}

/// The WSS relay transport — the current path, behaviourally identical to the
/// pre-seam code. Wraps the tungstenite WebSocket (TLS to the relay); the server
/// relays frames opaquely between the two session peers.
pub struct WsRelay {
    ws: WebSocketStream<MaybeTlsStream<TcpStream>>,
}

impl WsRelay {
    /// Dial the relay (TLS). Mirrors the old `connect_async` + error mapping.
    pub async fn connect(ws_url: &str) -> Result<Self> {
        let (ws, _resp) = connect_async(ws_url)
            .await
            .map_err(|e| NdError::Connect(e.to_string()))?;
        Ok(Self { ws })
    }

    /// Wrap an already-connected WebSocket (e.g. one opened during a handshake
    /// elsewhere).
    pub fn from_ws(ws: WebSocketStream<MaybeTlsStream<TcpStream>>) -> Self {
        Self { ws }
    }
}

#[async_trait]
impl Transport for WsRelay {
    async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
        self.ws
            .send(Message::Binary(data))
            .await
            .map_err(|e| NdError::Connect(e.to_string()))
    }

    async fn send_text(&mut self, text: String) -> Result<()> {
        self.ws
            .send(Message::Text(text))
            .await
            .map_err(|e| NdError::Connect(e.to_string()))
    }

    async fn recv(&mut self) -> Option<Result<TransportMsg>> {
        loop {
            match self.ws.next().await {
                Some(Ok(Message::Binary(b))) => return Some(Ok(TransportMsg::Binary(b))),
                Some(Ok(Message::Text(t))) => return Some(Ok(TransportMsg::Text(t))),
                Some(Ok(Message::Close(_))) | None => return None,
                // Ping/Pong/Frame are handled by tungstenite; skip at this layer.
                Some(Ok(_)) => continue,
                Some(Err(e)) => return Some(Err(NdError::Connect(e.to_string()))),
            }
        }
    }
}

// ============================ P2P WebRTC DataChannel ============================

/// 9-byte fragment header prepended to every binary chunk:
/// `[frame_id:u32 BE][idx:u16 BE][count:u16 BE][flags:u8]`. `flags` bit0 is
/// reserved for a keyframe hint (set by a higher layer in a later phase); the
/// reassembler ignores it.
const FRAG_HEADER: usize = 9;
/// Each SCTP message stays at or below 16 KiB (header + payload) — safely under
/// the negotiated max-message-size, and small enough that one slow frame never
/// head-of-line-blocks the queue for long.
const MAX_CHUNK_PAYLOAD: usize = 16 * 1024 - FRAG_HEADER;
/// Pause sending once this much is queued in the SCTP send buffer...
const HIGH_WATER: usize = 1024 * 1024;
/// ...and resume when it drains to here (`on_buffered_amount_low` fires).
const LOW_WATER: usize = 512 * 1024;

/// Reassembly accumulator for one multi-chunk binary frame.
struct FrameAsm {
    count: u16,
    got: u16,
    parts: Vec<Option<Vec<u8>>>,
}

impl FrameAsm {
    fn new(count: u16) -> Self {
        Self {
            count,
            got: 0,
            parts: vec![None; count as usize],
        }
    }
    fn put(&mut self, idx: u16, payload: &[u8]) {
        let i = idx as usize;
        if i < self.parts.len() && self.parts[i].is_none() {
            self.parts[i] = Some(payload.to_vec());
            self.got += 1;
        }
    }
    fn complete(&self) -> bool {
        self.got == self.count
    }
    /// Concatenate the chunks in index order. Because the chunks are contiguous
    /// slices of the original nonce||ciphertext||tag blob, the rejoined bytes
    /// are byte-identical — a mis-ordered or missing chunk would corrupt the
    /// blob and the viewer AES-GCM decrypt_frame would fail the auth tag, which
    /// is exactly the integrity guarantee we want.
    fn concat(self) -> Vec<u8> {
        let mut out = Vec::new();
        for p in self.parts.into_iter().flatten() {
            out.extend_from_slice(&p);
        }
        out
    }
}

/// Install the inbound reassembly handler on `dc`: SCTP string messages become
/// `TransportMsg::Text`; fragmented binary chunks are rejoined by `frame_id`
/// and pushed as exactly one contiguous `TransportMsg::Binary` (so the E2E
/// `nonce||ct||tag` blob stays intact for AES-GCM). Shared by `WebRtcChannel`
/// (direct P2P) and the relay/P2P multiplexer so both frame the wire identically.
pub fn install_reassembler(dc: &Arc<RTCDataChannel>, out: mpsc::UnboundedSender<TransportMsg>) {
    let reasm: Arc<Mutex<HashMap<u32, FrameAsm>>> = Arc::new(Mutex::new(HashMap::new()));
    dc.on_message(Box::new(move |msg: DataChannelMessage| {
        let out = out.clone();
        let reasm = reasm.clone();
        Box::pin(async move {
            if msg.is_string {
                if let Ok(s) = String::from_utf8(msg.data.to_vec()) {
                    let _ = out.send(TransportMsg::Text(s));
                }
                return;
            }
            let data = &msg.data;
            if data.len() < FRAG_HEADER {
                return; // malformed; drop
            }
            let frame_id = u32::from_be_bytes([data[0], data[1], data[2], data[3]]);
            let idx = u16::from_be_bytes([data[4], data[5]]);
            let count = u16::from_be_bytes([data[6], data[7]]);
            // data[8] = flags (bit0 = keyframe), reserved; unused for reassembly.
            let payload = &data[FRAG_HEADER..];
            match count {
                0 => {}
                1 => {
                    let _ = out.send(TransportMsg::Binary(payload.to_vec()));
                }
                _ => {
                    let mut map = reasm.lock().unwrap();
                    let entry = map.entry(frame_id).or_insert_with(|| FrameAsm::new(count));
                    entry.put(idx, payload);
                    if entry.complete() {
                        if let Some(asm) = map.remove(&frame_id) {
                            drop(map);
                            let _ = out.send(TransportMsg::Binary(asm.concat()));
                        }
                    }
                }
            }
        })
    }));
}

/// Fragment `data` into `<=16 KiB` SCTP messages (9-byte header each) under one
/// `frame_id` and send them in order on `dc`. No back-pressure here — the caller
/// applies its own policy (wait vs drop-oldest) around this. Returns the error
/// from the first failed `dc.send`.
pub async fn send_fragmented(dc: &Arc<RTCDataChannel>, frame_id: u32, data: &[u8]) -> Result<()> {
    let total = data.len();
    let count = ((total + MAX_CHUNK_PAYLOAD - 1) / MAX_CHUNK_PAYLOAD).max(1);
    if count > u16::MAX as usize {
        return Err(NdError::Connect(format!(
            "frame too large to fragment: {total} bytes"
        )));
    }
    for idx in 0..count {
        let start = idx * MAX_CHUNK_PAYLOAD;
        let end = (start + MAX_CHUNK_PAYLOAD).min(total);
        let chunk = &data[start..end];
        let mut framed = Vec::with_capacity(FRAG_HEADER + chunk.len());
        framed.extend_from_slice(&frame_id.to_be_bytes());
        framed.extend_from_slice(&(idx as u16).to_be_bytes());
        framed.extend_from_slice(&(count as u16).to_be_bytes());
        framed.push(0u8); // flags (bit0 = keyframe) reserved
        framed.extend_from_slice(chunk);
        dc.send(&Bytes::from(framed))
            .await
            .map_err(|e| NdError::Connect(e.to_string()))?;
    }
    Ok(())
}

/// A [`Transport`] over a single reliable+ordered WebRTC DataChannel — the
/// native-to-native direct P2P fast path. The E2E AES-GCM framing is unchanged
/// above this layer; this channel only moves opaque binary frames (fragmented
/// to fit SCTP, reassembled whole so the GCM blob stays contiguous) and opaque
/// text control JSON (sent as SCTP string messages, so the peer tells the two
/// apart with `DataChannelMessage::is_string`).
pub struct WebRtcChannel {
    dc: Arc<RTCDataChannel>,
    inbound: mpsc::UnboundedReceiver<TransportMsg>,
    /// Drains one token whenever the SCTP send buffer falls below `LOW_WATER`.
    low_water_rx: mpsc::Receiver<()>,
    /// Set true by `on_close`; makes `recv` return `None` once inbound is empty.
    closed_flag: Arc<AtomicBool>,
    closed_rx: mpsc::Receiver<()>,
    next_frame_id: u32,
}

impl WebRtcChannel {
    /// Wrap an open DataChannel. Registers the inbound reassembly handler, the
    /// back-pressure low-water signal, and the close signal. Call once the
    /// channel is open (sends before open error at the SCTP layer).
    pub async fn new(dc: Arc<RTCDataChannel>) -> Self {
        let (in_tx, in_rx) = mpsc::unbounded_channel::<TransportMsg>();

        // --- inbound: reassemble binary frames, pass text straight through.
        install_reassembler(&dc, in_tx);

        // --- back-pressure: a size-1 mpsc is race-free vs Notify here. The
        // low-water event only fires when the buffer crosses DOWN through the
        // threshold, so a token can be pending at most once; try_send drops
        // extras.
        let (low_tx, low_rx) = mpsc::channel::<()>(1);
        dc.set_buffered_amount_low_threshold(LOW_WATER).await;
        dc.on_buffered_amount_low(Box::new(move || {
            let low_tx = low_tx.clone();
            Box::pin(async move {
                let _ = low_tx.try_send(());
            })
        }))
        .await;

        // --- close signal so recv() can return None on a dropped channel.
        let closed_flag = Arc::new(AtomicBool::new(false));
        let (closed_tx, closed_rx) = mpsc::channel::<()>(1);
        let cf = closed_flag.clone();
        dc.on_close(Box::new(move || {
            let cf = cf.clone();
            let closed_tx = closed_tx.clone();
            Box::pin(async move {
                cf.store(true, Ordering::SeqCst);
                let _ = closed_tx.try_send(());
            })
        }));

        Self {
            dc,
            inbound: in_rx,
            low_water_rx: low_rx,
            closed_flag,
            closed_rx,
            next_frame_id: 0,
        }
    }

    /// The underlying channel, for stats / state inspection (upgrade logic).
    pub fn data_channel(&self) -> &Arc<RTCDataChannel> {
        &self.dc
    }
}

#[async_trait]
impl Transport for WebRtcChannel {
    async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
        // Pause if the SCTP send queue is already deep: wait for it to drain
        // below LOW_WATER (fresh frames never pile up unbounded in RAM).
        while self.dc.buffered_amount().await + FRAG_HEADER + data.len() > HIGH_WATER {
            if self.low_water_rx.recv().await.is_none() {
                break;
            }
        }
        let frame_id = self.next_frame_id;
        self.next_frame_id = self.next_frame_id.wrapping_add(1);
        send_fragmented(&self.dc, frame_id, &data).await
    }

    async fn send_text(&mut self, text: String) -> Result<()> {
        self.dc
            .send_text(text)
            .await
            .map_err(|e| NdError::Connect(e.to_string()))?;
        Ok(())
    }

    async fn recv(&mut self) -> Option<Result<TransportMsg>> {
        if self.closed_flag.load(Ordering::SeqCst) {
            // Drain anything already reassembled, then report closed.
            return self.inbound.try_recv().ok().map(Ok);
        }
        tokio::select! {
            biased;
            msg = self.inbound.recv() => msg.map(Ok),
            _ = self.closed_rx.recv() => self.inbound.try_recv().ok().map(Ok),
        }
    }
}

/// Direct-IP transport — a raw TCP socket instead of the relay. The viewer dials
/// the host's `IP:port` and the session runs the SAME challenge/auth + E2E
/// AES-GCM above this layer (so direct-IP is encrypted, unlike RustDesk's).
/// Framing is `[type:u8][len:u32 BE][payload]` (type 0 = binary, 1 = text). A
/// dedicated reader task parses frames into a channel, so `recv` is
/// cancellation-safe in the caller's `select!` (a bare `read_exact` is not).
pub struct DirectTcp {
    writer: tokio::net::tcp::OwnedWriteHalf,
    rx: mpsc::UnboundedReceiver<Result<TransportMsg>>,
}

/// Cap on a single framed message (generous for a 4K keyframe).
const DIRECT_MAX_FRAME: usize = 64 * 1024 * 1024;

impl DirectTcp {
    /// Wrap an accepted / dialed TCP stream.
    pub fn new(stream: TcpStream) -> Self {
        let _ = stream.set_nodelay(true);
        let (mut read, write) = stream.into_split();
        let (tx, rx) = mpsc::unbounded_channel();
        tokio::spawn(async move {
            use tokio::io::AsyncReadExt;
            let mut hdr = [0u8; 5];
            loop {
                if read.read_exact(&mut hdr).await.is_err() {
                    break;
                }
                let ty = hdr[0];
                let len = u32::from_be_bytes([hdr[1], hdr[2], hdr[3], hdr[4]]) as usize;
                if len > DIRECT_MAX_FRAME {
                    break;
                }
                let mut buf = vec![0u8; len];
                if read.read_exact(&mut buf).await.is_err() {
                    break;
                }
                let msg = if ty == 1 {
                    TransportMsg::Text(String::from_utf8_lossy(&buf).into_owned())
                } else {
                    TransportMsg::Binary(buf)
                };
                if tx.send(Ok(msg)).is_err() {
                    break;
                }
            }
        });
        DirectTcp { writer: write, rx }
    }

    async fn write_frame(&mut self, ty: u8, payload: &[u8]) -> Result<()> {
        use tokio::io::AsyncWriteExt;
        let mut hdr = [0u8; 5];
        hdr[0] = ty;
        hdr[1..5].copy_from_slice(&(payload.len() as u32).to_be_bytes());
        self.writer
            .write_all(&hdr)
            .await
            .map_err(|e| NdError::Connect(e.to_string()))?;
        self.writer
            .write_all(payload)
            .await
            .map_err(|e| NdError::Connect(e.to_string()))?;
        Ok(())
    }
}

#[async_trait]
impl Transport for DirectTcp {
    async fn send_binary(&mut self, data: Vec<u8>) -> Result<()> {
        self.write_frame(0, &data).await
    }
    async fn send_text(&mut self, text: String) -> Result<()> {
        self.write_frame(1, text.as_bytes()).await
    }
    async fn recv(&mut self) -> Option<Result<TransportMsg>> {
        self.rx.recv().await
    }
}

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

    // DirectTcp framing must round-trip text + binary, in order, over a real
    // loopback socket (the foundation of encrypted direct-IP access).
    #[tokio::test]
    async fn direct_tcp_roundtrips_text_and_binary_in_order() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let srv = tokio::spawn(async move {
            let (s, _) = listener.accept().await.unwrap();
            let mut t = DirectTcp::new(s);
            for _ in 0..3 {
                match t.recv().await.unwrap().unwrap() {
                    TransportMsg::Text(s) => t.send_text(format!("echo:{s}")).await.unwrap(),
                    TransportMsg::Binary(b) => {
                        let mut r = b.clone();
                        r.reverse();
                        t.send_binary(r).await.unwrap();
                    }
                }
            }
        });
        let cs = tokio::net::TcpStream::connect(addr).await.unwrap();
        let mut c = DirectTcp::new(cs);
        c.send_text("hello".into()).await.unwrap();
        c.send_binary(vec![1, 2, 3, 4]).await.unwrap();
        c.send_text("world".into()).await.unwrap();
        assert!(matches!(c.recv().await.unwrap().unwrap(), TransportMsg::Text(s) if s == "echo:hello"));
        assert!(matches!(c.recv().await.unwrap().unwrap(), TransportMsg::Binary(b) if b == vec![4, 3, 2, 1]));
        assert!(matches!(c.recv().await.unwrap().unwrap(), TransportMsg::Text(s) if s == "echo:world"));
        srv.await.unwrap();
    }
}