pdk_classy/extract/

context.rs

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

//! The different contexts where data can be extracted from.

use std::any::TypeId;
use std::collections::HashSet;
use std::{cell::RefCell, convert::Infallible, rc::Rc};

use crate::host::clock::Clock;
use crate::host::grpc::GrpcHost;
use crate::host::shared_data::SharedData;
use crate::middleware::EventHandlerStack;
use crate::{
    context,
    event::{After, Exchange, Start},
    host::Host,
    reactor::{http::HttpReactor, root::RootReactor},
    types::{HttpCid, RootCid},
    BoxFuture,
};

use super::{extractability, AlreadyExtracted, Exclusive, FromContext, FromContextOnce};

#[cfg(feature = "experimental_websocket")]
use crate::reactor::websocket::WebSocketReactor;

#[cfg(feature = "experimental_websocket")]
pub(crate) type WebSocketHandlerFn =
    Box<dyn Fn(Rc<WebSocketReactor>) -> BoxFuture<'static, Result<(), crate::BoxError>> + 'static>;

/// Context representing the configuration of the plugin.
/// The configuration function is the one annotated with the `#[entrypoint]` annotation.
pub struct ConfigureContext {
    pub(crate) host: Rc<dyn Host>,
    pub(crate) clock: Rc<dyn Clock>,
    pub(crate) grpc_host: Rc<dyn GrpcHost>,
    pub(crate) shared_data: Rc<dyn SharedData>,
    #[cfg(feature = "experimental_metrics")]
    pub(crate) metrics: Rc<dyn crate::host::metrics::MetricsHost>,
    pub(crate) root_reactor: Rc<RootReactor>,
    pub(crate) event_handlers: Rc<RefCell<EventHandlerStack>>,
    #[cfg(feature = "experimental_websocket")]
    pub(crate) websocket_upstream_handler: Rc<RefCell<Option<WebSocketHandlerFn>>>,
    #[cfg(feature = "experimental_websocket")]
    pub(crate) websocket_downstream_handler: Rc<RefCell<Option<WebSocketHandlerFn>>>,
    #[cfg(feature = "experimental_websocket")]
    pub(crate) shared_state: Rc<RefCell<Option<Rc<dyn std::any::Any>>>>,
    pub(crate) unique_extractions: RefCell<HashSet<TypeId>>,
}

context!(ConfigureContext);

impl ConfigureContext {
    pub(crate) fn extract_unique<T: 'static>(
        &self,
        op: impl FnOnce() -> T,
    ) -> Result<T, AlreadyExtracted<T>> {
        self.unique_extractions
            .borrow_mut()
            .insert(TypeId::of::<T>())
            .then(op)
            .ok_or_else(AlreadyExtracted::default)
    }
}

impl FromContext<ConfigureContext, extractability::Transitive> for Rc<dyn Host> {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        Ok(context.host.clone())
    }
}

impl FromContext<ConfigureContext, extractability::Transitive> for Rc<dyn Clock> {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        Ok(context.clock.clone())
    }
}

impl FromContext<ConfigureContext, extractability::Transitive> for Rc<dyn SharedData> {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        Ok(context.shared_data.clone())
    }
}

#[cfg(feature = "experimental_metrics")]
impl FromContext<ConfigureContext, extractability::Transitive>
    for Rc<dyn crate::host::metrics::MetricsHost>
{
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        Ok(context.metrics.clone())
    }
}

impl FromContext<ConfigureContext, extractability::Transitive> for Rc<RootReactor> {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        Ok(context.root_reactor.clone())
    }
}

impl FromContext<ConfigureContext, extractability::Transitive> for RootCid {
    type Error = Infallible;

    fn from_context(context: &ConfigureContext) -> Result<Self, Self::Error> {
        Ok(context.root_reactor.context_id())
    }
}

/// Context representing the request filter function of the plugin.
/// The filter function is the one that is provided to the [`Launcher`](crate::bootstrap::Launcher).
pub struct FilterContext {
    parent: Rc<ConfigureContext>,
    http_reactor: Rc<HttpReactor>,
}

impl FilterContext {
    pub(crate) fn new(parent: Rc<ConfigureContext>, http_reactor: Rc<HttpReactor>) -> Self {
        Self {
            parent,
            http_reactor,
        }
    }

    fn parent(&self) -> &ConfigureContext {
        &self.parent
    }

    #[cfg(feature = "experimental_websocket")]
    pub(crate) fn parent_rc(&self) -> &Rc<ConfigureContext> {
        &self.parent
    }
}

context!(FilterContext => ConfigureContext {FilterContext::parent});

impl<E> FromContextOnce<FilterContext> for Exchange<E>
where
    E: After<Start> + 'static,
{
    type Error = AlreadyExtracted<Exchange<E>>;

    type Future<'c> = BoxFuture<'c, Result<Self, Self::Error>>;

    fn from_context_once(context: Exclusive<FilterContext>) -> Self::Future<'_> {
        Box::pin(async move {
            let exchange = Exchange::new(
                context.http_reactor.clone(),
                context.parent.host.clone(),
                Some(Start {
                    _context_id: context.http_reactor.context_id(),
                }),
            );
            Ok(exchange.wait_for_event::<E>().await)
        })
    }
}

impl FromContext<FilterContext> for Rc<HttpReactor> {
    type Error = Infallible;

    fn from_context(context: &FilterContext) -> Result<Self, Self::Error> {
        Ok(context.http_reactor.clone())
    }
}

impl FromContext<FilterContext> for HttpCid {
    type Error = Infallible;

    fn from_context(context: &FilterContext) -> Result<Self, Self::Error> {
        Ok(context.http_reactor.context_id())
    }
}

/// Context representing a WebSocket upstream filter (client→server).
/// Provides access to all PDK injectables for WebSocket frame handling.
#[cfg(feature = "experimental_websocket")]
pub struct UpgradeUpstreamContext<S = ()> {
    parent: Rc<ConfigureContext>,
    state: Rc<S>,
    reactor: Rc<WebSocketReactor>,
}

#[cfg(feature = "experimental_websocket")]
impl<S> UpgradeUpstreamContext<S> {
    pub(crate) fn new(
        parent: Rc<ConfigureContext>,
        state: Rc<S>,
        reactor: Rc<WebSocketReactor>,
    ) -> Self {
        Self {
            parent,
            state,
            reactor,
        }
    }

    pub(crate) fn state(&self) -> &Rc<S> {
        &self.state
    }

    pub(crate) fn reactor(&self) -> &Rc<WebSocketReactor> {
        &self.reactor
    }

    fn parent(&self) -> &ConfigureContext {
        &self.parent
    }
}

#[cfg(feature = "experimental_websocket")]
context!(<S> UpgradeUpstreamContext<S> => ConfigureContext {UpgradeUpstreamContext::parent});

/// Context representing a WebSocket downstream filter (server→client).
/// Provides access to all PDK injectables for WebSocket frame handling.
#[cfg(feature = "experimental_websocket")]
pub struct UpgradeDownstreamContext<S = ()> {
    parent: Rc<ConfigureContext>,
    state: Rc<S>,
    reactor: Rc<WebSocketReactor>,
}

#[cfg(feature = "experimental_websocket")]
impl<S> UpgradeDownstreamContext<S> {
    pub(crate) fn new(
        parent: Rc<ConfigureContext>,
        state: Rc<S>,
        reactor: Rc<WebSocketReactor>,
    ) -> Self {
        Self {
            parent,
            state,
            reactor,
        }
    }

    pub(crate) fn state(&self) -> &Rc<S> {
        &self.state
    }

    pub(crate) fn reactor(&self) -> &Rc<WebSocketReactor> {
        &self.reactor
    }

    fn parent(&self) -> &ConfigureContext {
        &self.parent
    }
}

#[cfg(feature = "experimental_websocket")]
context!(<S> UpgradeDownstreamContext<S> => ConfigureContext {UpgradeDownstreamContext::parent});