wtx 0.50.0

use crate::{
  _AFTER_CLOSE_TIMEOUT_MS,
  collections::{MaybeUninitSlice, ShortBoxSliceU16},
  futures::Sleep,
  misc::{ConnectionState, Either},
  stream::{BufStreamReader, StreamCommon, StreamReader},
  sync::{Arc, AtomicU8, AtomicWaker},
  tls::{
    TlsMode, TlsStreamBridge, TlsStreamBridgeData,
    key_schedule::KeyScheduleRead,
    misc::read_after_handshake_data,
    protocol::{alert::Alert, key_update::KeyUpdate, new_session_ticket::NewSessionTicket},
  },
};
use alloc::boxed::Box;
use core::{
  future::poll_fn,
  hint::cold_path,
  num::NonZeroUsize,
  pin::{Pin, pin},
  sync::atomic::Ordering,
  task::{Poll, ready},
  time::Duration,
};

/// Reader that can be used in concurrent scenarios.
#[derive(Debug)]
pub struct TlsStreamReader<SR, TM, const IS_CLIENT: bool> {
  connection_state: Arc<AtomicU8>,
  ksr: KeyScheduleRead,
  max_fragment_length: u16,
  new_session_ticket: Option<NewSessionTicket<ShortBoxSliceU16<u8>>>,
  plaintext_consumed: usize,
  plaintext_len: usize,
  reader_buffer: BufStreamReader,
  reader_waker: Arc<AtomicWaker>,
  stream_bridge: TlsStreamBridge<IS_CLIENT>,
  stream_reader: SR,
  timer: Pin<Box<Sleep>>,
  _tm: TM,
}

impl<SR, TM, const IS_CLIENT: bool> TlsStreamReader<SR, TM, IS_CLIENT> {
  #[inline]
  pub(crate) fn new(
    connection_state: Arc<AtomicU8>,
    ksr: KeyScheduleRead,
    max_fragment_length: u16,
    new_session_ticket: Option<NewSessionTicket<ShortBoxSliceU16<u8>>>,
    plaintext_consumed: usize,
    plaintext_len: usize,
    reader_buffer: BufStreamReader,
    reader_waker: Arc<AtomicWaker>,
    stream_bridge: TlsStreamBridge<IS_CLIENT>,
    stream_reader: SR,
    _tm: TM,
  ) -> crate::Result<Self> {
    Ok(Self {
      connection_state,
      ksr,
      max_fragment_length,
      new_session_ticket,
      plaintext_consumed,
      plaintext_len,
      reader_buffer,
      reader_waker,
      stream_bridge,
      stream_reader,
      timer: Box::pin(Sleep::new(Duration::from_millis(_AFTER_CLOSE_TIMEOUT_MS))?),
      _tm,
    })
  }

  /// Closes itself as well as the write part without a graceful shutdown.
  //
  // There is nothing to write to the peer so we don't wake the writer side.
  #[inline]
  pub fn close_abruptly(&self) {
    self.connection_state.store(ConnectionState::ClosedAbruptly.into(), Ordering::Relaxed);
  }

  /// See [`ConnectionState`].
  #[inline]
  pub fn connection_state(&self) -> ConnectionState {
    self.connection_state.load(Ordering::Relaxed).into()
  }

  /// Sends a warning alert of type `CloseNotify`, gracefully closing the connection.
  #[inline]
  pub fn send_close_notify(&self) -> crate::Result<()> {
    self.connection_state.store(ConnectionState::WriteClosed.into(), Ordering::Relaxed);
    self.stream_bridge.update(TlsStreamBridgeData::new(Either::Left(Alert::close_notify())));
    Ok(())
  }

  #[cfg(any(feature = "http2", feature = "web-socket"))]
  #[inline]
  pub(crate) const fn connection_state_raw(&self) -> &Arc<AtomicU8> {
    &self.connection_state
  }
}

impl<SR, TM, const IS_CLIENT: bool> StreamCommon for TlsStreamReader<SR, TM, IS_CLIENT> {}

impl<SR, TM, const IS_CLIENT: bool> StreamReader for TlsStreamReader<SR, TM, IS_CLIENT>
where
  SR: StreamReader,
  TM: TlsMode,
{
  #[inline]
  async fn read(&mut self, bytes: MaybeUninitSlice<'_, u8>) -> crate::Result<Option<NonZeroUsize>> {
    let Self {
      connection_state,
      ksr,
      max_fragment_length,
      new_session_ticket,
      plaintext_consumed,
      plaintext_len,
      reader_buffer,
      reader_waker,
      stream_bridge,
      stream_reader,
      timer,
      _tm,
    } = self;
    let mut read_fut = pin!(async {
      if TM::TY.is_plain_text() {
        return stream_reader.read(bytes).await;
      }
      read_after_handshake_data::<_, _, IS_CLIENT>(
        (&*stream_bridge, &*connection_state),
        bytes,
        ksr,
        *max_fragment_length,
        new_session_ticket,
        plaintext_consumed,
        plaintext_len,
        reader_buffer,
        stream_reader,
        alert_cb,
        closed_conn_cb,
        key_update_cb,
      )
      .await
    });
    poll_fn(|cx| match read_fut.as_mut().poll(cx) {
      Poll::Ready(res) => Poll::Ready(res),
      Poll::Pending => {
        reader_waker.register(cx.waker());
        let current_state = connection_state.load(Ordering::Relaxed);
        match ConnectionState::from(current_state) {
          // Normal operation
          ConnectionState::Draining | ConnectionState::Open => Poll::Pending,
          // * An abrupt close signal was received/generated by us or the writer side abruptly
          // closed the connection.
          // * A graceful close signal was received (We CANNOT read close ourselves). The writer
          // side should have received a closing signal
          ConnectionState::ClosedAbruptly
          | ConnectionState::ClosedGracefully
          | ConnectionState::ReadClosed => {
            cold_path();
            Poll::Ready(Ok(None))
          }
          // After user interaction the writer side set itself as write closed then woke us. This
          // is a graceful stop.
          ConnectionState::WriteClosed => {
            cold_path();
            let _rslt = ready!(timer.as_mut().poll(cx));
            connection_state.store(ConnectionState::ClosedGracefully.into(), Ordering::Relaxed);
            Poll::Ready(Ok(None))
          }
        }
      }
    })
    .await
  }
}

async fn alert_cb<SR, const IS_CLIENT: bool>(
  aux: &mut (&TlsStreamBridge<IS_CLIENT>, &Arc<AtomicU8>),
  alert: Alert,
  _: &mut SR,
) -> crate::Result<bool> {
  if alert.is_close_notify() {
    aux.1.store(ConnectionState::ReadClosed.into(), Ordering::Relaxed);
    aux.0.update(TlsStreamBridgeData::new(Either::Left(alert)));
  } else {
    aux.1.store(ConnectionState::ClosedAbruptly.into(), Ordering::Relaxed);
  }
  Ok(true)
}

// This branch is only entered when the peer closed the connection without an alert.
fn closed_conn_cb<const IS_CLIENT: bool>(aux: &mut (&TlsStreamBridge<IS_CLIENT>, &Arc<AtomicU8>)) {
  aux.1.store(ConnectionState::ClosedAbruptly.into(), Ordering::Relaxed);
}

async fn key_update_cb<SR, const IS_CLIENT: bool>(
  aux: &mut (&TlsStreamBridge<IS_CLIENT>, &Arc<AtomicU8>),
  key_update: Option<KeyUpdate>,
  _: &mut SR,
) -> crate::Result<()> {
  if let Some(elem) = key_update {
    aux.0.update(TlsStreamBridgeData::new(Either::Right(elem)));
  }
  Ok(())
}