dope-tls 0.1.1

use std::cell::RefCell;
use std::rc::Rc;

use dope_session::{CodecLayer, WriteBufArr};

use crate::state::TlsState;
use crate::wire_bufs::WireBufs;

thread_local! {
    static CURRENT_CLIENT_CONFIG: RefCell<Option<Rc<shin::client::Config>>> =
        const { RefCell::new(None) };
}

pub struct LazyTlsConnState {
    pub tls: Option<TlsState>,
    pub bufs: WireBufs,
    pub pending_app: Vec<u8>,
}

impl Default for LazyTlsConnState {
    fn default() -> Self {
        Self {
            tls: None,
            bufs: WireBufs::new(),
            pending_app: Vec::new(),
        }
    }
}

#[derive(Debug)]
pub enum UpgradeError {
    AlreadyHandshaking,
    AlreadyEstablished,
    Closed,
    Inner(crate::state::TlsError),
}

pub struct LazyTlsCodec;

impl LazyTlsCodec {
    pub fn install_client_config(cfg: shin::client::Config) {
        CURRENT_CLIENT_CONFIG.with(|slot| {
            *slot.borrow_mut() = Some(Rc::new(cfg));
        });
    }

    pub fn current_client_config() -> Option<Rc<shin::client::Config>> {
        CURRENT_CLIENT_CONFIG.with(|slot| slot.borrow().clone())
    }

    pub fn start_handshake(
        state: &mut LazyTlsConnState,
        config: shin::client::Config,
    ) -> Result<(), UpgradeError> {
        if state.bufs.close {
            return Err(UpgradeError::Closed);
        }
        if let Some(tls) = state.tls.as_ref() {
            return Err(if tls.is_established() {
                UpgradeError::AlreadyEstablished
            } else {
                UpgradeError::AlreadyHandshaking
            });
        }
        let mut tls = TlsState::new_client(config).map_err(UpgradeError::Inner)?;
        tls.pull_send_into(&mut state.bufs.wire);
        state.tls = Some(tls);
        Ok(())
    }
}

impl CodecLayer for LazyTlsCodec {
    type ConnState = LazyTlsConnState;
    type Staging = WriteBufArr;
    const IS_PASSTHROUGH: bool = false;

    #[inline(always)]
    fn is_active(state: &Self::ConnState) -> bool {
        state.tls.as_ref().is_none_or(|s| s.is_established())
    }

    fn process_inbound(state: &mut Self::ConnState, wire_in: &[u8]) {
        match state.tls.as_mut() {
            None => {
                state.bufs.plaintext.extend_from_slice(wire_in);
            }
            Some(tls) => {
                if !tls.try_read_tcp(wire_in) {
                    state.bufs.close = true;
                    return;
                }
                tls.pull_send_into(&mut state.bufs.wire);
                while let Some(chunk) = tls.pull_app() {
                    state.bufs.plaintext.extend_from_slice(&chunk);
                }
                if tls.is_established() && !state.pending_app.is_empty() {
                    let pending = std::mem::take(&mut state.pending_app);
                    if tls.write_app(&pending).is_err() {
                        state.bufs.close = true;
                        return;
                    }
                    tls.pull_send_into(&mut state.bufs.wire);
                }
                if tls.is_closed() {
                    state.bufs.close = true;
                }
            }
        }
    }

    fn process_outbound(state: &mut Self::ConnState, plaintext_in: &[u8]) {
        match state.tls.as_mut() {
            None => {
                state.bufs.wire.extend_from_slice(plaintext_in);
            }
            Some(tls) => {
                if !tls.is_established() {
                    state.pending_app.extend_from_slice(plaintext_in);
                    return;
                }
                if tls.write_app(plaintext_in).is_err() {
                    state.bufs.close = true;
                    return;
                }
                tls.pull_send_into(&mut state.bufs.wire);
            }
        }
    }

    #[inline(always)]
    fn wire_slice(state: &Self::ConnState) -> &[u8] {
        state.bufs.wire_slice()
    }

    #[inline(always)]
    fn take_plaintext(state: &mut Self::ConnState) -> Vec<u8> {
        state.bufs.take_plaintext()
    }

    #[inline(always)]
    fn install_plaintext(state: &mut Self::ConnState, plaintext: Vec<u8>) {
        state.bufs.install_plaintext(plaintext);
    }

    fn take_inflight(state: &mut Self::ConnState) -> Option<(*const u8, u32)> {
        state.bufs.take_inflight()
    }

    fn consume_inflight(state: &mut Self::ConnState, n: usize) -> bool {
        state.bufs.consume_inflight(n)
    }

    #[inline(always)]
    fn has_inflight(state: &Self::ConnState) -> bool {
        state.bufs.has_inflight()
    }

    #[inline(always)]
    fn close_pending(state: &Self::ConnState) -> bool {
        state.bufs.close_pending()
    }
}