dope-tls 0.18.0

use std::iter::from_fn;

use dope_net::wire::buffered::{Buffer, Scratch};
use dope_net::wire::send::{Plain, Prepared, SendStorage, Sent, Storage, Vectored};

pub(crate) struct Cursor<'a, I>
where
    I: Iterator<Item = &'a [u8]>,
{
    source: I,
    current: &'a [u8],
    consumed: usize,
}

impl<'a, I> Cursor<'a, I>
where
    I: Iterator<Item = &'a [u8]>,
{
    pub(crate) fn new(source: I) -> Self {
        Self {
            source,
            current: &[],
            consumed: 0,
        }
    }

    pub(crate) fn take(&mut self, len: usize) -> impl Iterator<Item = &'a [u8]> + '_ {
        let mut remaining = len;
        from_fn(move || {
            if remaining == 0 {
                return None;
            }
            let part = self.next_part(remaining)?;
            remaining -= part.len();
            Some(part)
        })
    }

    pub(crate) fn consumed(&self) -> usize {
        self.consumed
    }

    fn next_part(&mut self, max: usize) -> Option<&'a [u8]> {
        while self.current.is_empty() {
            self.current = self.source.next()?;
        }
        let take = self.current.len().min(max);
        let (part, rest) = self.current.split_at(take);
        self.current = rest;
        self.consumed += take;
        Some(part)
    }
}

pub(crate) trait SendProtocol {
    type Storage: SendBuffer;

    fn needs_buffer(&self) -> bool;
    fn encrypt(&mut self, egress: &mut Buffer<Scratch>, plain: &[u8]) -> usize;
    fn encrypt_vectored(&mut self, egress: &mut Buffer<Scratch>, vectored: &Vectored<'_>) -> usize;
    fn propagate_close(&mut self, egress: &mut Buffer<Scratch>) -> bool;
    fn drain_to_egress(&mut self, egress: &mut Buffer<Scratch>);
    fn send_inflight(&mut self) -> &mut bool;
}

pub(crate) trait SendBuffer: SendStorage {
    fn buffer_mut(&mut self) -> Option<&mut Buffer<Scratch>>;
    fn try_buffer(&mut self) -> Option<&mut Buffer<Scratch>>;
    fn release_if_empty(&mut self) {}
}

pub(crate) struct Sender<'a, P> {
    protocol: &'a mut P,
}

impl<'a, P: SendProtocol> Sender<'a, P> {
    pub(crate) fn new(protocol: &'a mut P) -> Self {
        Self { protocol }
    }

    pub(crate) fn prepare<'b>(
        &mut self,
        mut send: Storage<'b, P::Storage>,
        plain: Plain<'b>,
    ) -> Prepared<'b> {
        if plain.is_empty() && !self.protocol.needs_buffer() {
            return send.empty(0);
        }
        let Some(egress) = send.try_buffer() else {
            return send.empty(0);
        };
        let consumed = self.protocol.encrypt(egress, plain.as_slice());
        let close = self.protocol.propagate_close(egress);
        self.finish(send, consumed, close)
    }

    pub(crate) fn prepare_vectored<'b>(
        &mut self,
        mut send: Storage<'b, P::Storage>,
        vectored: Vectored<'b>,
    ) -> Prepared<'b> {
        if vectored.is_empty() && !self.protocol.needs_buffer() {
            return send.empty(0);
        }
        let Some(egress) = send.try_buffer() else {
            return send.empty(0);
        };
        let consumed = self.protocol.encrypt_vectored(egress, &vectored);
        let close = self.protocol.propagate_close(egress);
        self.finish(send, consumed, close)
    }

    pub(crate) fn after_send<'b>(
        &mut self,
        mut send: Storage<'b, P::Storage>,
        sent: Sent,
    ) -> Prepared<'b> {
        *self.protocol.send_inflight() = false;
        let Some(egress) = send.buffer_mut() else {
            return send.empty(0);
        };
        let Ok(prefix) = egress.try_consume_prefix(sent.get()) else {
            return send.empty(0).close_after();
        };
        prefix.commit();
        self.protocol.drain_to_egress(egress);
        let close = self.protocol.propagate_close(egress);
        self.finish(send, 0, close)
    }

    pub(crate) fn flush<'b>(&mut self, mut send: Storage<'b, P::Storage>) -> Prepared<'b> {
        if !self.protocol.needs_buffer() {
            return send.empty(0);
        }
        let Some(egress) = send.try_buffer() else {
            return send.empty(0);
        };
        self.protocol.encrypt(egress, &[]);
        let close = self.protocol.propagate_close(egress);
        self.finish(send, 0, close)
    }

    pub(crate) fn finish<'b>(
        &mut self,
        mut send: Storage<'b, P::Storage>,
        consumed: usize,
        close: bool,
    ) -> Prepared<'b> {
        send.release_if_empty();
        *self.protocol.send_inflight() = !send.as_slice().is_empty();
        let prepared = send.buffered(consumed);
        if close {
            prepared.close_after()
        } else {
            prepared
        }
    }
}