pdk_classy/

event.rs

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

//! Abstractions of the proxy wasm underlying events.

use std::{
    cell::{Cell, RefCell},
    convert::TryFrom,
    marker::PhantomData,
    ops::Deref,
    rc::Rc,
    task::{Poll, Waker},
};

use futures::{Stream, StreamExt};

use crate::{
    host::Host,
    reactor::http::{HttpReactor, WakerId},
    types::HttpCid,
};
use crate::{
    http_constants::{
        DEFAULT_PATH, HEADER_AUTHORITY, HEADER_METHOD, HEADER_PATH, HEADER_SCHEME, HEADER_STATUS,
    },
    reactor::http::ExchangePhase,
};

mod private {
    use crate::host::Host;

    pub trait Sealed {}

    pub trait BodyAccessor {
        fn read_body(host: &dyn Host, offset: usize, max_size: usize) -> Option<Vec<u8>>;

        fn write_body(host: &dyn Host, offset: usize, size: usize, value: &[u8]);
    }
}

use private::{BodyAccessor, Sealed};

/// Common interface that all underlying events share.
pub trait Event: Sealed + Clone + Into<FiniteEvent> + TryFrom<FiniteEvent> + Unpin {
    fn kind() -> EventKind;

    fn body_size(&self) -> usize;

    fn should_pause(&self) -> bool;
}

/// Imposes an order of succession between two Events.
pub trait After<S: Event>: Event {}

/// Imposes an order of precedence between two Events.
pub trait Before<S: Event>: Event {}

impl<A, B> Before<B> for A
where
    B: After<A>,
    A: Event,
{
}

/// Events that can access the body and verify if more body remains.
pub trait BodyEvent: Event + BodyAccessor {
    fn end_of_stream(&self) -> bool;
}

/// Events that can access the headers.
pub trait HeadersEvent: Event {}

/// Alias name for CreateContext event
#[derive(Clone, Debug)]
pub struct Start {
    pub(crate) _context_id: HttpCid,
}

#[derive(Clone, Debug)]
/// Represents the state where the request headers are available.
pub struct RequestHeaders {
    pub(crate) _num_headers: usize,
    pub(crate) end_of_stream: bool,
}

impl HeadersEvent for RequestHeaders {}

#[derive(Clone, Debug)]
/// Represents the state where the request body is available.
pub struct RequestBody {
    pub(crate) body_size: usize,
    pub(crate) end_of_stream: bool,
}

impl BodyEvent for RequestBody {
    fn end_of_stream(&self) -> bool {
        self.end_of_stream
    }
}

impl BodyAccessor for RequestBody {
    fn read_body(host: &dyn Host, offset: usize, max_size: usize) -> Option<Vec<u8>> {
        host.get_http_request_body(offset, max_size)
    }

    fn write_body(host: &dyn Host, offset: usize, size: usize, value: &[u8]) {
        host.set_http_request_body(offset, size, value)
    }
}

#[cfg(feature = "experimental_enable_stop_iteration")]
impl HeadersAccessor for EventDataStream<'_, RequestBody> {
    fn header(&self, name: &str) -> Option<String> {
        self.exchange.host.get_http_request_header(name)
    }

    fn headers(&self) -> Vec<(String, String)> {
        self.exchange.host.get_http_request_headers()
    }

    fn add_header(&self, name: &str, value: &str) {
        self.exchange.host.add_http_request_header(name, value);
    }

    fn set_header(&self, name: &str, value: &str) {
        self.exchange
            .host
            .set_http_request_header(name, Some(value));
    }

    fn set_headers(&self, headers: Vec<(&str, &str)>) {
        self.exchange.host.set_http_request_headers(headers);
    }

    fn remove_header(&self, name: &str) {
        self.exchange.host.set_http_request_header(name, None);
    }
}

#[cfg(feature = "experimental_enable_stop_iteration")]
impl HeadersAccessor for EventDataStream<'_, ResponseBody> {
    fn header(&self, name: &str) -> Option<String> {
        self.exchange.host.get_http_response_header(name)
    }

    fn headers(&self) -> Vec<(String, String)> {
        self.exchange.host.get_http_response_headers()
    }

    fn add_header(&self, name: &str, value: &str) {
        self.exchange.host.add_http_response_header(name, value);
    }

    fn set_header(&self, name: &str, value: &str) {
        self.exchange
            .host
            .set_http_response_header(name, Some(value));
    }

    fn set_headers(&self, headers: Vec<(&str, &str)>) {
        self.exchange.host.set_http_response_headers(headers);
    }

    fn remove_header(&self, name: &str) {
        self.exchange.host.set_http_response_header(name, None);
    }
}

#[derive(Clone, Debug)]
/// Represents the state where the request trailers are available.
pub struct RequestTrailers {
    pub(crate) _num_trailers: usize,
}

impl HeadersEvent for RequestTrailers {}

#[derive(Clone, Debug)]
/// Represents the state where the response headers are available.
pub struct ResponseHeaders {
    pub(crate) _num_headers: usize,
    pub(crate) end_of_stream: bool,
}

impl HeadersEvent for ResponseHeaders {}

#[derive(Clone, Debug)]
/// Represents the state where the response body is available.
pub struct ResponseBody {
    pub(crate) body_size: usize,
    pub(crate) end_of_stream: bool,
}

impl BodyEvent for ResponseBody {
    fn end_of_stream(&self) -> bool {
        self.end_of_stream
    }
}

impl BodyAccessor for ResponseBody {
    fn read_body(host: &dyn Host, offset: usize, max_size: usize) -> Option<Vec<u8>> {
        host.get_http_response_body(offset, max_size)
    }

    fn write_body(host: &dyn Host, offset: usize, size: usize, value: &[u8]) {
        host.set_http_response_body(offset, size, value)
    }
}

#[derive(Clone, Debug)]
/// Represents the state where the response trailers are available.
pub struct ResponseTrailers {
    pub(crate) _num_trailers: usize,
}

impl HeadersEvent for ResponseTrailers {}

#[derive(Clone, Debug)]
/// Represents the state where the request has finished processing.
pub struct ExchangeComplete {}

macro_rules! should_pause {
    (RequestBody, $value:expr) => {
        !$value.end_of_stream
    };

    (ResponseBody, $value:expr) => {
        !$value.end_of_stream
    };

    ($event:ty, $value:expr) => {
        false
    };
}

macro_rules! body_size {
    (RequestBody, $value:expr) => {
        $value.body_size
    };

    (ResponseBody, $value:expr) => {
        $value.body_size
    };

    ($event:ty, $value:expr) => {
        0
    };
}

// Implements After trait for an Event
macro_rules! impl_after {

    ($event:ty) => {};

    ($event:ty, $($after:ty),*) => {
        $(impl After<$after> for $event {})*
        impl_after!($($after),*);
    };
}

// Implements After trait for an ordered sequence of events
macro_rules! after {

    ([] $($reversed:tt)*) => {
        impl_after!($($reversed),*); // base case
    };

    ([$first:tt $($rest:tt)*] $($reversed:tt)*) => {
        after!([$($rest)*] $first $($reversed)*);  // recursion
    };
}

macro_rules! finite_events {

    ($($event:ident,)+) => {

        #[derive(Clone, Copy, PartialEq, Eq, Debug)]
        /// Represents the data types of the events that can be processed.
        pub enum EventKind {
            $($event),+
        }

        #[derive(Clone, Debug)]
        /// Types of the events that can be processed.
        pub enum FiniteEvent {
            $($event($event)),+
        }

        impl FiniteEvent {
            pub fn kind(&self) -> EventKind {
                match self {
                    $(Self::$event(_) => EventKind::$event),+
                }
            }
        }

        after!([$($event)*]);

        $(
            impl Sealed for $event {}

            impl Event for $event {
                fn kind() -> EventKind {
                    EventKind::$event
                }

                fn should_pause(&self) -> bool {
                    should_pause!($event, self)
                }

                fn body_size(&self) -> usize {
                    body_size!($event, self)
                }
            }

            impl From<$event> for FiniteEvent {
                fn from(event: $event) -> Self {
                    Self::$event(event)
                }
            }

            impl TryFrom<FiniteEvent> for $event {
                type Error = FiniteEvent;

                fn try_from(finite_event: FiniteEvent) -> Result<Self, FiniteEvent> {
                    match finite_event {
                        FiniteEvent::$event(e) => Ok(e),
                        e => Err(e),
                    }
                }
            }
        )*
    };
}

finite_events! {
    Start,
    RequestHeaders,
    RequestBody,
    RequestTrailers,
    ResponseHeaders,
    ResponseBody,
    ResponseTrailers,
    ExchangeComplete,
}

impl PartialOrd for EventKind {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for EventKind {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.index().cmp(&other.index())
    }
}

impl EventKind {
    fn index(&self) -> usize {
        *self as usize
    }
}

struct InnerExchange<S: Event> {
    reactor: Rc<HttpReactor>,
    host: Rc<dyn Host>,
    first_event: RefCell<Option<S>>,
    empty_stream: Cell<bool>,
    offset: Cell<usize>,
    event_count: Cell<usize>,
}

impl<S: Event> InnerExchange<S> {
    fn take_first_event(&self) -> Option<S> {
        self.first_event.borrow_mut().take()
    }
}

/// A state machine which tracks the current state of the filter.
pub struct Exchange<S: Event = Start> {
    pub(crate) reactor: Rc<HttpReactor>,
    pub(crate) host: Rc<dyn Host>,
    inner: Rc<InnerExchange<S>>,
}

/// Manages the access to the data of each event.
pub struct EventData<'a, S: Event> {
    exchange: Rc<InnerExchange<S>>,
    pub(crate) event: S,
    offset: usize,
    _lifetime: PhantomData<&'a ()>,
}

impl<'a, S: Event> EventData<'a, S> {
    pub(crate) fn new(exchange: &'a Exchange<S>, event: S, offset: usize) -> Self {
        Self {
            exchange: exchange.inner.clone(),
            event,
            offset,
            _lifetime: PhantomData,
        }
    }

    fn from_inner(exchange: Rc<InnerExchange<S>>, event: S, offset: usize) -> Self {
        Self {
            exchange,
            event,
            offset,
            _lifetime: PhantomData,
        }
    }
}

/// Trait that provides access to the headers of an event.
pub trait HeadersAccessor {
    /// Return the first header value for the given name.
    /// Known Limitations: The header value will be converted to an utf-8 String
    /// If the bytes correspond to a non utf-8 string they will be parsed as an iso_8859_1 encoding.
    fn header(&self, name: &str) -> Option<String>;

    /// Returns a copy of all the headers.
    /// Known Limitations: The header values will be converted to utf-8 Strings
    /// If the bytes correspond to a non utf-8 string they will be parsed as an iso_8859_1 encoding.
    fn headers(&self) -> Vec<(String, String)>;

    /// Adds a new value of the header to the event.
    fn add_header(&self, name: &str, value: &str);

    /// Replaces value of the header to the event.
    fn set_header(&self, name: &str, value: &str);

    /// Replaces all the headers of the event.
    fn set_headers(&self, headers: Vec<(&str, &str)>);

    /// Removes the specified header.
    fn remove_header(&self, name: &str);
}

impl EventData<'_, RequestHeaders> {
    pub fn method(&self) -> String {
        self.header(HEADER_METHOD).unwrap_or_default()
    }

    pub fn scheme(&self) -> String {
        self.header(HEADER_SCHEME).unwrap_or_default()
    }

    pub fn authority(&self) -> String {
        self.header(HEADER_AUTHORITY).unwrap_or_default()
    }

    pub fn path(&self) -> String {
        self.header(HEADER_PATH)
            .unwrap_or_else(|| DEFAULT_PATH.to_string())
    }
}

impl EventData<'_, ResponseHeaders> {
    pub fn status_code(&self) -> u32 {
        self.header(HEADER_STATUS)
            .and_then(|status| status.parse::<u32>().ok())
            .unwrap_or_default()
    }
}

impl HeadersAccessor for EventData<'_, RequestHeaders> {
    fn header(&self, name: &str) -> Option<String> {
        self.exchange.host.get_http_request_header(name)
    }

    fn headers(&self) -> Vec<(String, String)> {
        self.exchange.host.get_http_request_headers()
    }

    fn add_header(&self, name: &str, value: &str) {
        self.exchange.host.add_http_request_header(name, value);
    }

    fn set_header(&self, name: &str, value: &str) {
        self.exchange
            .host
            .set_http_request_header(name, Some(value));
    }

    fn set_headers(&self, headers: Vec<(&str, &str)>) {
        self.exchange.host.set_http_request_headers(headers);
    }

    fn remove_header(&self, name: &str) {
        self.exchange.host.set_http_request_header(name, None);
    }
}

impl EventData<'_, RequestTrailers> {
    pub fn header(&self, name: &str) -> Option<String> {
        self.exchange.host.get_http_request_trailer(name)
    }

    pub fn headers(&self) -> Vec<(String, String)> {
        self.exchange.host.get_http_request_trailers()
    }
}

impl HeadersAccessor for EventData<'_, ResponseHeaders> {
    fn header(&self, name: &str) -> Option<String> {
        self.exchange.host.get_http_response_header(name)
    }

    fn headers(&self) -> Vec<(String, String)> {
        self.exchange.host.get_http_response_headers()
    }

    fn add_header(&self, name: &str, value: &str) {
        self.exchange.host.add_http_response_header(name, value);
    }

    fn set_header(&self, name: &str, value: &str) {
        self.exchange
            .host
            .set_http_response_header(name, Some(value));
    }

    fn set_headers(&self, headers: Vec<(&str, &str)>) {
        self.exchange.host.set_http_response_headers(headers);
    }

    fn remove_header(&self, name: &str) {
        self.exchange.host.set_http_response_header(name, None);
    }
}

impl<S: Event> Exchange<S> {
    pub(crate) fn new(
        reactor: Rc<HttpReactor>,
        host: Rc<dyn Host>,
        first_event: Option<S>,
    ) -> Self {
        let empty_stream = first_event
            .as_ref()
            .map(|e| !e.should_pause())
            .unwrap_or(false);
        Self {
            reactor: reactor.clone(),
            host: host.clone(),
            inner: Rc::new(InnerExchange {
                reactor,
                host,
                first_event: RefCell::new(first_event),
                empty_stream: Cell::new(empty_stream),
                event_count: Cell::new(0),
                offset: Cell::new(0),
            }),
        }
    }

    pub fn event_data(&self) -> Option<EventData<S>>
    where
        S: HeadersEvent,
    {
        let finite_event = self.reactor.cloned_finite_event();
        S::try_from(finite_event)
            .ok()
            .map(|e| EventData::new(self, e, 0))
    }

    #[must_use]
    pub fn event_data_stream(&self) -> EventDataStream<S>
    where
        S: BodyEvent,
    {
        EventDataStream {
            id_and_waker: None,
            exchange: self.inner.clone(),
            buffering: true,
            _lifetime: PhantomData,
        }
    }

    pub(crate) fn static_event_data_stream(&self) -> EventDataStream<'static, S>
    where
        S: BodyEvent,
    {
        EventDataStream {
            id_and_waker: None,
            exchange: self.inner.clone(),
            buffering: false,
            _lifetime: PhantomData,
        }
    }

    pub(crate) async fn wait_for_event<E>(self) -> Exchange<E>
    where
        E: Event,
        S: Before<E>,
    {
        self.wait_for_event_buffering(true).await
    }

    pub(crate) async fn wait_for_event_buffering<E>(self, buffering: bool) -> Exchange<E>
    where
        E: Event,
        S: Before<E>,
    {
        let exchange: Exchange<E> =
            Exchange::new(Rc::clone(&self.reactor), Rc::clone(&self.host), None);

        // Ensure flow resume
        drop(self);

        let mut stream = EventDataStream {
            id_and_waker: None,
            exchange: exchange.inner.clone(),
            buffering,
            _lifetime: PhantomData,
        };
        let first_event = stream.next().await.map(|ed| ed.event);

        *exchange.inner.first_event.borrow_mut() = first_event;
        exchange.inner.offset.set(0);
        exchange.inner.empty_stream.set(false);

        exchange
    }

    pub async fn wait_for_request_headers(self) -> Exchange<RequestHeaders>
    where
        S: Before<RequestHeaders>,
    {
        self.wait_for_event().await
    }

    pub async fn wait_for_request_body(self) -> Exchange<RequestBody>
    where
        S: Before<RequestBody>,
    {
        self.wait_for_event().await
    }

    pub(crate) async fn _wait_for_request_trailers(self) -> Exchange<RequestTrailers>
    where
        S: Before<RequestTrailers>,
    {
        self.wait_for_event().await
    }

    pub async fn wait_for_response_headers(self) -> Exchange<ResponseHeaders>
    where
        S: Before<ResponseHeaders>,
    {
        self.wait_for_event().await
    }

    pub async fn wait_for_response_body(self) -> Exchange<ResponseBody>
    where
        S: Before<ResponseBody>,
    {
        self.wait_for_event().await
    }

    pub(crate) async fn _wait_for_response_trailers(self) -> Exchange<ResponseTrailers>
    where
        S: Before<ResponseTrailers>,
    {
        self.wait_for_event().await
    }

    pub(crate) async fn _wait_for_exchange_complete(self) -> Exchange<ExchangeComplete>
    where
        S: Before<ExchangeComplete>,
    {
        self.wait_for_event().await
    }

    pub fn send_response(self, status_code: u32, headers: Vec<(&str, &str)>, body: Option<&[u8]>)
    where
        S: After<Start> + Before<ResponseHeaders>,
    {
        self.host
            .set_effective_context(self.reactor.context_id().into());
        self.reactor.set_paused(true);
        self.reactor.cancel_request();
        self.host.send_http_response(status_code, headers, body);
    }
}

impl<S: Event> Drop for Exchange<S> {
    fn drop(&mut self) {
        let reactor = &self.reactor;
        let host = &self.host;
        if !reactor.is_done() && reactor.paused() && !reactor.cancelled_request() {
            reactor.set_paused(false);

            host.set_effective_context(reactor.context_id().into());

            match reactor.phase() {
                ExchangePhase::Request => host.resume_http_request(),
                ExchangePhase::Response => host.resume_http_response(),
            }
        }
    }
}

impl<S> EventData<'_, S>
where
    S: BodyEvent,
{
    pub fn offset(&self) -> usize {
        self.offset
    }

    pub fn chunk_size(&self) -> usize {
        self.event.body_size()
    }

    pub fn read_body(&self, offset: usize, max_size: usize) -> Vec<u8> {
        S::read_body(self.exchange.host.deref(), offset, max_size).unwrap_or_default()
    }

    pub fn read_chunk(&self) -> Vec<u8> {
        self.read_body(self.offset, self.event.body_size())
    }

    pub fn read_payload(&self) -> Vec<u8> {
        self.read_body(0, self.event.body_size())
    }
}

/// Manages the access to the data for streamed events.
pub struct EventDataStream<'e, S: Event> {
    exchange: Rc<InnerExchange<S>>,
    id_and_waker: Option<(WakerId, Waker)>,
    buffering: bool,
    _lifetime: PhantomData<&'e ()>,
}

impl<'e, S: Event> EventDataStream<'e, S> {
    fn process_event(&mut self, event: S) -> EventData<'e, S> {
        let exchange = self.exchange.clone();
        let reactor = exchange.reactor.as_ref();
        let pause = self.buffering && event.should_pause();
        reactor.set_paused(pause);
        let offset = exchange.offset.get();
        exchange.offset.set(event.body_size());
        exchange.empty_stream.set(!event.should_pause());
        EventData::from_inner(exchange, event, offset)
    }
}

impl<'e, S: Event> Stream for EventDataStream<'e, S> {
    type Item = EventData<'e, S>;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> Poll<Option<Self::Item>> {
        let reactor = self.exchange.reactor.clone();
        let exchange = self.exchange.clone();

        if exchange.empty_stream.get() {
            if let Some((id, _)) = self.id_and_waker.take() {
                // Deregister the waker from the reactor.
                reactor.remove_waker(S::kind(), id);
            }
            return Poll::Ready(None);
        }

        if reactor.current_event() >= S::kind() {
            let event_data = if let Some(event) = exchange.take_first_event() {
                if let Some((id, _)) = self.id_and_waker.take() {
                    // Deregister the waker from the reactor.
                    reactor.remove_waker(S::kind(), id);
                }
                Some(self.process_event(event))
            } else {
                let event_count = reactor.event_count();

                // Ensure that repeated events are not the same
                if event_count > exchange.event_count.get() {
                    exchange.event_count.set(event_count);
                    if let Ok(event) = S::try_from(reactor.cloned_finite_event()) {
                        Some(self.process_event(event))
                    } else {
                        None
                    }
                } else {
                    return Poll::Pending;
                }
            };

            Poll::Ready(event_data)
        } else {
            match &self.id_and_waker {
                None => {
                    // Register the waker in the reactor.
                    let id = reactor.insert_waker(S::kind(), cx.waker().clone());
                    self.id_and_waker = Some((id, cx.waker().clone()));
                }
                Some((id, w)) if !w.will_wake(cx.waker()) => {
                    // Deregister the waker from the reactor to remove the old waker.
                    reactor.remove_waker(S::kind(), *id);

                    // Register the waker in the reactor with the new waker.
                    let id = reactor.insert_waker(S::kind(), cx.waker().clone());
                    self.id_and_waker = Some((id, cx.waker().clone()));
                }
                Some(_) => {}
            }
            Poll::Pending
        }
    }
}