pdk-classy 1.10.0

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

use std::cell::RefCell;
use std::future::Future;
use std::pin::pin;
use std::rc::Rc;
use std::task::Poll;

use crate::{
    event::{Exchange, ExchangeComplete, RequestHeaders, ResponseHeaders, Start},
    extract::{context::FilterContext, AlreadyExtracted, Exclusive, FromContext, FromContextOnce},
    handler::{ExtractionError, Handler, IntoHandler},
    hl::state::{CreateHandler, DoneHandler, EmptyCreateHandler},
    host::Host,
    reactor::http::{FlowStatus, HttpReactor},
    BoxFuture,
};

#[cfg(feature = "experimental_websocket")]
use crate::{handler::IntoHandlerResult, BoxError};

#[cfg(feature = "experimental_websocket")]
use crate::extract::context::{ConfigureContext, UpgradeDownstreamContext, UpgradeUpstreamContext};

use super::{
    context::{RequestContext, ResponseContext},
    dynamic_exchange::DynamicExchange,
    request_data::RequestData,
    Flow, IntoFlow,
};

pub struct RequestFilter<ReqHnd, Sf = EmptyCreateHandler, Done = ()> {
    pub(super) request_handler: ReqHnd,
    pub(super) state_factory: Sf,
    pub(super) done_handler: Done,
}

impl<ReqHnd, Sf, Done> RequestFilter<ReqHnd, Sf, Done>
where
    ReqHnd: Handler<RequestContext<Sf::State>>,
    ReqHnd::Output: IntoFlow,
    Sf: CreateHandler,
{
    /// Creates a Response filter from a handler.
    pub fn on_response<ResHnd, I>(
        self,
        response_handler: ResHnd,
    ) -> DualFilter<ReqHnd, ResHnd::Handler, Sf, Done>
    where
        ResHnd: IntoHandler<
            ResponseContext<<ReqHnd::Output as IntoFlow>::RequestData, Sf::State>,
            I,
            Output = (),
        >,
    {
        DualFilter {
            request_handler: self.request_handler,
            response_handler: response_handler.into_handler(),
            state_factory: self.state_factory,
            done_handler: self.done_handler,
        }
    }
}

impl<ReqHnd, T, Sf, Done> Handler<FilterContext> for RequestFilter<ReqHnd, Sf, Done>
where
    ReqHnd: Handler<RequestContext<Sf::State>>,
    ReqHnd::Output: IntoFlow<RequestData = T>,
    Sf: CreateHandler<State: Clone + 'static>,
    Done: DoneHandler<Sf::State>,
{
    type Output = ();

    type Future<'h>
        = BoxFuture<'h, Result<Self::Output, ExtractionError>>
    where
        Self: 'h;

    fn call<'h>(&'h self, context: FilterContext) -> Self::Future<'h>
    where
        Self: 'h,
    {
        #[allow(clippy::await_holding_refcell_ref)]
        Box::pin(async move {
            let context = Rc::new(context);
            let exclusive_context = Exclusive::new(context.as_ref());
            let exchange = <Exchange<RequestHeaders>>::from_context_once(exclusive_context)
                .await
                .map_err(|e| ExtractionError(e.into()))?;

            let reactor = Rc::clone(&exchange.reactor);
            let exchange = Rc::new(RefCell::new(DynamicExchange::new(exchange)));
            let state = Rc::new(self.state_factory.create());

            // Store state for WebSocket handlers
            #[cfg(feature = "experimental_websocket")]
            {
                *context.parent_rc().shared_state.borrow_mut() =
                    Some(Rc::clone(&state) as Rc<dyn std::any::Any>);
            }

            // Needed by wait_for_on_done
            let host: Rc<dyn Host> =
                FromContext::<_, crate::extract::extractability::Transitive>::from_context(
                    &*context,
                )
                .unwrap();

            let result = may_suspend_request(&reactor, async {
                let request_context =
                    RequestContext::new(context, exchange.clone(), Rc::clone(&state));

                let flow = self
                    .request_handler
                    .call(request_context)
                    .await?
                    .into_flow();
                if let Flow::Break(response) = flow {
                    let mut exchange = exchange.borrow_mut();
                    exchange.send_response(
                        response.status_code(),
                        response.headers(),
                        response.body(),
                    );
                }

                Ok(())
            })
            .await
            .unwrap_or(Ok(()));

            // Wait for ExchangeComplete event (triggered by on_done from proxy-wasm)
            wait_for_on_done(Rc::clone(&reactor), host).await;

            // Execute done handler. If try_unwrap fails (WebSocket closures still hold Rc refs),
            // clone the state — done_handler receives a snapshot at HTTP exchange completion.
            match Rc::try_unwrap(state) {
                Ok(state) => self.done_handler.done(state),
                Err(rc) => self.done_handler.done((*rc).clone()),
            }

            result
        })
    }
}

/// Creates a Request filter from a handler.
pub fn on_request<ReqHnd, I>(request_handler: ReqHnd) -> RequestFilter<ReqHnd::Handler>
where
    ReqHnd: IntoHandler<RequestContext<()>, I>,
    ReqHnd::Output: IntoFlow,
{
    RequestFilter {
        request_handler: request_handler.into_handler(),
        state_factory: EmptyCreateHandler,
        done_handler: (),
    }
}

pub struct ResponseFilter<ResHnd, Sf = EmptyCreateHandler, Done = ()> {
    pub(super) response_handler: ResHnd,
    pub(super) state_factory: Sf,
    pub(super) done_handler: Done,
}

impl<ResHnd, Sf, Done> Handler<FilterContext> for ResponseFilter<ResHnd, Sf, Done>
where
    ResHnd: Handler<ResponseContext<(), Sf::State>, Output = ()>,
    Sf: CreateHandler<State: Clone + 'static>,
    Done: DoneHandler<Sf::State>,
{
    type Output = ();

    type Future<'h>
        = BoxFuture<'h, Result<Self::Output, ExtractionError>>
    where
        Self: 'h;

    fn call<'h>(&'h self, context: FilterContext) -> Self::Future<'h>
    where
        Self: 'h,
    {
        #[allow(clippy::await_holding_refcell_ref)]
        Box::pin(async move {
            let context = Rc::new(context);
            let exclusive_context = Exclusive::new(context.as_ref());
            let exchange = <Exchange<ResponseHeaders>>::from_context_once(exclusive_context)
                .await
                .map_err(|e| ExtractionError(e.into()))?;
            let reactor = Rc::clone(&exchange.reactor);

            let exchange = Rc::new(RefCell::new(DynamicExchange::new(exchange)));
            let state = Rc::new(self.state_factory.create());

            // Store state for WebSocket handlers
            #[cfg(feature = "experimental_websocket")]
            {
                *context.parent_rc().shared_state.borrow_mut() =
                    Some(Rc::clone(&state) as Rc<dyn std::any::Any>);
            }

            // Needed by wait_for_on_done
            let host: Rc<dyn Host> =
                FromContext::<_, crate::extract::extractability::Transitive>::from_context(
                    &*context,
                )
                .unwrap();

            let result = may_suspend_response(&reactor, async {
                let response_context = ResponseContext::new(
                    context,
                    exchange.clone(),
                    RequestData::Break,
                    Rc::clone(&state),
                );
                self.response_handler.call(response_context).await?;

                Ok(())
            })
            .await
            .unwrap_or(Ok(()));

            // Wait for ExchangeComplete event (triggered by on_done from proxy-wasm)
            wait_for_on_done(Rc::clone(&reactor), host).await;

            // Execute done handler. If try_unwrap fails (WebSocket closures still hold Rc refs),
            // clone the state — done_handler receives a snapshot at HTTP exchange completion.
            match Rc::try_unwrap(state) {
                Ok(state) => self.done_handler.done(state),
                Err(rc) => self.done_handler.done((*rc).clone()),
            }

            result
        })
    }
}

/// Creates a Response filter from a handler.
pub fn on_response<ResHnd, I>(response_handler: ResHnd) -> ResponseFilter<ResHnd::Handler>
where
    ResHnd: IntoHandler<ResponseContext<(), ()>, I, Output = ()>,
{
    ResponseFilter {
        response_handler: response_handler.into_handler(),
        state_factory: EmptyCreateHandler,
        done_handler: (),
    }
}

pub struct DualFilter<ReqHnd, ResHnd, Sf = EmptyCreateHandler, Done = ()> {
    pub(super) request_handler: ReqHnd,
    pub(super) response_handler: ResHnd,
    pub(super) state_factory: Sf,
    pub(super) done_handler: Done,
}

impl<ReqHnd, ResHnd, Sf, Done> Handler<FilterContext> for DualFilter<ReqHnd, ResHnd, Sf, Done>
where
    ReqHnd: Handler<RequestContext<Sf::State>>,
    ReqHnd::Output: IntoFlow,
    ResHnd:
        Handler<ResponseContext<<ReqHnd::Output as IntoFlow>::RequestData, Sf::State>, Output = ()>,
    Sf: CreateHandler<State: Clone + 'static>,
    Done: DoneHandler<Sf::State>,
{
    type Output = ();

    type Future<'h>
        = BoxFuture<'h, Result<Self::Output, ExtractionError>>
    where
        Self: 'h;

    fn call<'h>(&'h self, context: FilterContext) -> Self::Future<'h>
    where
        Self: 'h,
    {
        #[allow(clippy::await_holding_refcell_ref)]
        Box::pin(async move {
            let context = Rc::new(context);
            let exclusive_context = Exclusive::new(context.as_ref());
            let exchange = <Exchange<RequestHeaders>>::from_context_once(exclusive_context)
                .await
                .map_err(|_| {
                    ExtractionError(AlreadyExtracted::<Exchange<RequestHeaders>>::default().into())
                })?;
            let reactor = Rc::clone(&exchange.reactor);

            let state = Rc::new(self.state_factory.create());

            // Store a strong Rc so WebSocket handler closures keep
            // state alive after the configure future completes.
            #[cfg(feature = "experimental_websocket")]
            {
                *context.parent_rc().shared_state.borrow_mut() =
                    Some(Rc::clone(&state) as Rc<dyn std::any::Any>);
            }

            let exchange = Rc::new(RefCell::new(DynamicExchange::new(exchange)));

            // Needed by wait_for_on_done
            let host: Rc<dyn Host> =
                FromContext::<_, crate::extract::extractability::Transitive>::from_context(
                    &*context,
                )
                .unwrap();

            let request_data = may_suspend_request(&reactor, async {
                let request_context =
                    RequestContext::new(context.clone(), exchange.clone(), Rc::clone(&state));

                let flow = self
                    .request_handler
                    .call(request_context)
                    .await?
                    .into_flow();
                match flow {
                    Flow::Break(response) => {
                        exchange.borrow_mut().send_response(
                            response.status_code(),
                            response.headers(),
                            response.body(),
                        );
                        Ok(RequestData::Break)
                    }
                    Flow::Continue(data) => Ok(RequestData::Continue(data)),
                }
            })
            .await
            .unwrap_or(Ok(RequestData::Cancel))?;

            let exclusive_context = Exclusive::new(context.as_ref());
            let exchange = <Exchange<ResponseHeaders>>::from_context_once(exclusive_context)
                .await
                .map_err(|_| {
                    ExtractionError(AlreadyExtracted::<Exchange<ResponseHeaders>>::default().into())
                })?;

            let exchange = Rc::new(RefCell::new(DynamicExchange::new(exchange)));

            let result = may_suspend_response(
                &reactor,
                self.response_handler.call(ResponseContext::new(
                    context.clone(),
                    exchange.clone(),
                    request_data,
                    state.clone(),
                )),
            )
            .await
            .unwrap_or(Ok(()));

            // Wait for ExchangeComplete event (triggered by on_done from proxy-wasm)
            wait_for_on_done(Rc::clone(&reactor), host).await;

            // Execute done handler. If try_unwrap fails (WebSocket closures still hold Rc refs),
            // clone the state — done_handler receives a snapshot at HTTP exchange completion.
            match Rc::try_unwrap(state) {
                Ok(state) => self.done_handler.done(state),
                Err(rc) => self.done_handler.done((*rc).clone()),
            }

            result
        })
    }
}

async fn wait_for_on_done(reactor: Rc<HttpReactor>, host: Rc<dyn Host>) {
    let complete_exchange: Exchange<Start> = Exchange::new(reactor, host, None);
    let _ = complete_exchange.wait_for_event::<ExchangeComplete>().await;
}

async fn may_suspend_request<F: Future>(reactor: &HttpReactor, task: F) -> Option<F::Output> {
    let mut task = pin!(task);

    std::future::poll_fn(move |cx| match reactor.request_status() {
        FlowStatus::Suspended => {
            let id: u32 = reactor.context_id().into();
            log::debug!("Request for filter with context id {id} has been suspended.");

            Poll::Ready(None)
        }
        FlowStatus::Unsuspended => task.as_mut().poll(cx).map(Some),
    })
    .await
}

async fn may_suspend_response<F: Future>(reactor: &HttpReactor, task: F) -> Option<F::Output> {
    let mut task = pin!(task);

    std::future::poll_fn(move |cx| match reactor.response_status() {
        FlowStatus::Suspended => {
            let id: u32 = reactor.context_id().into();
            log::debug!("Response for filter with context id {id} has been suspended.");

            Poll::Ready(None)
        }
        FlowStatus::Unsuspended => task.as_mut().poll(cx).map(Some),
    })
    .await
}

/// Filter that wraps a base HTTP filter and adds upstream/downstream WebSocket frame handlers after a successful upgrade.
#[cfg(feature = "experimental_websocket")]
pub struct WebSocketFilter<BaseFilter, Sf, WsUp = (), WsDown = ()> {
    pub(super) base_filter: BaseFilter,
    pub(super) websocket_upstream: WsUp,
    pub(super) websocket_downstream: WsDown,
    pub(super) _sf: std::marker::PhantomData<Sf>,
}

#[cfg(feature = "experimental_websocket")]
impl<BaseFilter, Sf, WsUp, WsDown> WebSocketFilter<BaseFilter, Sf, WsUp, WsDown>
where
    Sf: CreateHandler,
{
    /// Adds an upstream (client→server) WebSocket frame handler.
    ///
    /// The handler may return `()` (its frame loop runs until it completes) or a
    /// `Result<(), E>` (returning `Err` ends the loop; the error is logged, see the
    /// `Handler` impl below).
    pub fn on_upgrade_upstream<WsUpHnd, I>(
        self,
        websocket_upstream: WsUpHnd,
    ) -> WebSocketFilter<BaseFilter, Sf, WsUpHnd::Handler, WsDown>
    where
        WsUpHnd: IntoHandler<UpgradeUpstreamContext<Sf::State>, I, Output: IntoHandlerResult>,
    {
        WebSocketFilter {
            base_filter: self.base_filter,
            websocket_upstream: websocket_upstream.into_handler(),
            websocket_downstream: self.websocket_downstream,
            _sf: std::marker::PhantomData,
        }
    }

    /// Adds a downstream (server→client) WebSocket frame handler.
    ///
    /// The handler may return `()` (its frame loop runs until it completes) or a
    /// `Result<(), E>` (returning `Err` ends the loop; the error is logged, see the
    /// `Handler` impl below).
    pub fn on_upgrade_downstream<WsDownHnd, I>(
        self,
        websocket_downstream: WsDownHnd,
    ) -> WebSocketFilter<BaseFilter, Sf, WsUp, WsDownHnd::Handler>
    where
        WsDownHnd: IntoHandler<UpgradeDownstreamContext<Sf::State>, I, Output: IntoHandlerResult>,
    {
        WebSocketFilter {
            base_filter: self.base_filter,
            websocket_upstream: self.websocket_upstream,
            websocket_downstream: websocket_downstream.into_handler(),
            _sf: std::marker::PhantomData,
        }
    }
}

// RFC 6455 1011 ("internal error") Close frame. The upstream plane is client→server,
// so its Close must be MASKED (MASK bit set, 4-byte key; a zero key is legal). The
// downstream plane is server→client and must be UNMASKED. Handlers write fully-framed
// bytes; the proxy forwards them verbatim (no masking added by the host).
#[cfg(feature = "experimental_websocket")]
const WS_CLOSE_1011_UPSTREAM: [u8; 8] = [0x88, 0x82, 0x00, 0x00, 0x00, 0x00, 0x03, 0xF3];
#[cfg(feature = "experimental_websocket")]
const WS_CLOSE_1011_DOWNSTREAM: [u8; 4] = [0x88, 0x02, 0x03, 0xF3];

/// Parks until this side's upgrade gate opens, registering `register_waker` so the HTTP context
/// can wake it. `is_upgraded` checks the per-direction gate; returns immediately if already open.
#[cfg(feature = "experimental_websocket")]
async fn await_upgrade(is_upgraded: impl Fn() -> bool, register_waker: impl Fn(std::task::Waker)) {
    std::future::poll_fn(|cx| {
        if is_upgraded() {
            Poll::Ready(())
        } else {
            register_waker(cx.waker().clone());
            Poll::Pending
        }
    })
    .await
}

/// Recovers the per-connection state (stashed by the base filter) from the configure context.
#[cfg(feature = "experimental_websocket")]
fn recover_state<S: 'static>(configure_context: &ConfigureContext) -> Option<Rc<S>> {
    configure_context
        .shared_state
        .borrow()
        .as_ref()
        .and_then(|s| s.clone().downcast::<S>().ok())
}

// Handler implementation that delegates to the base filter and drives the WebSocket frame handlers
// by reference, concurrently with the base filter, so their closures can borrow the `configure`
// locals (like `on_request`/`on_response`) instead of requiring `Clone + 'static`.
#[cfg(feature = "experimental_websocket")]
impl<BaseFilter, Sf, WsUp, WsDown> Handler<FilterContext>
    for WebSocketFilter<BaseFilter, Sf, WsUp, WsDown>
where
    BaseFilter: Handler<FilterContext, Output = ()>,
    Sf: CreateHandler,
    Sf::State: Clone + 'static,
    WsUp: Handler<UpgradeUpstreamContext<Sf::State>, Output: IntoHandlerResult>,
    WsDown: Handler<UpgradeDownstreamContext<Sf::State>, Output: IntoHandlerResult>,
{
    type Output = ();

    type Future<'h>
        = BoxFuture<'h, Result<Self::Output, ExtractionError>>
    where
        Self: 'h;

    fn call<'h>(&'h self, context: FilterContext) -> Self::Future<'h>
    where
        Self: 'h,
    {
        Box::pin(async move {
            let configure_context = Rc::clone(context.parent_rc());
            let ws_reactor = Rc::clone(context.websocket_reactor());

            // Host used to write the Close frame into the held buffer on handler error. Extracted
            // the same way UpstreamState/DownstreamState do (see hl/websocket.rs), via the
            // Transitive ConfigureContext -> Rc<dyn Host> extractor (extract/context.rs).
            let host: Rc<dyn Host> =
                FromContext::<_, crate::extract::extractability::Transitive>::from_context(
                    &*configure_context,
                )
                .unwrap();

            // Upstream (client→server) pump: parks on the upgrade gate, then runs the handler once.
            // The handler owns its own frame loop via UpstreamState::next()/accumulate().
            let upstream = async {
                await_upgrade(
                    || ws_reactor.upstream_upgraded(),
                    |w| ws_reactor.register_upstream_upgrade_waker(w),
                )
                .await;
                let result = match recover_state::<Sf::State>(&configure_context) {
                    Some(state) => {
                        let ctx = UpgradeUpstreamContext::new(
                            Rc::clone(&configure_context),
                            state,
                            Rc::clone(&ws_reactor),
                        );
                        self.websocket_upstream
                            .call(ctx)
                            .await
                            .map_err(|e| Box::new(e) as BoxError)
                            .and_then(|r| r.into_handler_result())
                    }
                    None => Err(Box::from("WebSocket upstream: state unavailable") as BoxError),
                };
                if let Err(e) = result {
                    log::error!("WebSocket upstream handler error: {e:?}");
                    // Cut-on-error: overwrite the held upstream buffer with a 1011 Close frame,
                    // release the pause so the in-flight on_http_request_body dispatch forwards it,
                    // and mark the connection closing so the downstream pump emits its own Close.
                    host.set_http_request_body(0, usize::MAX, &WS_CLOSE_1011_UPSTREAM);
                    ws_reactor.set_upstream_paused(false);
                    ws_reactor.mark_closing();
                }
            };

            // Downstream (server→client) pump, symmetric to the upstream pump.
            let downstream = async {
                await_upgrade(
                    || ws_reactor.downstream_upgraded(),
                    |w| ws_reactor.register_downstream_upgrade_waker(w),
                )
                .await;
                let result = match recover_state::<Sf::State>(&configure_context) {
                    Some(state) => {
                        let ctx = UpgradeDownstreamContext::new(
                            Rc::clone(&configure_context),
                            state,
                            Rc::clone(&ws_reactor),
                        );
                        self.websocket_downstream
                            .call(ctx)
                            .await
                            .map_err(|e| Box::new(e) as BoxError)
                            .and_then(|r| r.into_handler_result())
                    }
                    None => Err(Box::from("WebSocket downstream: state unavailable") as BoxError),
                };
                if let Err(e) = result {
                    log::error!("WebSocket downstream handler error: {e:?}");
                    // Cut-on-error: overwrite the held downstream buffer with a 1011 Close frame,
                    // release the pause so the in-flight on_http_response_body dispatch forwards it,
                    // and mark the connection closing so the upstream pump emits its own Close.
                    host.set_http_response_body(0, usize::MAX, &WS_CLOSE_1011_DOWNSTREAM);
                    ws_reactor.set_downstream_paused(false);
                    ws_reactor.mark_closing();
                }
            };

            // Drive the base filter and both pumps concurrently. The base filter completes when
            // the HTTP exchange ends (ExchangeComplete); at that point the connection is closed and
            // the still-parked or looping pumps are dropped. Each pump is fused so a completed pump
            // (handler returned, or the no-op `()` handler) is not polled again.
            //
            // Error handling is cut-on-error: a handler that returns `Err` (or whose state is
            // unavailable) has its error logged above, writes a 1011 Close frame on its own plane,
            // and marks the connection closing. mark_closing also unpauses both planes so neither
            // holds a buffer. Once closing, the pumps are treated as done so no user handler code is
            // re-entered; the base filter keeps being polled to completion (ExchangeComplete). The
            // opposite plane does not emit its own Close: it drains (unpaused), and the connection
            // tears down via the erroring plane's Close plus the peer's Close echo (RFC 6455).
            let mut base = pin!(self.base_filter.call(context));
            let mut upstream = pin!(upstream);
            let mut downstream = pin!(downstream);
            let mut upstream_done = false;
            let mut downstream_done = false;

            std::future::poll_fn(|cx| {
                if ws_reactor.is_closing() {
                    // Do not re-enter user pump code once closing; drive base to ExchangeComplete.
                    upstream_done = true;
                    downstream_done = true;
                } else {
                    if !upstream_done && upstream.as_mut().poll(cx).is_ready() {
                        upstream_done = true;
                    }
                    if !downstream_done && downstream.as_mut().poll(cx).is_ready() {
                        downstream_done = true;
                    }
                }
                base.as_mut().poll(cx)
            })
            .await
        })
    }
}