dope_session/protocol/

server.rs

use crate::{CodecLayer, SlotGen, SlotId};
use std::pin::Pin;

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum RouteRequestBodyKind {
    Inline,
    Spilled,
    Stream,
}

/// Server-side protocol.
///
/// One trait owns parsing AND handling — there is no separate `Framer` trait
/// or `framer()` accessor. Per-connection parsing state lives in
/// [`ConnState`]; shared protocol state (rate limiters, metrics, etc.) lives
/// in `&mut self`.
///
/// `Head<'buf>` is a GAT so parsers can return zero-copy views into `buf`.
///
/// [`ConnState`]: ServerProtocol::ConnState
pub trait ServerProtocol: 'static {
    type Head<'buf>;

    type ConnState: Default + 'static;

    type CodecLayer: CodecLayer;

    const SUPPORTS_PARK: bool = false;

    /// Try to parse one frame's head from `buf`. Returns
    /// `(head, head_len, body_len)` on success — frame ring waits for
    /// `head_len + body_len` bytes total before dispatching to `handle`.
    /// `body_len = 0` means a head-only frame.
    fn parse<'buf>(
        &self,
        state: &mut Self::ConnState,
        buf: &'buf [u8],
    ) -> Option<(Self::Head<'buf>, u32, u32)>;

    fn handle<'buf>(
        &mut self,
        req: &'buf [u8],
        head: Self::Head<'buf>,
        write: &mut [u8],
    ) -> HandlerAction;

    fn handle_park<'buf>(
        &mut self,
        conn_id: SlotId,
        conn_state: &mut Self::ConnState,
        req: &'buf [u8],
        head: Self::Head<'buf>,
        write: &mut [u8],
    ) -> ParkAction {
        let _ = (conn_id, conn_state, req, head, write);
        unreachable!("ServerProtocol::handle_park called with SUPPORTS_PARK = true but no override")
    }

    fn oversize_response(&self) -> &'static [u8] {
        b""
    }

    fn route_request_body_kind(&self, head: &Self::Head<'_>) -> RouteRequestBodyKind {
        let _ = head;
        RouteRequestBodyKind::Inline
    }

    fn handle_request_stream<'buf>(
        &mut self,
        head_bytes: &'buf [u8],
        head: Self::Head<'buf>,
        body_stream: RequestBodyStream,
        write: &mut [u8],
    ) -> RequestStreamAction {
        let _ = (head_bytes, head, body_stream, write);
        unreachable!(
            "ServerProtocol::handle_request_stream called without an override; \
             a route returned RouteRequestBodyKind::Stream but the \
             handler does not implement the stream dispatch path"
        )
    }

    fn on_register(
        &mut self,
        conn_id: SlotId,
        conn_state: &mut Self::ConnState,
    ) -> Option<Pin<Box<dyn std::future::Future<Output = ()> + 'static>>> {
        let _ = (conn_id, conn_state);
        None
    }

    fn drain_outbound(&mut self, conn_state: &mut Self::ConnState, write: &mut [u8]) -> u32 {
        let _ = (conn_state, write);
        0
    }

    fn close_pending(&self, conn_state: &Self::ConnState) -> bool {
        let _ = conn_state;
        false
    }
}

pub enum ParkAction {
    Send { written: u16, close_after: bool },
    Park,
    Close(&'static [u8]),
}

pub enum RequestStreamAction {
    Pending(dambi::InlineFuture<'static, PendingResponse, 64>),
    Close(&'static [u8]),
}

pub enum HandlerAction {
    Send {
        written: u16,
        close_after: bool,
    },
    SendStatic {
        hdr_written: u32,
        body: &'static [u8],
        close_after: bool,
    },

    SendVectored {
        iovs: [dope_core::IoVec; 4],
        iov_count: u8,
        cookie: *const u8,
        total_bytes: u32,
        close_after: bool,
    },
    SendStream {
        hdr_written: u32,
        stream: Pin<Box<dyn ResponseChunkStream>>,
        close_after: bool,
    },
    Pending(dambi::InlineFuture<'static, PendingResponse, 64>),
    Close(&'static [u8]),
}

pub trait ResponseChunkStream: 'static {
    fn poll_chunk(
        self: Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<dambi::Bytes>>;
}

pub struct RequestBodyStream {
    inner: std::rc::Rc<std::cell::RefCell<RequestStreamInner>>,
}

#[derive(Default)]
pub(crate) struct RequestStreamInner {
    pub(crate) pending: std::collections::VecDeque<dambi::Bytes>,
    pub(crate) waker: Option<std::task::Waker>,
    pub(crate) received: u32,
    pub(crate) target: u32,
    pub(crate) eof: bool,
}

impl RequestBodyStream {
    pub(crate) fn new(target: u32) -> (Self, RequestStreamHandle) {
        let inner = std::rc::Rc::new(std::cell::RefCell::new(RequestStreamInner {
            pending: std::collections::VecDeque::new(),
            waker: None,
            received: 0,
            target,
            eof: target == 0,
        }));
        (
            Self {
                inner: inner.clone(),
            },
            RequestStreamHandle { inner },
        )
    }

    pub fn poll_chunk(
        &mut self,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<dambi::Bytes>> {
        let mut state = self.inner.borrow_mut();
        if let Some(chunk) = state.pending.pop_front() {
            return std::task::Poll::Ready(Some(chunk));
        }
        if state.eof {
            return std::task::Poll::Ready(None);
        }
        state.waker = Some(cx.waker().clone());
        std::task::Poll::Pending
    }
}

pub(crate) struct RequestStreamHandle {
    pub(crate) inner: std::rc::Rc<std::cell::RefCell<RequestStreamInner>>,
}

const REQUEST_STREAM_QUEUE_CAP: usize = 8 * 1024 * 1024;

impl RequestStreamHandle {
    pub(crate) fn push_slice(&self, src: &[u8]) -> usize {
        let mut state = self.inner.borrow_mut();
        if src.is_empty() {
            return 0;
        }
        let queued: usize = state.pending.iter().map(|b| b.len()).sum();
        if queued >= REQUEST_STREAM_QUEUE_CAP {
            return 0;
        }
        let needed = (state.target.saturating_sub(state.received)) as usize;
        let take = src.len().min(needed);
        if take > 0 {
            state
                .pending
                .push_back(dambi::Bytes::copy_from_slice(&src[..take]));
            state.received = state.received.saturating_add(take as u32);
        }
        if state.received >= state.target {
            state.eof = true;
        }
        if let Some(waker) = state.waker.take() {
            drop(state);
            waker.wake();
        }
        take
    }

    pub(crate) fn close(&self) {
        let mut state = self.inner.borrow_mut();
        state.eof = true;
        if let Some(waker) = state.waker.take() {
            drop(state);
            waker.wake();
        }
    }

    pub(crate) fn is_complete(&self) -> bool {
        let state = self.inner.borrow();
        state.eof
    }
}

type ResponseSerializer = Box<dyn FnOnce(&mut [u8]) -> PendingOutcome + 'static>;

pub struct PendingResponse {
    pub serializer: ResponseSerializer,
}

pub enum PendingOutcome {
    Inline {
        written: u16,
        close_after: bool,
    },
    Static {
        hdr_written: u32,
        body: &'static [u8],
        close_after: bool,
    },
    Stream {
        hdr_written: u32,
        stream: Pin<Box<dyn ResponseChunkStream>>,
        close_after: bool,
    },
}

pub(crate) struct PendingFuture {
    pub(crate) generation: SlotGen,
    pub(crate) future: dambi::InlineFuture<'static, PendingResponse, 64>,
}

pub(crate) struct LongFuture {
    pub(crate) generation: SlotGen,
    pub(crate) future: Pin<Box<dyn std::future::Future<Output = ()> + 'static>>,
}