leyline-http 0.1.5

use tokio::sync::mpsc::error::TrySendError;

use super::*;

enum Push {
    Sent,
    Stalled,
    Closed,
}

enum Pull {
    Drained,
    Stalled,
    Closed,
}

pub(in crate::quic::pool) fn forward_stream_body(
    h3: &mut quiche::h3::Connection,
    conn: &mut quiche::Connection,
    stream_id: u64,
    stream: &mut H3Stream,
    scratch: &mut [u8],
) -> bool {
    let Some(tx) = stream.stream_tx.clone() else {
        return false;
    };

    if let Some(chunk) = stream.stalled.take() {
        match push_chunk(h3, conn, stream_id, stream, &tx, chunk) {
            Push::Sent => {}
            Push::Stalled => return false,
            Push::Closed => return true,
        }
    }

    match pull_body(h3, conn, stream_id, stream, scratch, &tx) {
        Pull::Drained => finish_stream(conn, stream_id, stream),
        Pull::Stalled => false,
        Pull::Closed => true,
    }
}

fn push_chunk(
    h3: &mut quiche::h3::Connection,
    conn: &mut quiche::Connection,
    stream_id: u64,
    stream: &mut H3Stream,
    tx: &mpsc::Sender<std::io::Result<Bytes>>,
    chunk: Bytes,
) -> Push {
    match tx.try_send(Ok(chunk)) {
        Ok(()) => Push::Sent,
        Err(TrySendError::Full(item)) => {
            if let Ok(returned) = item {
                stream.stalled = Some(returned);
            }
            Push::Stalled
        }
        Err(TrySendError::Closed(_)) => {
            abort_stream(
                h3,
                conn,
                stream_id,
                stream,
                quiche::h3::WireErrorCode::RequestCancelled,
            );
            Push::Closed
        }
    }
}

fn pull_body(
    h3: &mut quiche::h3::Connection,
    conn: &mut quiche::Connection,
    stream_id: u64,
    stream: &mut H3Stream,
    scratch: &mut [u8],
    tx: &mpsc::Sender<std::io::Result<Bytes>>,
) -> Pull {
    loop {
        match h3.recv_body(conn, stream_id, scratch) {
            Ok(0) | Err(quiche::h3::Error::Done) => return Pull::Drained,
            Ok(n) => {
                stream.body_bytes_seen += n;
                let chunk = Bytes::copy_from_slice(&scratch[..n]);
                match push_chunk(h3, conn, stream_id, stream, tx, chunk) {
                    Push::Sent => {}
                    Push::Stalled => return Pull::Stalled,
                    Push::Closed => return Pull::Closed,
                }
            }
            Err(e) => {
                let message = on_body_read_error(h3, conn, stream_id, stream, e);
                stream.deliver_terminal(std::io::Error::other(message));
                return Pull::Closed;
            }
        }
    }
}

fn finish_stream(conn: &mut quiche::Connection, stream_id: u64, stream: &mut H3Stream) -> bool {
    let finished = stream.stalled.is_none()
        && (stream.peer_finished
            || (conn.stream_finished(stream_id) && !conn.stream_readable(stream_id)));
    if !finished {
        return false;
    }
    reset_upload_half(
        conn,
        stream_id,
        stream,
        quiche::h3::WireErrorCode::RequestCancelled,
    );
    if stream.length_mismatch() {
        stream.deliver_terminal(std::io::Error::other(LENGTH_MISMATCH));
    }
    stream.stream_tx = None;
    stream.terminal = None;
    true
}

pub(in crate::quic::pool) fn pump_streaming_bodies(
    h3: &mut quiche::h3::Connection,
    conn: &mut quiche::Connection,
    streams: &mut HashMap<u64, H3Stream>,
    scratch: &mut [u8],
) -> bool {
    let mut to_remove: Vec<u64> = Vec::new();
    for (stream_id, stream) in streams.iter_mut() {
        if stream.stream_tx.is_none() {
            continue;
        }
        if forward_stream_body(h3, conn, *stream_id, stream, scratch) {
            to_remove.push(*stream_id);
        }
    }

    for id in to_remove {
        streams.remove(&id);
    }

    streams.values().any(|s| s.stalled.is_some())
}