pdk_classy/hl/

request.rs

// Copyright (c) 2025, Salesforce, Inc.,
// All rights reserved.
// For full license text, see the LICENSE.txt file

use std::cell::RefCell;
use std::future::{ready, Ready};
use std::rc::Rc;

use crate::event::{EventKind, RequestBody, RequestHeaders};
use crate::BoxFuture;

use super::body::BodyExchange;
use super::body_stream::{BodyStream, BodyStreamExchange};
use super::dynamic_exchange::DynamicExchange;
use super::entity::{BodyStreamState, IntoBodyState, IntoBodyStreamState};
use super::headers::HeadersExchange;

#[cfg(feature = "experimental_enable_stop_iteration")]
use super::{
    entity::{HeadersBodyState, IntoHeadersBodyState},
    headers_body::HeadersBodyExchange,
};

#[cfg(feature = "experimental")]
use super::{BodyError, Chunk};
use super::{BodyState, EntityState, HeadersState};

/// Initial state of the Request state-machine.
/// This type is intended to be injected as a parameter in Request filters.
pub struct RequestState {
    contains_body: bool,
    exchange: Rc<RefCell<DynamicExchange>>,
}

/// Error that represents an invalid state in the Request state-machine.
#[derive(thiserror::Error, Debug)]
pub enum InvalidRequestState {
    /// Unexpected low-level event.
    #[error("Invalid request state: {0:?}")]
    Event(EventKind),

    /// Unexpected early response.
    #[error("Early response")]
    EarlyResponse,
}

impl RequestState {
    #[allow(clippy::await_holding_refcell_ref)]
    pub(crate) async fn new(
        exchange: Rc<RefCell<DynamicExchange>>,
    ) -> Result<Self, InvalidRequestState> {
        let contains_body;
        {
            let mut exchange_ref = exchange.borrow_mut();
            let Some(exchange) = exchange_ref.wait_for_event::<RequestHeaders>().await else {
                let error = exchange
                    .borrow()
                    .current_event()
                    .map(InvalidRequestState::Event)
                    .unwrap_or(InvalidRequestState::EarlyResponse);
                return Err(error);
            };
            contains_body = !exchange
                .event_data()
                .expect("Must contain headers")
                .event
                .end_of_stream;
        }

        Ok(Self {
            contains_body,
            exchange,
        })
    }
}

impl IntoBodyState for RequestState {
    type BodyState = RequestBodyState;
    type BodyStateFuture = BoxFuture<'static, Self::BodyState>;

    fn into_body_state(self) -> Self::BodyStateFuture {
        Box::pin(async {
            RequestBodyState {
                inner: BodyExchange::new(self.exchange, self.contains_body).await,
            }
        })
    }
}

impl IntoBodyStreamState for RequestState {
    type BodyStreamState = RequestBodyStreamState;

    type BodyStreamStateFuture = BoxFuture<'static, Self::BodyStreamState>;

    fn into_body_stream_state(self) -> Self::BodyStreamStateFuture {
        Box::pin(async {
            RequestBodyStreamState {
                inner: BodyStreamExchange::new(self.exchange, self.contains_body).await,
            }
        })
    }
}

#[cfg(feature = "experimental_enable_stop_iteration")]
impl IntoHeadersBodyState for RequestState {
    type HeadersBodyState = RequestHeadersBodyState;
    type HeadersBodyStateFuture = BoxFuture<'static, Self::HeadersBodyState>;

    fn into_headers_body_state(self) -> Self::HeadersBodyStateFuture {
        Box::pin(async {
            RequestHeadersBodyState {
                inner: HeadersBodyExchange::new(self.exchange, self.contains_body).await,
            }
        })
    }
}

impl EntityState for RequestState {
    type HeadersState = RequestHeadersState;
    type HeadersStateFuture = Ready<Self::HeadersState>;

    fn into_headers_state(self) -> Self::HeadersStateFuture {
        ready(RequestHeadersState {
            inner: HeadersExchange::new(self.exchange, self.contains_body),
        })
    }
}

/// Headers state of the Request state-machine.
pub struct RequestHeadersState {
    inner: HeadersExchange<RequestHeaders, RequestBody>,
}

impl RequestHeadersState {
    /// Returns the HTTP method of the current Request.
    pub fn method(&self) -> String {
        self.inner.with_event_data(|e| e.method())
    }

    /// Returns the HTTP scheme of the current Request.
    pub fn scheme(&self) -> String {
        self.inner.with_event_data(|e| e.scheme())
    }

    /// Returns the HTTP authority of the current Request.
    pub fn authority(&self) -> String {
        self.inner.with_event_data(|e| e.authority())
    }

    /// Returns the HTTP path of the current Request.
    pub fn path(&self) -> String {
        self.inner.with_event_data(|e| e.path())
    }
}

impl IntoBodyState for RequestHeadersState {
    type BodyState = RequestBodyState;

    type BodyStateFuture = BoxFuture<'static, Self::BodyState>;

    fn into_body_state(self) -> Self::BodyStateFuture {
        Box::pin(async {
            RequestBodyState {
                inner: self.inner.into_body_state().await,
            }
        })
    }
}

impl IntoBodyStreamState for RequestHeadersState {
    type BodyStreamState = RequestBodyStreamState;

    type BodyStreamStateFuture = BoxFuture<'static, Self::BodyStreamState>;

    fn into_body_stream_state(self) -> Self::BodyStreamStateFuture {
        Box::pin(async {
            RequestBodyStreamState {
                inner: self.inner.into_body_stream_state().await,
            }
        })
    }
}

#[cfg(feature = "experimental_enable_stop_iteration")]
impl IntoHeadersBodyState for RequestHeadersState {
    type HeadersBodyState = RequestHeadersBodyState;
    type HeadersBodyStateFuture = BoxFuture<'static, Self::HeadersBodyState>;

    fn into_headers_body_state(self) -> Self::HeadersBodyStateFuture {
        Box::pin(async {
            RequestHeadersBodyState {
                inner: self.inner.into_headers_body_state().await,
            }
        })
    }
}

impl HeadersState for RequestHeadersState {
    fn handler(&self) -> &dyn super::HeadersHandler {
        self.inner.handler()
    }

    fn contains_body(&self) -> bool {
        self.inner.contains_body
    }
}

/// Body state of the Request state-machine.
pub struct RequestBodyState {
    inner: BodyExchange<RequestBody>,
}

impl BodyState for RequestBodyState {
    fn handler(&self) -> &dyn super::BodyHandler {
        &self.inner
    }

    fn contains_body(&self) -> bool {
        self.inner.contains_body()
    }
}

pub struct RequestBodyStreamState {
    inner: BodyStreamExchange<RequestBody>,
}

impl BodyStreamState for RequestBodyStreamState {
    type Stream<'b>
        = BodyStream<'b>
    where
        Self: 'b;

    fn contains_body(&self) -> bool {
        self.inner.contains_body()
    }

    fn stream(&self) -> Self::Stream<'_> {
        self.inner.stream()
    }

    #[cfg(feature = "experimental")]
    fn write_chunk(&self, chunk: Chunk) -> Result<(), BodyError> {
        self.inner.write_chunk(chunk)
    }
}

#[cfg(feature = "experimental_enable_stop_iteration")]
/// HeadersBody state of the Request state-machine.
pub struct RequestHeadersBodyState {
    inner: HeadersBodyExchange<RequestBody>,
}

#[cfg(feature = "experimental_enable_stop_iteration")]
impl HeadersBodyState for RequestHeadersBodyState {
    fn handler(&self) -> &dyn super::HeadersBodyHandler {
        self.inner.handler()
    }

    fn contains_body(&self) -> bool {
        self.inner.contains_body()
    }
}