pdk-classy 1.10.0

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

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

use super::{
    dynamic_exchange::{DynamicExchange, ExchangeEvent},
    entity::{BodyError, BodyHandler},
};
use crate::{event::BodyEvent, host::Host};

#[derive(PartialEq, Debug)]
enum MaxBodySize {
    Unset,
    Unavailable,
    Set(usize),
}

thread_local! {
    static MAX_BODY_SIZE: RefCell<MaxBodySize> = const { RefCell::new(MaxBodySize::Unset) };
}

pub(super) const DEFAULT_MAX_BODY_SIZE: usize = 1024 * 1024;

/// If it was set by [`__set_buffer_size`] returns the size of the buffer size that was set for body validations.
pub fn __buffer_size() -> Option<usize> {
    MAX_BODY_SIZE.with_borrow(|max_size| match max_size {
        MaxBodySize::Set(max) => Some(*max),
        _ => None,
    })
}

/// Checks if [`__set_buffer_size`] was called for this thread.
pub fn __buffer_size_set() -> bool {
    MAX_BODY_SIZE.with_borrow(|max_size| !MaxBodySize::Unset.eq(max_size))
}

/// Set information about the buffer size of the underlying proxy_wasm host. This way we can
/// do size checks when trying to write to the buffer.  
pub fn __set_buffer_size(size: Option<usize>) {
    match size {
        None => {
            MAX_BODY_SIZE.replace(MaxBodySize::Unavailable);
        }
        Some(size) => {
            MAX_BODY_SIZE.replace(MaxBodySize::Set(size));
        }
    }
}

pub(super) fn validate_body_size(size: usize) -> Result<(), BodyError> {
    if let Some(max_size) = __buffer_size() {
        if size >= max_size {
            return Err(BodyError::ExceededBodySize(size));
        }
    } else if size >= DEFAULT_MAX_BODY_SIZE {
        #[cfg(not(feature = "experimental_disable_body_limit_check"))]
        return Err(BodyError::ExceededBodySize(size));
    }
    Ok(())
}

struct Inner<B> {
    host: Rc<dyn Host>,
    event: B,
}

pub struct BodyExchange<B> {
    inner: Option<Inner<B>>,
}

impl<B> BodyExchange<B>
where
    B: BodyEvent + ExchangeEvent,
{
    #[allow(clippy::await_holding_refcell_ref)]
    pub(super) async fn new(exchange: Rc<RefCell<DynamicExchange>>, contains_body: bool) -> Self {
        if !contains_body {
            return Self { inner: None };
        }
        let mut exchange = exchange.borrow_mut();
        let exchange = exchange
            .wait_for_event::<B>()
            .await
            .expect("Must contain body");

        let event = exchange
            .event_data_stream()
            .map(|e| e.event)
            .skip_while(|e| ready(!e.end_of_stream()))
            .next()
            .await
            .expect("End of stream");

        Self {
            inner: Some(Inner {
                host: exchange.host.clone(),
                event,
            }),
        }
    }

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

impl<B: BodyEvent> BodyHandler for BodyExchange<B> {
    fn body(&self) -> Vec<u8> {
        let Some(inner) = self.inner.as_ref() else {
            return Vec::new();
        };
        B::read_body(inner.host.as_ref(), 0, inner.event.body_size()).unwrap_or_default()
    }

    fn set_body(&self, body: &[u8]) -> Result<(), BodyError> {
        let Some(inner) = self.inner.as_ref() else {
            return Err(BodyError::BodyNotSent);
        };

        validate_body_size(body.len())?;

        B::write_body(inner.host.as_ref(), 0, usize::MAX, body);
        Ok(())
    }
}