swactor/crates/iroh-driver/src/datastream_transport.rs

//! Iroh/QUIC transport adapter for datastream subscriptions.

use std::error::Error;
use std::sync::Arc;
use std::time::Duration;

use crossbeam_channel::TryRecvError;
use datastream::{
    ChannelDescriptor, ChannelId, ChannelRef, DatastreamEvent, DatastreamSnapshot,
    DatastreamSubscription, FrameDelivery, Position, StreamDescriptor,
};
use iroh::endpoint::{Connection, RecvStream, SendStream};
use iroh::{Endpoint, EndpointAddr};
use tokio::runtime::Handle;

pub const DATASTREAM_ALPN: &[u8] = b"swactor/datastream/0";

const MAGIC: &[u8; 4] = b"DSQ1";
const TAG_CHANNEL_DECLARED: u8 = 0x01;
const TAG_FRAME: u8 = 0x02;
const TAG_STREAM_ENDED: u8 = 0x03;
const MAX_RECORD_BYTES: usize = 16 * 1024 * 1024;

type BoxError = Box<dyn Error + Send + Sync + 'static>;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatastreamQuicHeader {
    pub flow_id: [u8; 16],
    pub token: Vec<u8>,
    pub stream: StreamDescriptor,
    pub channels: Vec<ChannelDescriptor>,
}

impl DatastreamQuicHeader {
    pub fn new(
        flow_id: [u8; 16],
        token: impl Into<Vec<u8>>,
        stream: StreamDescriptor,
        channels: Vec<ChannelDescriptor>,
    ) -> Self {
        Self {
            flow_id,
            token: token.into(),
            stream,
            channels,
        }
    }

    pub fn from_snapshot(
        flow_id: [u8; 16],
        token: impl Into<Vec<u8>>,
        snapshot: &DatastreamSnapshot,
    ) -> Result<Self, BoxError> {
        let stream = snapshot
            .streams
            .first()
            .cloned()
            .ok_or("datastream subscription snapshot has no stream")?;
        Ok(Self::new(flow_id, token, stream, snapshot.channels.clone()))
    }
}

#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct DatastreamQuicWriteStats {
    pub events: usize,
    pub bytes: usize,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatastreamQuicRead {
    pub header: DatastreamQuicHeader,
    pub events: Vec<DatastreamEvent>,
}

pub fn spawn_subscription_writer(
    handle: &Handle,
    endpoint: Endpoint,
    peer: EndpointAddr,
    header: DatastreamQuicHeader,
    subscription: DatastreamSubscription,
    idle_sleep: Duration,
) -> tokio::task::JoinHandle<Result<DatastreamQuicWriteStats, String>> {
    handle.spawn(async move {
        let conn = endpoint
            .connect(peer, DATASTREAM_ALPN)
            .await
            .map_err(|error| error.to_string())?;
        let send = conn.open_uni().await.map_err(|error| error.to_string())?;
        write_subscription_until_closed(send, header, subscription, idle_sleep)
            .await
            .map_err(|error| error.to_string())
    })
}

pub async fn write_available_subscription(
    send: SendStream,
    header: &DatastreamQuicHeader,
    subscription: &DatastreamSubscription,
) -> Result<DatastreamQuicWriteStats, BoxError> {
    write_subscription_inner(send, header, subscription, None).await
}

pub async fn write_subscription_until_closed(
    mut send: SendStream,
    header: DatastreamQuicHeader,
    subscription: DatastreamSubscription,
    idle_sleep: Duration,
) -> Result<DatastreamQuicWriteStats, BoxError> {
    write_header(&mut send, &header).await?;
    let mut stats = DatastreamQuicWriteStats::default();
    loop {
        match subscription.try_recv() {
            Ok(event) => {
                let bytes = write_event(&mut send, &event).await?;
                if bytes > 0 {
                    stats.bytes += bytes;
                    stats.events += 1;
                }
            }
            Err(TryRecvError::Empty) => {
                tokio::time::sleep(idle_sleep).await;
            }
            Err(TryRecvError::Disconnected) => break,
        }
    }
    send.finish()?;
    Ok(stats)
}

async fn write_subscription_inner(
    mut send: SendStream,
    header: &DatastreamQuicHeader,
    subscription: &DatastreamSubscription,
    idle_sleep: Option<Duration>,
) -> Result<DatastreamQuicWriteStats, BoxError> {
    write_header(&mut send, header).await?;
    let mut stats = DatastreamQuicWriteStats::default();
    loop {
        match subscription.try_recv() {
            Ok(event) => {
                let bytes = write_event(&mut send, &event).await?;
                if bytes > 0 {
                    stats.bytes += bytes;
                    stats.events += 1;
                }
            }
            Err(TryRecvError::Empty) => match idle_sleep {
                Some(delay) => tokio::time::sleep(delay).await,
                None => break,
            },
            Err(TryRecvError::Disconnected) => break,
        }
    }
    send.finish()?;
    Ok(stats)
}

pub async fn write_event(
    send: &mut SendStream,
    event: &DatastreamEvent,
) -> Result<usize, BoxError> {
    let mut bytes = Vec::new();
    match event {
        DatastreamEvent::StreamDeclared(_) => return Ok(0),
        DatastreamEvent::ChannelDeclared(descriptor) => {
            bytes.push(TAG_CHANNEL_DECLARED);
            put_json(&mut bytes, descriptor)?;
        }
        DatastreamEvent::Frame(delivery) => {
            bytes.push(TAG_FRAME);
            bytes.extend_from_slice(&delivery.channel.channel.0.to_le_bytes());
            bytes.extend_from_slice(&delivery.position.0.to_le_bytes());
            put_bytes(&mut bytes, &delivery.payload)?;
        }
        DatastreamEvent::StreamEnded(_) => {
            bytes.push(TAG_STREAM_ENDED);
        }
    }
    if bytes.len() > MAX_RECORD_BYTES {
        return Err("datastream QUIC record exceeds max size".into());
    }
    send.write_all(&(bytes.len() as u32).to_le_bytes()).await?;
    send.write_all(&bytes).await?;
    Ok(4 + bytes.len())
}

pub async fn read_stream_header(recv: &mut RecvStream) -> Result<DatastreamQuicHeader, BoxError> {
    read_header(recv).await
}

pub async fn read_events_from_stream(mut recv: RecvStream) -> Result<DatastreamQuicRead, BoxError> {
    let header = read_header(&mut recv).await?;
    let mut events = Vec::new();
    loop {
        match read_next_event(&mut recv, &header.stream).await? {
            Some(event) => events.push(event),
            None => break,
        }
    }
    Ok(DatastreamQuicRead { header, events })
}

pub async fn read_next_uni_from_connection(
    conn: &Connection,
) -> Result<DatastreamQuicRead, BoxError> {
    let recv = conn.accept_uni().await?;
    read_events_from_stream(recv).await
}

pub fn spawn_connection_reader(
    handle: &Handle,
    conn: Connection,
    sink: std::sync::mpsc::Sender<DatastreamEvent>,
) -> tokio::task::JoinHandle<Result<(), String>> {
    handle.spawn(async move {
        loop {
            let recv = match conn.accept_uni().await {
                Ok(recv) => recv,
                Err(error) => return Err(error.to_string()),
            };
            let read = read_events_from_stream(recv)
                .await
                .map_err(|error| error.to_string())?;
            for event in read.events {
                if sink.send(event).is_err() {
                    return Ok(());
                }
            }
        }
    })
}

async fn write_header(
    send: &mut SendStream,
    header: &DatastreamQuicHeader,
) -> Result<(), BoxError> {
    if header.token.len() > u16::MAX as usize {
        return Err("datastream token exceeds u16 length prefix".into());
    }
    send.write_all(MAGIC).await?;
    send.write_all(&header.flow_id).await?;
    send.write_all(&(header.token.len() as u16).to_le_bytes())
        .await?;
    send.write_all(&header.token).await?;
    write_json(send, &header.stream).await?;
    write_json(send, &header.channels).await?;
    Ok(())
}

async fn read_header(recv: &mut RecvStream) -> Result<DatastreamQuicHeader, BoxError> {
    let mut magic = [0u8; 4];
    recv.read_exact(&mut magic).await?;
    if &magic != MAGIC {
        return Err("invalid datastream QUIC magic".into());
    }
    let mut flow_id = [0u8; 16];
    recv.read_exact(&mut flow_id).await?;
    let mut token_len = [0u8; 2];
    recv.read_exact(&mut token_len).await?;
    let token_len = u16::from_le_bytes(token_len) as usize;
    let mut token = vec![0u8; token_len];
    recv.read_exact(&mut token).await?;
    let stream = read_json(recv).await?;
    let channels = read_json(recv).await?;
    Ok(DatastreamQuicHeader {
        flow_id,
        token,
        stream,
        channels,
    })
}

pub async fn read_next_event(
    recv: &mut RecvStream,
    stream: &StreamDescriptor,
) -> Result<Option<DatastreamEvent>, BoxError> {
    let mut len = [0u8; 4];
    if recv.read_exact(&mut len).await.is_err() {
        return Ok(None);
    }
    let len = u32::from_le_bytes(len) as usize;
    if len > MAX_RECORD_BYTES {
        return Err("datastream QUIC record exceeds max size".into());
    }
    let mut buf = vec![0u8; len];
    recv.read_exact(&mut buf).await?;
    decode_record(&buf, stream).map(Some)
}

fn decode_record(buf: &[u8], stream: &StreamDescriptor) -> Result<DatastreamEvent, BoxError> {
    if buf.is_empty() {
        return Err("empty datastream QUIC record".into());
    }
    match buf[0] {
        TAG_CHANNEL_DECLARED => {
            let descriptor: ChannelDescriptor = serde_json::from_slice(&buf[1..])?;
            Ok(DatastreamEvent::ChannelDeclared(descriptor))
        }
        TAG_FRAME => {
            if buf.len() < 1 + 4 + 8 + 4 {
                return Err("datastream QUIC frame record truncated".into());
            }
            let channel = ChannelId(u32::from_le_bytes([buf[1], buf[2], buf[3], buf[4]]));
            let position = Position(u64::from_le_bytes([
                buf[5], buf[6], buf[7], buf[8], buf[9], buf[10], buf[11], buf[12],
            ]));
            let mut len = [0u8; 4];
            len.copy_from_slice(&buf[13..17]);
            let payload_len = u32::from_le_bytes(len) as usize;
            let payload = buf
                .get(17..17 + payload_len)
                .ok_or("datastream QUIC frame payload truncated")?;
            if 17 + payload_len != buf.len() {
                return Err("bytes remain after datastream QUIC frame record".into());
            }
            Ok(DatastreamEvent::Frame(FrameDelivery {
                channel: ChannelRef {
                    stream: stream.stream.clone(),
                    channel,
                },
                position,
                payload: payload.to_vec(),
            }))
        }
        TAG_STREAM_ENDED => Ok(DatastreamEvent::StreamEnded(stream.stream.clone())),
        _ => Err("unknown datastream QUIC record tag".into()),
    }
}

pub async fn read_stream_into_fanout(
    recv: RecvStream,
    fanout: Arc<datastream::DeliveryFanout>,
) -> Result<DatastreamQuicHeader, BoxError> {
    let read = read_events_from_stream(recv).await?;
    fanout.publish_batch(read.events);
    Ok(read.header)
}

fn put_json<T: serde::Serialize>(out: &mut Vec<u8>, value: &T) -> Result<(), BoxError> {
    out.extend_from_slice(&serde_json::to_vec(value)?);
    Ok(())
}

fn put_bytes(out: &mut Vec<u8>, bytes: &[u8]) -> Result<(), BoxError> {
    if bytes.len() > u32::MAX as usize {
        return Err("datastream delivery exceeds u32 length prefix".into());
    }
    out.extend_from_slice(&(bytes.len() as u32).to_le_bytes());
    out.extend_from_slice(bytes);
    Ok(())
}

async fn write_json<T: serde::Serialize>(send: &mut SendStream, value: &T) -> Result<(), BoxError> {
    let bytes = serde_json::to_vec(value)?;
    if bytes.len() > u32::MAX as usize {
        return Err("datastream header JSON exceeds u32 length prefix".into());
    }
    send.write_all(&(bytes.len() as u32).to_le_bytes()).await?;
    send.write_all(&bytes).await?;
    Ok(())
}

async fn read_json<T: serde::de::DeserializeOwned>(recv: &mut RecvStream) -> Result<T, BoxError> {
    let mut len = [0u8; 4];
    recv.read_exact(&mut len).await?;
    let len = u32::from_le_bytes(len) as usize;
    if len > MAX_RECORD_BYTES {
        return Err("datastream QUIC record exceeds max size".into());
    }
    let mut bytes = vec![0u8; len];
    recv.read_exact(&mut bytes).await?;
    Ok(serde_json::from_slice(&bytes)?)
}