pdk-classy 1.3.0-alpha.0

// Copyright 2023 Salesforce, Inc. All rights reserved.
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::dynamic_exchange::DynamicExchange;
use super::entity::IntoBodyState;
use super::headers::HeadersExchange;
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 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 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()
    }
}