dosh/src/transport.rs

//! Embeddable reliable stream machinery for Dosh-native application transports.
//!
//! The terminal client/server binaries use the same wire packets exposed here:
//! `StreamOpen`, `StreamData`, `StreamWindowAdjust`, and `StreamClose`. This
//! module keeps the stream state machine independent from UDP sockets, native
//! auth, terminal rendering, or any particular service protocol. An embedding
//! application is expected to:
//!
//! 1. authenticate and maintain a Dosh session using [`crate::protocol`],
//! 2. encrypt each [`OutgoingStreamPacket::body`] as the given packet kind,
//! 3. pass received stream packet bodies back into [`StreamMux`],
//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data.
//!
//! That split is intentional: Dosh roaming/reconnect stays at the authenticated
//! session layer, while this module provides the reusable fast persistent byte
//! streams that higher-level app protocols can embed.

use crate::protocol::{
    self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData,
    StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust,
};
use anyhow::{Result, anyhow, bail};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::net::UdpSocket;

pub const DEFAULT_INITIAL_WINDOW: usize = 1024 * 1024;
pub const DEFAULT_RETRANSMIT_AFTER: Duration = Duration::from_millis(200);
pub const DEFAULT_KEEPALIVE_AFTER: Duration = Duration::from_secs(2);
pub const SERVICE_TARGET_PREFIX: &str = "@dosh-";

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TransportConfig {
    pub initial_window: usize,
    pub retransmit_after: Duration,
    pub keepalive_after: Duration,
}

impl Default for TransportConfig {
    fn default() -> Self {
        Self {
            initial_window: DEFAULT_INITIAL_WINDOW,
            retransmit_after: DEFAULT_RETRANSMIT_AFTER,
            keepalive_after: DEFAULT_KEEPALIVE_AFTER,
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SessionRole {
    Client,
    Server,
}

impl SessionRole {
    fn send_direction(self) -> u32 {
        match self {
            Self::Client => CLIENT_TO_SERVER,
            Self::Server => SERVER_TO_CLIENT,
        }
    }

    fn recv_direction(self) -> u32 {
        match self {
            Self::Client => SERVER_TO_CLIENT,
            Self::Server => CLIENT_TO_SERVER,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct SessionTransportConfig {
    pub role: SessionRole,
    pub conn_id: [u8; 16],
    pub session_key: [u8; 32],
    pub peer_addr: SocketAddr,
    pub initial_send_seq: u64,
    pub initial_ack: u64,
    pub stream: TransportConfig,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum SessionEvent {
    Stream(TransportEvent),
    Ping,
    Pong,
    Ignored,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct OutgoingStreamPacket {
    pub kind: PacketKind,
    pub body: Vec<u8>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IncomingStreamData {
    pub stream_id: u64,
    pub chunks: Vec<Vec<u8>>,
    pub received_offset: u64,
    pub consumed: usize,
    pub window_adjust: OutgoingStreamPacket,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IncomingStreamOpen {
    pub stream_id: u64,
    pub target_host: String,
    pub target_port: u16,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct AcceptedService {
    pub name: String,
    pub stream_id: u64,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TransportEvent {
    Open(IncomingStreamOpen),
    OpenOk {
        stream_id: u64,
        flushed: Vec<OutgoingStreamPacket>,
    },
    OpenReject {
        stream_id: u64,
        reason: String,
    },
    Data(IncomingStreamData),
    WindowAdjust {
        stream_id: u64,
        flushed: Vec<OutgoingStreamPacket>,
    },
    Close {
        stream_id: u64,
    },
}

#[derive(Debug, Clone)]
struct PendingStreamChunk {
    offset: u64,
    bytes: Vec<u8>,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Debug, Clone)]
struct PendingStreamOpen {
    target_host: String,
    target_port: u16,
    last_sent: Instant,
    attempts: u32,
}

#[derive(Debug, Clone)]
pub struct StreamMux {
    config: TransportConfig,
    opened_streams: HashSet<u64>,
    pending_opens: HashMap<u64, PendingStreamOpen>,
    send_credit: HashMap<u64, usize>,
    pending_data: HashMap<u64, VecDeque<Vec<u8>>>,
    sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
    next_send_offset: HashMap<u64, u64>,
    next_recv_offset: HashMap<u64, u64>,
    recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>>,
}

#[derive(Debug, Clone, Default)]
pub struct ServiceRegistry {
    services: HashSet<String>,
}

impl ServiceRegistry {
    pub fn new() -> Self {
        Self::default()
    }

    pub fn register(&mut self, name: impl Into<String>) -> Result<()> {
        let name = name.into();
        service_target(&name)?;
        self.services.insert(name);
        Ok(())
    }

    pub fn is_registered(&self, name: &str) -> bool {
        self.services.contains(name)
    }

    pub fn match_open(&self, open: &IncomingStreamOpen) -> Result<AcceptedService> {
        let Some(name) = service_name_from_target(&open.target_host) else {
            bail!("stream target {} is not a Dosh service", open.target_host);
        };
        if !self.services.contains(name) {
            bail!("service {name:?} is not registered");
        }
        if open.target_port != 0 {
            bail!(
                "service stream {} used non-zero target port",
                open.stream_id
            );
        }
        Ok(AcceptedService {
            name: name.to_string(),
            stream_id: open.stream_id,
        })
    }
}

impl StreamMux {
    pub fn new(config: TransportConfig) -> Self {
        Self {
            config,
            opened_streams: HashSet::new(),
            pending_opens: HashMap::new(),
            send_credit: HashMap::new(),
            pending_data: HashMap::new(),
            sent_data: HashMap::new(),
            next_send_offset: HashMap::new(),
            next_recv_offset: HashMap::new(),
            recv_pending: HashMap::new(),
        }
    }

    pub fn open_stream(
        &mut self,
        stream_id: u64,
        target_host: impl Into<String>,
        target_port: u16,
    ) -> Result<OutgoingStreamPacket> {
        let target_host = target_host.into();
        if self.opened_streams.contains(&stream_id) || self.pending_opens.contains_key(&stream_id) {
            bail!("stream {stream_id} already exists");
        }
        self.send_credit
            .insert(stream_id, self.config.initial_window);
        self.next_send_offset.entry(stream_id).or_insert(0);
        self.next_recv_offset.entry(stream_id).or_insert(0);
        self.pending_opens.insert(
            stream_id,
            PendingStreamOpen {
                target_host: target_host.clone(),
                target_port,
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        encode_packet(
            PacketKind::StreamOpen,
            &StreamOpen {
                stream_id,
                target_host,
                target_port,
            },
        )
    }

    pub fn accept_open(&mut self, open: StreamOpen) -> Result<IncomingStreamOpen> {
        self.opened_streams.insert(open.stream_id);
        self.send_credit
            .entry(open.stream_id)
            .or_insert(self.config.initial_window);
        self.next_send_offset.entry(open.stream_id).or_insert(0);
        self.next_recv_offset.entry(open.stream_id).or_insert(0);
        Ok(IncomingStreamOpen {
            stream_id: open.stream_id,
            target_host: open.target_host,
            target_port: open.target_port,
        })
    }

    pub fn open_ok(stream_id: u64) -> Result<OutgoingStreamPacket> {
        encode_packet(PacketKind::StreamOpenOk, &StreamOpenOk { stream_id })
    }

    pub fn open_reject(stream_id: u64, reason: impl Into<String>) -> Result<OutgoingStreamPacket> {
        encode_packet(
            PacketKind::StreamOpenReject,
            &StreamOpenReject {
                stream_id,
                reason: reason.into(),
            },
        )
    }

    pub fn handle_open_ok(&mut self, ok: StreamOpenOk) -> Result<Vec<OutgoingStreamPacket>> {
        self.pending_opens.remove(&ok.stream_id);
        self.opened_streams.insert(ok.stream_id);
        self.send_credit
            .entry(ok.stream_id)
            .or_insert(self.config.initial_window);
        self.next_send_offset.entry(ok.stream_id).or_insert(0);
        self.next_recv_offset.entry(ok.stream_id).or_insert(0);
        self.flush_pending_data(ok.stream_id)
    }

    pub fn handle_open_reject(&mut self, reject: StreamOpenReject) {
        self.cleanup_stream(reject.stream_id);
    }

    pub fn send_data(
        &mut self,
        stream_id: u64,
        bytes: impl Into<Vec<u8>>,
    ) -> Result<Vec<OutgoingStreamPacket>> {
        let bytes = bytes.into();
        if bytes.is_empty() {
            return Ok(Vec::new());
        }
        if !self.opened_streams.contains(&stream_id)
            || self.send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len()
            || self
                .pending_data
                .get(&stream_id)
                .is_some_and(|pending| !pending.is_empty())
        {
            self.pending_data
                .entry(stream_id)
                .or_default()
                .push_back(bytes);
            return Ok(Vec::new());
        }
        Ok(vec![self.send_data_now(stream_id, bytes)?])
    }

    pub fn handle_data(&mut self, data: StreamData) -> Result<IncomingStreamData> {
        let stream_id = data.stream_id;
        let (chunks, consumed, received_offset) = self.accept_data(data);
        let window_adjust = encode_packet(
            PacketKind::StreamWindowAdjust,
            &StreamWindowAdjust {
                stream_id,
                received_offset,
                bytes: consumed.min(u32::MAX as usize) as u32,
            },
        )?;
        Ok(IncomingStreamData {
            stream_id,
            chunks,
            consumed,
            received_offset,
            window_adjust,
        })
    }

    pub fn handle_window_adjust(
        &mut self,
        adjust: StreamWindowAdjust,
    ) -> Result<Vec<OutgoingStreamPacket>> {
        self.ack_data(adjust.stream_id, adjust.received_offset);
        self.add_credit(adjust.stream_id, adjust.bytes as usize);
        self.flush_pending_data(adjust.stream_id)
    }

    pub fn close_stream(&mut self, stream_id: u64) -> Result<OutgoingStreamPacket> {
        self.cleanup_stream(stream_id);
        encode_packet(PacketKind::StreamClose, &StreamClose { stream_id })
    }

    pub fn handle_close(&mut self, close: StreamClose) {
        self.cleanup_stream(close.stream_id);
    }

    pub fn handle_packet(&mut self, kind: PacketKind, body: &[u8]) -> Result<TransportEvent> {
        match kind {
            PacketKind::StreamOpen => {
                let open: StreamOpen = protocol::from_body(body)?;
                self.accept_open(open).map(TransportEvent::Open)
            }
            PacketKind::StreamOpenOk => {
                let ok: StreamOpenOk = protocol::from_body(body)?;
                let stream_id = ok.stream_id;
                let flushed = self.handle_open_ok(ok)?;
                Ok(TransportEvent::OpenOk { stream_id, flushed })
            }
            PacketKind::StreamOpenReject => {
                let reject: StreamOpenReject = protocol::from_body(body)?;
                self.handle_open_reject(reject.clone());
                Ok(TransportEvent::OpenReject {
                    stream_id: reject.stream_id,
                    reason: reject.reason,
                })
            }
            PacketKind::StreamData => {
                let data: StreamData = protocol::from_body(body)?;
                self.handle_data(data).map(TransportEvent::Data)
            }
            PacketKind::StreamWindowAdjust => {
                let adjust: StreamWindowAdjust = protocol::from_body(body)?;
                let stream_id = adjust.stream_id;
                let flushed = self.handle_window_adjust(adjust)?;
                Ok(TransportEvent::WindowAdjust { stream_id, flushed })
            }
            PacketKind::StreamClose => {
                let close: StreamClose = protocol::from_body(body)?;
                let stream_id = close.stream_id;
                self.handle_close(close);
                Ok(TransportEvent::Close { stream_id })
            }
            _ => bail!("packet kind {kind:?} is not a stream transport packet"),
        }
    }

    pub fn tick(&mut self) -> Result<Vec<OutgoingStreamPacket>> {
        let now = Instant::now();
        let mut out = Vec::new();
        let mut retransmit_opens = Vec::new();
        for (stream_id, pending) in self.pending_opens.iter_mut() {
            if now.duration_since(pending.last_sent) < self.config.retransmit_after {
                continue;
            }
            pending.last_sent = now;
            pending.attempts = pending.attempts.saturating_add(1);
            retransmit_opens.push((*stream_id, pending.target_host.clone(), pending.target_port));
        }
        for (stream_id, target_host, target_port) in retransmit_opens {
            out.push(encode_packet(
                PacketKind::StreamOpen,
                &StreamOpen {
                    stream_id,
                    target_host,
                    target_port,
                },
            )?);
        }

        let mut retransmit_data = Vec::new();
        for (stream_id, chunks) in self.sent_data.iter_mut() {
            for chunk in chunks.values_mut() {
                if now.duration_since(chunk.last_sent) < self.config.retransmit_after {
                    continue;
                }
                chunk.last_sent = now;
                chunk.attempts = chunk.attempts.saturating_add(1);
                retransmit_data.push((*stream_id, chunk.offset, chunk.bytes.clone()));
            }
        }
        for (stream_id, offset, bytes) in retransmit_data {
            out.push(encode_packet(
                PacketKind::StreamData,
                &StreamData {
                    stream_id,
                    offset,
                    bytes,
                },
            )?);
        }
        Ok(out)
    }

    pub fn is_open(&self, stream_id: u64) -> bool {
        self.opened_streams.contains(&stream_id)
    }

    pub fn pending_bytes(&self, stream_id: u64) -> usize {
        self.pending_data
            .get(&stream_id)
            .map(|queue| queue.iter().map(Vec::len).sum())
            .unwrap_or(0)
    }

    fn send_data_now(&mut self, stream_id: u64, bytes: Vec<u8>) -> Result<OutgoingStreamPacket> {
        let credit = self
            .send_credit
            .get_mut(&stream_id)
            .ok_or_else(|| anyhow!("stream {stream_id} has no send window"))?;
        if *credit < bytes.len() {
            bail!("stream {stream_id} send window exhausted");
        }
        *credit -= bytes.len();
        let offset = *self.next_send_offset.entry(stream_id).or_default();
        self.next_send_offset
            .insert(stream_id, offset.saturating_add(bytes.len() as u64));
        self.sent_data.entry(stream_id).or_default().insert(
            offset,
            PendingStreamChunk {
                offset,
                bytes: bytes.clone(),
                last_sent: Instant::now(),
                attempts: 1,
            },
        );
        encode_packet(
            PacketKind::StreamData,
            &StreamData {
                stream_id,
                offset,
                bytes,
            },
        )
    }

    fn flush_pending_data(&mut self, stream_id: u64) -> Result<Vec<OutgoingStreamPacket>> {
        let mut out = Vec::new();
        while let Some(front_len) = self
            .pending_data
            .get(&stream_id)
            .and_then(|pending| pending.front().map(Vec::len))
        {
            if self.send_credit.get(&stream_id).copied().unwrap_or(0) < front_len {
                break;
            }
            let bytes = self
                .pending_data
                .get_mut(&stream_id)
                .and_then(VecDeque::pop_front)
                .expect("pending front exists");
            if self
                .pending_data
                .get(&stream_id)
                .is_some_and(VecDeque::is_empty)
            {
                self.pending_data.remove(&stream_id);
            }
            out.push(self.send_data_now(stream_id, bytes)?);
        }
        Ok(out)
    }

    fn accept_data(&mut self, data: StreamData) -> (Vec<Vec<u8>>, usize, u64) {
        let stream_id = data.stream_id;
        let expected = self.next_recv_offset.entry(stream_id).or_insert(0);
        if data.offset < *expected {
            return (Vec::new(), 0, *expected);
        }
        if data.offset > *expected {
            self.recv_pending
                .entry(stream_id)
                .or_default()
                .entry(data.offset)
                .or_insert(data.bytes);
            return (Vec::new(), 0, *expected);
        }

        let mut chunks = vec![data.bytes];
        let mut consumed = chunks[0].len();
        *expected = expected.saturating_add(consumed as u64);
        while let Some(bytes) = self
            .recv_pending
            .entry(stream_id)
            .or_default()
            .remove(expected)
        {
            consumed = consumed.saturating_add(bytes.len());
            *expected = expected.saturating_add(bytes.len() as u64);
            chunks.push(bytes);
        }
        if self
            .recv_pending
            .get(&stream_id)
            .is_some_and(BTreeMap::is_empty)
        {
            self.recv_pending.remove(&stream_id);
        }
        (chunks, consumed, *expected)
    }

    fn ack_data(&mut self, stream_id: u64, received_offset: u64) {
        let Some(sent) = self.sent_data.get_mut(&stream_id) else {
            return;
        };
        let acked_offsets = sent
            .iter()
            .filter_map(|(offset, chunk)| {
                let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
                (end <= received_offset).then_some(*offset)
            })
            .collect::<Vec<_>>();
        for offset in acked_offsets {
            sent.remove(&offset);
        }
        if sent.is_empty() {
            self.sent_data.remove(&stream_id);
        }
    }

    fn add_credit(&mut self, stream_id: u64, bytes: usize) {
        let credit = self.send_credit.entry(stream_id).or_default();
        *credit = credit.saturating_add(bytes).min(self.config.initial_window);
    }

    fn cleanup_stream(&mut self, stream_id: u64) {
        self.opened_streams.remove(&stream_id);
        self.pending_opens.remove(&stream_id);
        self.send_credit.remove(&stream_id);
        self.pending_data.remove(&stream_id);
        self.sent_data.remove(&stream_id);
        self.next_send_offset.remove(&stream_id);
        self.next_recv_offset.remove(&stream_id);
        self.recv_pending.remove(&stream_id);
    }
}

pub struct DoshTransport {
    socket: Arc<UdpSocket>,
    role: SessionRole,
    conn_id: [u8; 16],
    session_key: [u8; 32],
    peer_addr: SocketAddr,
    send_seq: u64,
    ack_seq: u64,
    replay: ReplayWindow,
    mux: StreamMux,
    next_stream_id: u64,
    last_contact: Instant,
    last_keepalive: Instant,
}

impl DoshTransport {
    pub fn new(socket: Arc<UdpSocket>, config: SessionTransportConfig) -> Self {
        let now = Instant::now();
        Self {
            socket,
            role: config.role,
            conn_id: config.conn_id,
            session_key: config.session_key,
            peer_addr: config.peer_addr,
            send_seq: config.initial_send_seq,
            ack_seq: config.initial_ack,
            replay: ReplayWindow::default(),
            mux: StreamMux::new(config.stream),
            next_stream_id: 1,
            last_contact: now,
            last_keepalive: now,
        }
    }

    pub fn new_owned(socket: UdpSocket, config: SessionTransportConfig) -> Self {
        Self::new(Arc::new(socket), config)
    }

    pub fn local_addr(&self) -> Result<SocketAddr> {
        Ok(self.socket.local_addr()?)
    }

    pub fn conn_id(&self) -> [u8; 16] {
        self.conn_id
    }

    pub fn mux(&self) -> &StreamMux {
        &self.mux
    }

    pub fn mux_mut(&mut self) -> &mut StreamMux {
        &mut self.mux
    }

    pub fn peer_addr(&self) -> SocketAddr {
        self.peer_addr
    }

    pub fn set_peer_addr(&mut self, peer_addr: SocketAddr) {
        self.peer_addr = peer_addr;
    }

    pub fn last_contact(&self) -> Instant {
        self.last_contact
    }

    pub fn stale_for(&self) -> Duration {
        self.last_contact.elapsed()
    }

    pub fn allocate_stream_id(&mut self) -> u64 {
        let id = self.next_stream_id;
        self.next_stream_id = self.next_stream_id.saturating_add(1).max(1);
        id
    }

    pub async fn open_service(&mut self, name: &str) -> Result<u64> {
        let stream_id = self.allocate_stream_id();
        let target = service_target(name)?;
        let packet = self.mux.open_stream(stream_id, target, 0)?;
        self.send_outgoing(packet).await?;
        Ok(stream_id)
    }

    pub async fn open_target(
        &mut self,
        target_host: impl Into<String>,
        target_port: u16,
    ) -> Result<u64> {
        let stream_id = self.allocate_stream_id();
        let packet = self
            .mux
            .open_stream(stream_id, target_host.into(), target_port)?;
        self.send_outgoing(packet).await?;
        Ok(stream_id)
    }

    pub async fn accept_stream(&mut self, stream_id: u64) -> Result<()> {
        self.send_outgoing(StreamMux::open_ok(stream_id)?).await
    }

    pub async fn reject_stream(&mut self, stream_id: u64, reason: impl Into<String>) -> Result<()> {
        self.send_outgoing(StreamMux::open_reject(stream_id, reason)?)
            .await
    }

    pub async fn send(&mut self, stream_id: u64, bytes: impl Into<Vec<u8>>) -> Result<()> {
        for packet in self.mux.send_data(stream_id, bytes)? {
            self.send_outgoing(packet).await?;
        }
        Ok(())
    }

    pub async fn close(&mut self, stream_id: u64) -> Result<()> {
        let packet = self.mux.close_stream(stream_id)?;
        self.send_outgoing(packet).await
    }

    pub async fn maintenance(&mut self) -> Result<()> {
        for packet in self.mux.tick()? {
            self.send_outgoing(packet).await?;
        }
        if self.last_keepalive.elapsed() >= self.mux.config.keepalive_after {
            self.send_kind(PacketKind::Ping, b"").await?;
            self.last_keepalive = Instant::now();
        }
        Ok(())
    }

    pub async fn recv(&mut self) -> Result<SessionEvent> {
        let mut buf = vec![0u8; 65535];
        loop {
            let (n, peer) = self.socket.recv_from(&mut buf).await?;
            let packet = match protocol::decode(&buf[..n]) {
                Ok(packet) => packet,
                Err(_) => continue,
            };
            if packet.header.conn_id != self.conn_id {
                continue;
            }
            if !self.replay.accept(packet.header.seq) {
                continue;
            }
            let plain =
                protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction())?;
            self.peer_addr = peer;
            self.ack_seq = self.ack_seq.max(packet.header.seq);
            self.last_contact = Instant::now();
            return self
                .handle_decrypted_packet(packet.header.kind, &plain)
                .await;
        }
    }

    pub async fn handle_datagram(
        &mut self,
        datagram: &[u8],
        peer: SocketAddr,
    ) -> Result<SessionEvent> {
        let packet = protocol::decode(datagram)?;
        if packet.header.conn_id != self.conn_id {
            return Ok(SessionEvent::Ignored);
        }
        if !self.replay.accept(packet.header.seq) {
            return Ok(SessionEvent::Ignored);
        }
        let plain = protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction())?;
        self.peer_addr = peer;
        self.ack_seq = self.ack_seq.max(packet.header.seq);
        self.last_contact = Instant::now();
        self.handle_decrypted_packet(packet.header.kind, &plain)
            .await
    }

    async fn handle_decrypted_packet(
        &mut self,
        kind: PacketKind,
        plain: &[u8],
    ) -> Result<SessionEvent> {
        match kind {
            PacketKind::Ping => {
                self.send_kind(PacketKind::Pong, b"").await?;
                Ok(SessionEvent::Ping)
            }
            PacketKind::Pong => Ok(SessionEvent::Pong),
            PacketKind::StreamOpen
            | PacketKind::StreamOpenOk
            | PacketKind::StreamOpenReject
            | PacketKind::StreamData
            | PacketKind::StreamWindowAdjust
            | PacketKind::StreamClose => {
                let event = self.mux.handle_packet(kind, plain)?;
                self.send_followups(&event).await?;
                Ok(SessionEvent::Stream(event))
            }
            _ => Ok(SessionEvent::Ignored),
        }
    }

    async fn send_followups(&mut self, event: &TransportEvent) -> Result<()> {
        match event {
            TransportEvent::OpenOk { flushed, .. }
            | TransportEvent::WindowAdjust { flushed, .. } => {
                for packet in flushed.clone() {
                    self.send_outgoing(packet).await?;
                }
            }
            TransportEvent::Data(data) => {
                self.send_outgoing(data.window_adjust.clone()).await?;
            }
            TransportEvent::Open(_)
            | TransportEvent::OpenReject { .. }
            | TransportEvent::Close { .. } => {}
        }
        Ok(())
    }

    async fn send_outgoing(&mut self, packet: OutgoingStreamPacket) -> Result<()> {
        self.send_kind(packet.kind, &packet.body).await
    }

    async fn send_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result<()> {
        let encoded = self.encode_kind(kind, body)?;
        self.socket.send_to(&encoded, self.peer_addr).await?;
        Ok(())
    }

    fn encode_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result<Vec<u8>> {
        let seq = self.send_seq;
        self.send_seq = self.send_seq.saturating_add(1).max(1);
        protocol::encode_encrypted(
            kind,
            self.conn_id,
            seq,
            self.ack_seq,
            &self.session_key,
            self.role.send_direction(),
            body,
        )
    }
}

pub fn service_target(name: &str) -> Result<String> {
    if name.is_empty()
        || !name
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_'))
    {
        bail!("invalid Dosh service name {name:?}");
    }
    Ok(format!("{SERVICE_TARGET_PREFIX}{name}"))
}

pub fn is_service_target(target: &str) -> bool {
    target.starts_with(SERVICE_TARGET_PREFIX)
}

pub fn service_name_from_target(target: &str) -> Option<&str> {
    target.strip_prefix(SERVICE_TARGET_PREFIX)
}

fn encode_packet<T: serde::Serialize>(kind: PacketKind, value: &T) -> Result<OutgoingStreamPacket> {
    Ok(OutgoingStreamPacket {
        kind,
        body: protocol::to_body(value)?,
    })
}

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

    fn decode<T: serde::de::DeserializeOwned>(packet: &OutgoingStreamPacket) -> T {
        protocol::from_body(&packet.body).unwrap()
    }

    #[test]
    fn service_targets_are_reserved_and_validated() {
        assert_eq!(service_target("chat_api").unwrap(), "@dosh-chat_api");
        assert!(is_service_target("@dosh-file"));
        assert_eq!(service_name_from_target("@dosh-chat"), Some("chat"));
        assert!(service_target("../bad").is_err());
    }

    #[test]
    fn service_registry_matches_registered_service_opens() {
        let mut registry = ServiceRegistry::new();
        registry.register("chat").unwrap();
        let accepted = registry
            .match_open(&IncomingStreamOpen {
                stream_id: 4,
                target_host: "@dosh-chat".to_string(),
                target_port: 0,
            })
            .unwrap();
        assert_eq!(
            accepted,
            AcceptedService {
                name: "chat".to_string(),
                stream_id: 4,
            }
        );
        assert!(
            registry
                .match_open(&IncomingStreamOpen {
                    stream_id: 5,
                    target_host: "@dosh-other".to_string(),
                    target_port: 0,
                })
                .is_err()
        );
    }

    #[test]
    fn opens_stream_and_flushes_queued_data_after_open_ok() {
        let mut mux = StreamMux::new(TransportConfig::default());
        let open = mux
            .open_stream(7, service_target("chat").unwrap(), 0)
            .unwrap();
        assert_eq!(open.kind, PacketKind::StreamOpen);
        assert!(mux.send_data(7, b"hello".to_vec()).unwrap().is_empty());
        assert_eq!(mux.pending_bytes(7), 5);

        let flushed = mux.handle_open_ok(StreamOpenOk { stream_id: 7 }).unwrap();
        assert_eq!(flushed.len(), 1);
        assert_eq!(flushed[0].kind, PacketKind::StreamData);
        let data: StreamData = decode(&flushed[0]);
        assert_eq!(data.offset, 0);
        assert_eq!(data.bytes, b"hello");
    }

    #[test]
    fn receives_out_of_order_data_in_order_and_acks_offsets() {
        let mut mux = StreamMux::new(TransportConfig::default());
        mux.accept_open(StreamOpen {
            stream_id: 3,
            target_host: "@dosh-test".to_string(),
            target_port: 0,
        })
        .unwrap();
        let first = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 5,
                bytes: b"world".to_vec(),
            })
            .unwrap();
        assert!(first.chunks.is_empty());
        assert_eq!(first.received_offset, 0);

        let second = mux
            .handle_data(StreamData {
                stream_id: 3,
                offset: 0,
                bytes: b"hello".to_vec(),
            })
            .unwrap();
        assert_eq!(second.chunks, vec![b"hello".to_vec(), b"world".to_vec()]);
        assert_eq!(second.consumed, 10);
        assert_eq!(second.received_offset, 10);
    }

    #[test]
    fn window_adjust_acks_sent_data_and_releases_credit() {
        let mut mux = StreamMux::new(TransportConfig {
            initial_window: 5,
            retransmit_after: Duration::from_secs(1),
            ..TransportConfig::default()
        });
        mux.open_stream(1, "@dosh-test", 0).unwrap();
        mux.handle_open_ok(StreamOpenOk { stream_id: 1 }).unwrap();
        let sent = mux.send_data(1, b"hello".to_vec()).unwrap();
        assert_eq!(sent.len(), 1);
        assert!(mux.send_data(1, b"!".to_vec()).unwrap().is_empty());
        assert_eq!(mux.pending_bytes(1), 1);
        let flushed = mux
            .handle_window_adjust(StreamWindowAdjust {
                stream_id: 1,
                received_offset: 5,
                bytes: 5,
            })
            .unwrap();
        assert_eq!(flushed.len(), 1);
        let data: StreamData = decode(&flushed[0]);
        assert_eq!(data.offset, 5);
        assert_eq!(data.bytes, b"!");
    }

    #[test]
    fn handle_packet_decodes_and_dispatches_stream_packets() {
        let mut mux = StreamMux::new(TransportConfig::default());
        let body = protocol::to_body(&StreamOpen {
            stream_id: 9,
            target_host: "@dosh-app".to_string(),
            target_port: 0,
        })
        .unwrap();
        let event = mux.handle_packet(PacketKind::StreamOpen, &body).unwrap();
        assert_eq!(
            event,
            TransportEvent::Open(IncomingStreamOpen {
                stream_id: 9,
                target_host: "@dosh-app".to_string(),
                target_port: 0,
            })
        );
    }

    #[tokio::test]
    async fn session_transport_opens_service_and_moves_data_over_encrypted_udp() {
        let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let client_addr = client_socket.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [42u8; 32];
        let conn_id = [7u8; 16];
        let mut client = DoshTransport::new_owned(
            client_socket,
            SessionTransportConfig {
                role: SessionRole::Client,
                conn_id,
                session_key: key,
                peer_addr: server_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: client_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );

        let stream_id = client.open_service("app").await.unwrap();
        let event = server.recv().await.unwrap();
        assert_eq!(server.peer_addr(), client_addr);
        match event {
            SessionEvent::Stream(TransportEvent::Open(open)) => {
                assert_eq!(open.stream_id, stream_id);
                assert_eq!(open.target_host, "@dosh-app");
                server.accept_stream(open.stream_id).await.unwrap();
            }
            other => panic!("unexpected event {other:?}"),
        }

        assert!(matches!(
            client.recv().await.unwrap(),
            SessionEvent::Stream(TransportEvent::OpenOk { .. })
        ));
        client.send(stream_id, b"hello".to_vec()).await.unwrap();
        let event = server.recv().await.unwrap();
        match event {
            SessionEvent::Stream(TransportEvent::Data(data)) => {
                assert_eq!(data.stream_id, stream_id);
                assert_eq!(data.chunks, vec![b"hello".to_vec()]);
            }
            other => panic!("unexpected event {other:?}"),
        }
    }

    #[tokio::test]
    async fn session_transport_roams_to_authenticated_new_peer_addr() {
        let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let original_client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let roaming_client = UdpSocket::bind("127.0.0.1:0").await.unwrap();
        let original_addr = original_client.local_addr().unwrap();
        let roaming_addr = roaming_client.local_addr().unwrap();
        let server_addr = server_socket.local_addr().unwrap();
        let key = [9u8; 32];
        let conn_id = [3u8; 16];
        let mut server = DoshTransport::new_owned(
            server_socket,
            SessionTransportConfig {
                role: SessionRole::Server,
                conn_id,
                session_key: key,
                peer_addr: original_addr,
                initial_send_seq: 1,
                initial_ack: 0,
                stream: TransportConfig::default(),
            },
        );
        let ping = protocol::encode_encrypted(
            PacketKind::Ping,
            conn_id,
            1,
            0,
            &key,
            CLIENT_TO_SERVER,
            b"",
        )
        .unwrap();
        roaming_client.send_to(&ping, server_addr).await.unwrap();
        assert!(matches!(server.recv().await.unwrap(), SessionEvent::Ping));
        assert_eq!(server.peer_addr(), roaming_addr);
    }
}