wtx 0.50.0

use crate::{
  _AFTER_CLOSE_TIMEOUT_MS,
  collections::{ArrayVectorU8, Vector},
  futures::{Sleep, TryJoinArrayVector},
  http::{Headers, StatusCode, U31},
  http2::{
    Http2Error, Http2Inner, Http2RecvStatus, Http2SendStatus,
    hpack_static_headers::{HpackStaticRequestHeaders, HpackStaticResponseHeaders},
    misc::{frame_reader_rslt, protocol_err, sorp_mut, status_recv, status_send, write_array},
    window::WindowsPair,
    write_functions::{encode_headers, push_data, push_headers, push_trailers, write_frames},
  },
  misc::{Usize, span::Span},
  stream::StreamWriter,
  sync::Arc,
  tls::TlsMode,
};
use core::{future::poll_fn, mem, pin::pin, task::Poll, time::Duration};

/// Groups common client and server operations as well as low level methods that deal with
/// individual frames.
#[derive(Debug)]
pub struct CommonStream<'instance, SW, TM, const IS_CLIENT: bool> {
  pub(crate) inner: &'instance Arc<Http2Inner<SW, TM, IS_CLIENT>>,
  pub(crate) linger: bool,
  pub(crate) span: &'instance Span,
  pub(crate) stream_id: U31,
}

impl<SW, TM, const IS_CLIENT: bool> CommonStream<'_, SW, TM, IS_CLIENT>
where
  SW: StreamWriter,
  TM: TlsMode,
{
  /// Removes internal elements that are no longer necessary after the end of the stream.
  #[inline]
  pub async fn clear(&self) -> crate::Result<()> {
    let Self { inner, linger, span: _, stream_id } = self;
    if *linger {
      Sleep::new(Duration::from_millis(_AFTER_CLOSE_TIMEOUT_MS))?.await?;
    }
    let mut hd_guard = inner.hd.lock().await;
    let hdpm = hd_guard.parts_mut();
    if let Some(elem) = hdpm.hb.scrps.remove(stream_id) {
      elem.waker.wake();
    }
    if let Some(elem) = hdpm.hb.sorps.remove(stream_id) {
      elem.waker.wake();
    }
    Ok(())
  }

  /// Receive Data.
  ///
  /// `cb` contains the bytes received from a ongoing stream and all operations inside of it should
  /// be finished as soon as possible to avoid long contentions. If the remote peer sends no data
  /// or only a single DATA frame, `cb` is never called.
  ///
  /// The remaining payload (if any) is returned in the [`Http2RecvStatus`] result.
  ///
  /// Shouldn't interact with higher operations that receive data.
  #[inline]
  pub async fn recv_data<ONG>(
    &mut self,
    mut cb: impl FnMut(&mut [u8]) -> crate::Result<ONG>,
  ) -> crate::Result<Http2RecvStatus<Vector<u8>, ONG>> {
    let Self { inner, linger: _, span, stream_id } = self;
    let _e = span.enter();
    _trace!("Fetching data");
    let mut hd_guard_pin = pin!(inner.hd.lock());
    poll_fn(|cx| {
      let mut hd_guard = lock_pin!(cx, inner.hd, hd_guard_pin);
      let hdpm = hd_guard.parts_mut();
      let sorp = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
      if let Some(elem) = status_recv(&inner.is_conn_open, sorp, |local_sorp| {
        Ok(mem::take(&mut local_sorp.msg_buffer.body))
      })? {
        return Poll::Ready(Ok(elem));
      }
      if sorp.has_one_or_more_data_frames && !sorp.msg_buffer.body.is_empty() {
        frame_reader_rslt(hdpm.frame_reader_error)?;
        let ongoing = cb(&mut sorp.msg_buffer.body)?;
        sorp.msg_buffer.body.clear();
        Poll::Ready(Ok(Http2RecvStatus::Ongoing(ongoing)))
      } else {
        sorp.waker.clone_from(cx.waker());
        Poll::Pending
      }
    })
    .await
  }

  /// Receive Trailers
  ///
  /// Low level operation that retrieves one or more frames that compose a header. Shouldn't interact
  /// with higher operations that receive data.
  #[inline]
  pub async fn recv_trailers(&mut self) -> crate::Result<Http2RecvStatus<Headers, ()>> {
    let Self { inner, linger: _, span, stream_id } = self;
    let _e = span.enter();
    _trace!("Fetching trailers");
    let mut hd_guard_pin = pin!(inner.hd.lock());
    poll_fn(|cx| {
      let mut hd_guard = lock_pin!(cx, inner.hd, hd_guard_pin);
      let hdpm = hd_guard.parts_mut();
      let sorp = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
      if let Some(elem) = status_recv(&inner.is_conn_open, sorp, |local_sorp| {
        Ok(mem::take(&mut local_sorp.msg_buffer.headers))
      })? {
        return Poll::Ready(Ok(elem));
      }
      sorp.waker.clone_from(cx.waker());
      frame_reader_rslt(hdpm.frame_reader_error)?;
      Poll::Pending
    })
    .await
  }

  /// Should be used when receiving data to re-evaluated flow control values. Both connection and
  /// stream capacities are modified.
  ///
  /// `value` is capped to an integer of 31 bits.
  #[inline]
  pub async fn release_capacity(&mut self, value: u32) -> crate::Result<()> {
    let Self { inner, linger: _, span: _, stream_id } = self;
    let frame = {
      let mut hd_guard = inner.hd.lock().await;
      let hdpm = hd_guard.parts_mut();
      let elem = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
      let mut wp = WindowsPair::new(hdpm.windows, &mut elem.windows);
      wp.withdrawn_recv(hdpm.hp, *stream_id, U31::from_u32(value))?
    };
    write_array([&frame], &mut *inner.wd.lock().await).await?;
    Ok(())
  }

  /// Should be used when sending data to re-evaluated flow control values. Both connection and
  /// stream capacities are modified.
  ///
  /// `value` is capped to an integer of 31 bits.
  #[inline]
  pub async fn reserve_capacity(&mut self, value: u32) -> crate::Result<()> {
    let Self { inner, linger: _, span: _, stream_id } = self;
    let mut lock = inner.hd.lock().await;
    let hdpm = lock.parts_mut();
    let elem = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
    let mut wp = WindowsPair::new(hdpm.windows, &mut elem.windows);
    wp.withdrawn_send(Some(*stream_id), U31::from_u32(value))
  }

  /// Low level operation that sends the content of `data` as one or more DATA frames. If `eos` is
  /// true, then the last frame is set with the end-of-stream flag. Shouldn't interact with
  /// higher operations that send data.
  ///
  /// This method will spin until the entirety of `data` is sent and such behavior depends on the
  /// current available window size as well as the negotiated maximum frame length.
  #[inline]
  pub async fn send_data(&self, data: &[u8], is_eos: bool) -> crate::Result<Http2SendStatus> {
    let Self { inner, linger: _, span, stream_id } = self;
    let _e = span.enter();
    _trace!("Sending data");
    let mut data_idx = 0;
    let mut frames = ArrayVectorU8::new();
    loop {
      let opt = {
        frames.clear();
        let mut hd_pin = pin!(inner.hd.lock());
        poll_fn(|cx| {
          let mut hd_guard = lock_pin!(cx, inner.hd, hd_pin);
          let hdpm = hd_guard.parts_mut();
          let sorp = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
          if let Some(elem) = status_send::<false>(&inner.is_conn_open, sorp) {
            return Poll::Ready(crate::Result::Ok(Some(elem)));
          }
          let mut wp = WindowsPair::new(hdpm.windows, &mut sorp.windows);
          let Ok(available_send @ 1..=u32::MAX) = u32::try_from(wp.available_send()) else {
            sorp.waker.clone_from(cx.waker());
            return Poll::Pending;
          };
          let _ = push_data(
            available_send,
            data,
            &mut data_idx,
            &mut frames,
            is_eos,
            hdpm.hps.max_frame_len,
            *stream_id,
            &mut wp,
          )?;
          Poll::Ready(Ok(None))
        })
        .await?
      };
      if let Some(el) = opt {
        return Ok(el);
      }
      write_frames((&[], data), &frames, &mut *inner.wd.lock().await).await?;
      if *Usize::from(data_idx) >= data.len() {
        return Ok(Http2SendStatus::Ok);
      }
    }
  }

  /// Calls [`Self::send_data`] concurrently on each data element terminating the stream in the
  /// final iteration.
  ///
  /// Returns [`None`] if `data` is empty. This method will abort early if any sending result is
  /// not [`Http2SendStatus::Ok`].
  #[inline]
  pub async fn send_data_concurrent<const N: usize>(
    &self,
    mut data: ArrayVectorU8<&[u8], N>,
  ) -> crate::Result<Option<Http2SendStatus>> {
    let Some(last) = data.pop() else {
      return Ok(None);
    };
    let mut futures = ArrayVectorU8::<_, N>::new();
    for elem in data {
      futures.push(self.send_data(elem, false))?;
    }
    drop(
      TryJoinArrayVector::new(futures, |hss| {
        if !matches!(hss, Http2SendStatus::Ok) {
          return Err(protocol_err(Http2Error::ClosedConnectionWhenSendingConcurrentData));
        }
        Ok(hss)
      })
      .await?,
    );
    Ok(Some(self.send_data(last, true).await?))
  }

  /// Calls [`Self::send_data`] sequentially on each data element terminating the stream in the
  /// final iteration.
  ///
  /// Returns [`None`] if `data` is empty. This method will abort early if any sending result is
  /// not [`Http2SendStatus::Ok`].
  #[inline]
  pub async fn send_data_sequential<'bytes, I>(
    &self,
    data: I,
  ) -> crate::Result<Option<Http2SendStatus>>
  where
    I: IntoIterator<Item = &'bytes [u8]>,
    I::IntoIter: ExactSizeIterator,
  {
    let mut iter = data.into_iter();
    let take = iter.len().saturating_sub(1);
    for elem in iter.by_ref().take(take) {
      let hss = self.send_data(elem, false).await?;
      if !matches!(hss, Http2SendStatus::Ok) {
        return Ok(Some(hss));
      }
    }
    if let Some(elem) = iter.next() {
      return Ok(Some(self.send_data(elem, true).await?));
    }
    Ok(None)
  }

  send_go_away_method!();

  /// Low level operation that sends the content of `headers` with at most two frames. If `is_eos`
  /// is true, then the last frame is set with the end-of-stream flag. Shouldn't interact with
  /// higher operations that send data.
  ///
  /// If two frames aren't enough for the contents of `headers`, try increasing the maximum frame
  /// length.
  #[inline]
  pub async fn send_headers(
    &mut self,
    enc_buffer: &mut Vector<u8>,
    headers: &Headers,
    is_eos: bool,
    status_code: StatusCode,
  ) -> crate::Result<Http2SendStatus> {
    let Self { inner, linger: _, span, stream_id } = self;
    let _e = span.enter();
    _trace!("Sending headers");
    let hsreh = HpackStaticResponseHeaders { status_code: Some(status_code) };
    let max_frame_len = {
      let mut hd_guard = inner.hd.lock().await;
      let hdpm = hd_guard.parts_mut();
      let sorp = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
      if let Some(elem) = status_send::<false>(&inner.is_conn_open, sorp) {
        return Ok(elem);
      }
      encode_headers::<false>(
        enc_buffer,
        headers,
        &mut hdpm.hb.hpack_enc,
        (HpackStaticRequestHeaders::EMPTY, hsreh),
      )?;
      hdpm.hp.max_frame_len()
    };
    let mut frames = ArrayVectorU8::new();
    let _ = push_headers::<IS_CLIENT>(
      enc_buffer,
      &mut frames,
      &mut 0,
      (HpackStaticRequestHeaders::EMPTY, hsreh),
      is_eos,
      max_frame_len,
      *stream_id,
    )?;
    write_frames((enc_buffer, &[]), &frames, &mut *inner.wd.lock().await).await?;
    Ok(Http2SendStatus::Ok)
  }

  /// Sends a reset frame to the peer, which cancels this stream.
  #[inline]
  pub async fn send_reset(&self, error_code: crate::http2::Http2ErrorCode) {
    let Self { inner, linger: _, span: _, stream_id } = self;
    let _ = crate::http2::misc::send_reset_stream(error_code, inner, *stream_id).await;
  }

  /// Low level operation that sends headers that are preceded by DATA frames and then closes
  /// the stream. Shouldn't interact with higher operations that send data.
  ///
  /// An error will probably be returned if the end-of-stream flag was set in previous operations.
  ///
  /// Returns `false` if the stream is already closed.
  #[inline]
  pub async fn send_trailers(
    &mut self,
    enc_buffer: &mut Vector<u8>,
    trailers: &Headers,
  ) -> crate::Result<Http2SendStatus> {
    let Self { inner, linger: _, span, stream_id } = self;
    let _e = span.enter();
    _trace!("Sending {} trailers", trailers.headers_len());
    let mut frames = ArrayVectorU8::new();
    {
      let mut hd_guard = inner.hd.lock().await;
      let hdpm = hd_guard.parts_mut();
      let sorp = sorp_mut(&mut hdpm.hb.sorps, *stream_id)?;
      if let Some(elem) = status_send::<false>(&inner.is_conn_open, sorp) {
        return Ok(elem);
      }
      let _ = push_trailers(
        enc_buffer,
        &mut frames,
        trailers,
        &mut hdpm.hb.hpack_enc,
        &mut 0,
        hdpm.hps.max_frame_len,
        *stream_id,
      )?;
    }
    write_frames((enc_buffer, &[]), &frames, &mut *inner.wd.lock().await).await?;
    Ok(Http2SendStatus::Ok)
  }

  /// Stream ID
  #[inline]
  pub const fn stream_id(&self) -> u32 {
    self.stream_id.u32()
  }

  /// Low level operation that returns the current flow control parameters of the stream.
  #[inline]
  pub async fn windows(&self) -> crate::Result<crate::http2::Windows> {
    let Self { inner, linger: _, span: _, stream_id } = self;
    Ok(sorp_mut(&mut inner.hd.lock().await.parts_mut().hb.sorps, *stream_id)?.windows)
  }
}