dope-tls 0.1.1

use std::io;

use dope_futures::{AsyncRead, AsyncWrite};

use crate::state::{TlsError, TlsState};

const READ_SCRATCH_BYTES: usize = 8192;

pub struct TlsStream<IO> {
    io: IO,
    tls: TlsState,
    residual: Vec<u8>,
}

impl<IO> TlsStream<IO> {
    pub(crate) fn new(io: IO, tls: TlsState) -> Self {
        Self {
            io,
            tls,
            residual: Vec::new(),
        }
    }

    pub fn into_parts(self) -> (IO, TlsState) {
        (self.io, self.tls)
    }

    pub fn tls(&self) -> &TlsState {
        &self.tls
    }

    pub fn selected_alpn(&self) -> Option<&[u8]> {
        self.tls.selected_alpn()
    }
}

impl<IO> TlsStream<IO>
where
    IO: AsyncRead + AsyncWrite,
{
    async fn flush_pending(&mut self) -> Result<(), TlsError> {
        let pending = self.tls.pull_send();
        if !pending.is_empty() {
            let (res, _) = self.io.write(pending).await;
            res?;
        }
        Ok(())
    }

    pub async fn shutdown(&mut self) -> Result<(), TlsError> {
        if self.tls.is_closed() {
            return Ok(());
        }
        self.tls.send_close_notify()?;
        self.flush_pending().await
    }

    pub async fn handshake(&mut self) -> Result<(), TlsError> {
        loop {
            self.flush_pending().await?;

            if self.tls.is_established() {
                self.flush_pending().await?;
                return Ok(());
            }
            if self.tls.is_closed() {
                return Err(TlsError::Other);
            }

            let scratch = vec![0u8; READ_SCRATCH_BYTES];
            let (res, scratch) = self.io.read(scratch).await;
            let n = res?;
            if n == 0 {
                return Err(io::Error::new(
                    io::ErrorKind::UnexpectedEof,
                    "tls handshake eof",
                )
                .into());
            }
            self.tls.read_tcp(&scratch[..n])?;
        }
    }
}

impl<IO> AsyncRead for TlsStream<IO>
where
    IO: AsyncRead + AsyncWrite,
{
    async fn read(&mut self, mut buf: Vec<u8>) -> (io::Result<usize>, Vec<u8>) {
        loop {
            if !self.residual.is_empty() {
                let n = self.residual.len().min(buf.len());
                buf[..n].copy_from_slice(&self.residual[..n]);
                self.residual.drain(0..n);
                return (Ok(n), buf);
            }
            if let Some(data) = self.tls.pull_app() {
                let n = data.len().min(buf.len());
                buf[..n].copy_from_slice(&data[..n]);
                if data.len() > n {
                    self.residual.extend_from_slice(&data[n..]);
                }
                return (Ok(n), buf);
            }
            if self.tls.is_closed() {
                return (Ok(0), buf);
            }

            let scratch = vec![0u8; READ_SCRATCH_BYTES];
            let (res, scratch) = self.io.read(scratch).await;
            let n = match res {
                Ok(0) => return (Ok(0), buf),
                Ok(n) => n,
                Err(e) => return (Err(e), buf),
            };
            if let Err(e) = self.tls.read_tcp(&scratch[..n]) {
                return (Err(e.into()), buf);
            }
            if let Err(e) = self.flush_pending().await {
                return (Err(e.into()), buf);
            }
        }
    }
}

impl<IO> AsyncWrite for TlsStream<IO>
where
    IO: AsyncRead + AsyncWrite,
{
    async fn write(&mut self, buf: Vec<u8>) -> (io::Result<usize>, Vec<u8>) {
        if let Err(e) = self.tls.write_app(&buf) {
            return (Err(e.into()), buf);
        }
        let pending = self.tls.pull_send();
        let n = buf.len();
        let (res, _) = self.io.write(pending).await;
        match res {
            Ok(_) => (Ok(n), buf),
            Err(e) => (Err(e), buf),
        }
    }
}

impl From<TlsError> for io::Error {
    fn from(e: TlsError) -> Self {
        Self::other(format!("tls: {e:?}"))
    }
}

impl From<io::Error> for TlsError {
    fn from(e: io::Error) -> Self {
        Self::Io(e)
    }
}