fastly 0.12.0

//! HTTP bodies.

use crate::{
    abi::{self, FastlyStatus, MultiValueHostcallError},
    error::{BufferSizeError, HandleError, HandleKind},
    experimental::{BodyHandleError, BodyHandleExt},
};
use bytes::BytesMut;
use fastly_shared::BodyWriteEnd;
use http::header::{HeaderMap, HeaderName, HeaderValue};
use std::sync::atomic::{AtomicBool, Ordering};
use std::{
    io::{BufReader, Read, Write},
    mem::ManuallyDrop,
};

/// A low-level interface to HTTP bodies.
///
/// For most applications, you should use [`Body`][`crate::Body`] instead of this
/// interface. See the top-level [`handle`][`crate::handle`] documentation for more details.
///
/// This type implements [`Read`] to read bytes from the beginning of a body, and [`Write`] to write
/// to the end of a body. Note that these operations are unbuffered, unlike the same operations on
/// the higher-level [`Body`][`crate::Body`] type.
#[derive(Debug, Eq, Hash, PartialEq)]
#[repr(transparent)]
pub struct BodyHandle {
    pub(crate) handle: u32,
}

/// A flag representing whether or not we have taken the client body.
pub(crate) static GOT_CLIENT_BODY: AtomicBool = AtomicBool::new(false);

impl BodyHandle {
    /// An invalid body handle.
    ///
    /// This is primarily useful to represent uninitialized values when using the interfaces in
    /// [`fastly_sys`].
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub const INVALID: Self = Self {
        handle: fastly_shared::INVALID_BODY_HANDLE,
    };

    /// Returns `true` if the body handle is valid.
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub const fn is_valid(&self) -> bool {
        !self.is_invalid()
    }

    /// Returns `true` if the body handle is invalid.
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub const fn is_invalid(&self) -> bool {
        self.handle == fastly_shared::INVALID_BODY_HANDLE
    }

    /// Make a handle from its underlying representation.
    ///
    /// # Safety
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub unsafe fn from_u32(handle: u32) -> Self {
        Self { handle }
    }

    /// Get the underlying representation of the handle.
    ///
    /// # Safety
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub unsafe fn as_u32(&self) -> u32 {
        self.handle
    }

    /// Get a mutable reference to the underlying `u32` representation of the handle.
    ///
    /// # Safety
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub unsafe fn as_u32_mut(&mut self) -> &mut u32 {
        &mut self.handle
    }

    /// Turn a handle into its representation without closing the underlying resource.
    ///
    /// # Safety
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    #[cfg_attr(
        not(target_env = "p1"),
        deprecated(
            since = "0.11.6",
            note = "This code will need to be updated for wasip2."
        )
    )]
    pub fn into_u32(self) -> u32 {
        unsafe { ManuallyDrop::new(self).as_u32() }
    }

    /// Set `GOT_CLIENT_BODY` flag to show we've taken the client body.
    ///
    /// This will panic if the flag has already been set by someone else.
    pub(crate) fn set_got_client() {
        if GOT_CLIENT_BODY.swap(true, Ordering::SeqCst) {
            panic!("cannot get more than one handle to the client body per execution");
        }
    }

    /// Get a handle to the client request body.
    ///
    /// This handle may only be retrieved once per execution, either through this function or
    /// through [`client_request_and_body()`][`crate::handle::client_request_and_body()`].
    pub fn from_client() -> Self {
        Self::set_got_client();
        let mut handle = BodyHandle::INVALID;
        let status = unsafe {
            abi::fastly_http_req::body_downstream_get(std::ptr::null_mut(), handle.as_u32_mut())
        };
        match status.result().map(|_| handle) {
            Ok(h) if h.is_valid() => h,
            _ => panic!("fastly_http_req::body_downstream_get failed"),
        }
    }

    /// Acquire a new, empty body handle.
    #[allow(clippy::new_without_default)]
    pub fn new() -> Self {
        let mut handle = BodyHandle::INVALID;
        let status = unsafe { abi::fastly_http_body::new(handle.as_u32_mut()) };
        match status.result().map(|_| handle) {
            Ok(h) if h.is_valid() => h,
            _ => panic!("fastly_http_body::new failed"),
        }
    }

    /// Append another body onto the end of this body.
    ///
    #[doc = include_str!("../../../docs/snippets/body-append-constant-time.md")]
    ///
    /// The other body will no longer be valid after this call.
    pub fn append(&mut self, other: BodyHandle) {
        unsafe { abi::fastly_http_body::append(self.as_u32(), other.into_u32()) }
            .result()
            .expect("fastly_http_body::append failed")
    }

    /// Read the entirety of the body into a byte vector.
    ///
    #[doc = include_str!("../../../docs/snippets/buffers-body-handle.md")]
    pub fn into_bytes(self) -> Vec<u8> {
        let mut body = vec![];
        let mut bufread = BufReader::new(self);
        bufread
            .read_to_end(&mut body)
            .expect("fastly_http_body::read failed");
        body
    }

    /// Read the entirety of the body into a `String`, interpreting the bytes as UTF-8.
    ///
    #[doc = include_str!("../../../docs/snippets/buffers-body-handle.md")]
    ///
    /// # Panics
    ///
    #[doc = include_str!("../../../docs/snippets/panics-body-utf8.md")]
    pub fn into_string(self) -> String {
        let mut body = String::new();
        let mut bufread = BufReader::new(self);
        bufread
            .read_to_string(&mut body)
            .expect("fastly_http_body::read failed");
        body
    }

    pub(crate) fn write_front(&mut self, bytes: &[u8]) -> usize {
        let mut nwritten = 0;
        let status = unsafe {
            abi::fastly_http_body::write(
                self.as_u32(),
                bytes.as_ptr(),
                bytes.len(),
                BodyWriteEnd::Front,
                &mut nwritten,
            )
        };
        assert!(status.is_ok(), "fastly_http_body::write_front failed");
        assert!(
            nwritten == bytes.len(),
            "fastly_http_body::write_front didn't fully write"
        );
        nwritten
    }

    /// Close the BodyHandle by removing it from the host Session. This will close
    /// out streaming and non streaming bodies. Care must be taken to only
    /// close out the handle if the body is not being written to or read from
    /// anymore. If the handle has already been closed an error will be
    /// returned.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use std::io::Write;
    /// # use fastly::handle::BodyHandle;
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let mut body = BodyHandle::new();
    /// body.write_all(b"You're already closed.")?;
    /// // The handle is not being used in a request and doesn't refer to any
    /// // response body so we can close this out
    /// body.close()?;
    /// # Ok(())
    /// # }
    /// ```
    pub fn close(self) -> Result<(), HandleError> {
        match unsafe { abi::fastly_http_body::close(self.into_u32()) } {
            FastlyStatus::OK => Ok(()),
            _ => Err(HandleError::ClosedHandle(HandleKind::Body)),
        }
    }

    /// Append all trailers to the Body. This is used by `StreamingBody::finish_with_trailers()`.
    pub(crate) fn trailers_replace(&mut self, trailers: &HeaderMap) {
        for (n, v) in trailers.iter() {
            self.append_trailer(n, v);
        }
    }

    /// Add a body trailing header with given value.
    pub(crate) fn append_trailer(&mut self, name: &HeaderName, value: &HeaderValue) {
        let name_bytes: &[u8] = name.as_ref();
        let value_bytes: &[u8] = value.as_ref();
        unsafe {
            abi::fastly_http_body::trailer_append(
                self.as_u32(),
                name_bytes.as_ptr(),
                name_bytes.len(),
                value_bytes.as_ptr(),
                value_bytes.len(),
            )
        }
        .result()
        .expect("fastly_http_body::trailer_append returned error");
    }

    pub(crate) fn get_trailer_names_impl(
        &self,
        mut initial_buf_size: usize,
        max_buf_size: Option<usize>,
    ) -> Result<Box<impl Iterator<Item = Result<HeaderName, BufferSizeError>> + '_>, BodyHandleError>
    {
        if let Some(max) = max_buf_size {
            initial_buf_size = std::cmp::min(initial_buf_size, max);
        }

        let mut iter = abi::MultiValueHostcall::new(
            b'\0',
            initial_buf_size,
            max_buf_size,
            move |buf, buf_size, cursor, ending_cursor, nwritten| unsafe {
                abi::fastly_http_body::trailer_names_get(
                    self.as_u32(),
                    buf,
                    buf_size,
                    cursor,
                    ending_cursor,
                    nwritten,
                )
            },
        )
        .peekable();
        // NB: We know that the AGAIN error will only be returned upon the first element,
        // so we look for it up front and allow everything else to pass through to the block
        // below.
        if let Some(res) = iter.peek() {
            use MultiValueHostcallError::ClosureError;
            match res {
                Err(ClosureError(e)) if *e == FastlyStatus::AGAIN => {
                    return Err(BodyHandleError::TrailersNotReady)
                }
                _ => (),
            }
        }

        Ok(Box::new(iter.map(move |res| {
            use MultiValueHostcallError::{BufferTooSmall, ClosureError};
            match res {
                // we trust that the hostcall is giving us valid header bytes
                Ok(name_bytes) => Ok(HeaderName::try_from(Vec::from(name_bytes)).unwrap()),
                // return an error if the buffer was not large enough; this means there must be a
                // max buffer size
                Err(BufferTooSmall { needed_buf_size }) => Err(BufferSizeError::header_name(
                    max_buf_size
                        .expect("maximum buffer size must exist if a buffer size error occurs"),
                    needed_buf_size,
                )),
                // panic if the hostcall failed for some other reason
                Err(ClosureError(e)) => {
                    panic!("fastly_http_resp::trailer_names_get returned error: {e:?}",)
                }
            }
        })))
    }

    pub(crate) fn get_trailer_values_impl<'a>(
        &'a self,
        name: &'a HeaderName,
        mut initial_buf_size: usize,
        max_buf_size: Option<usize>,
    ) -> Result<Box<impl Iterator<Item = Result<HeaderValue, BufferSizeError>> + 'a>, BodyHandleError>
    {
        if let Some(max) = max_buf_size {
            initial_buf_size = std::cmp::min(initial_buf_size, max);
        }

        let mut iter = abi::MultiValueHostcall::new(
            b'\0',
            initial_buf_size,
            max_buf_size,
            move |buf, buf_size, cursor, ending_cursor, nwritten| unsafe {
                let name: &[u8] = name.as_ref();
                abi::fastly_http_body::trailer_values_get(
                    self.as_u32(),
                    name.as_ptr(),
                    name.len(),
                    buf,
                    buf_size,
                    cursor,
                    ending_cursor,
                    nwritten,
                )
            },
        )
        .peekable();
        // NB: We know that the AGAIN error will only be returned upon the first element,
        // so we look for it up front and allow everything else to pass through to the block
        // below.
        if let Some(res) = iter.peek() {
            use MultiValueHostcallError::ClosureError;
            match res {
                Err(ClosureError(e)) if *e == FastlyStatus::AGAIN => {
                    return Err(BodyHandleError::TrailersNotReady)
                }
                _ => (),
            }
        }

        Ok(Box::new(iter.map(move |res| {
            use MultiValueHostcallError::{BufferTooSmall, ClosureError};
            match res {
                Ok(value_bytes) => {
                    let header_value =
                        unsafe { HeaderValue::from_maybe_shared_unchecked(value_bytes) };
                    Ok(header_value)
                }
                Err(BufferTooSmall { needed_buf_size }) => Err(BufferSizeError::header_value(
                    max_buf_size
                        .expect("maximum buffer size must exist if a buffer size error occurs"),
                    needed_buf_size,
                )),
                Err(ClosureError(e)) => {
                    panic!(
                        "fastly_http_resp::trailer_values_get returned error: {:?}",
                        e
                    );
                }
            }
        })))
    }

    pub(crate) fn known_length(&self) -> Option<u64> {
        let mut length_out: u64 = 0;
        let status =
            unsafe { fastly_sys::fastly_http_body::known_length(self.as_u32(), &mut length_out) };
        if status == FastlyStatus::NONE {
            None
        } else {
            assert_eq!(status, FastlyStatus::OK);
            Some(length_out)
        }
    }
}

impl Read for BodyHandle {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        use std::io::{Error, ErrorKind};

        let mut nread = 0;
        let status = unsafe {
            abi::fastly_http_body::read(self.as_u32(), buf.as_mut_ptr(), buf.len(), &mut nread)
        };
        match status {
            FastlyStatus::OK => Ok(nread),
            FastlyStatus::HTTPINCOMPLETE => {
                Err(Error::new(ErrorKind::UnexpectedEof, "incomplete HTTP body"))
            }
            other => Err(Error::other(format!(
                "fastly_http_body::read failed with '{:?}'",
                other
            ))),
        }
    }
}

impl Write for BodyHandle {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        use std::io::{Error, ErrorKind};

        let mut nwritten = 0;
        let status = unsafe {
            abi::fastly_http_body::write(
                self.as_u32(),
                buf.as_ptr(),
                buf.len(),
                BodyWriteEnd::Back,
                &mut nwritten,
            )
        };
        match status {
            FastlyStatus::OK => Ok(nwritten),
            FastlyStatus::BADF => Err(Error::new(ErrorKind::InvalidInput, format!("{status:?}"))),
            _ => Err(Error::other(format!("{status:?}"))),
        }
    }

    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

impl Drop for BodyHandle {
    fn drop(&mut self) {
        if self.is_valid() {
            unsafe { abi::fastly_http_body::close(self.as_u32()) }
                .result()
                .expect("fastly_http_body::close failed");
        }
    }
}

impl From<&str> for BodyHandle {
    fn from(s: &str) -> Self {
        let mut handle = Self::new();
        handle
            .write_all(s.as_bytes())
            .expect("BodyHandle::from() write failed");
        handle
    }
}

impl From<String> for BodyHandle {
    fn from(s: String) -> Self {
        Self::from(s.as_str())
    }
}

impl From<&[u8]> for BodyHandle {
    fn from(s: &[u8]) -> Self {
        let mut handle = Self::new();
        handle
            .write_all(s)
            .expect("BodyHandle::from() write failed");
        handle
    }
}

impl From<Vec<u8>> for BodyHandle {
    fn from(s: Vec<u8>) -> Self {
        Self::from(s.as_slice())
    }
}

impl BodyHandleExt for BodyHandle {
    fn get_trailer_names<'a>(
        &'a self,
        buf_size: usize,
    ) -> Result<Box<dyn Iterator<Item = Result<HeaderName, BufferSizeError>> + 'a>, BodyHandleError>
    {
        Ok(self.get_trailer_names_impl(buf_size, Some(buf_size))?)
    }

    fn get_trailer_value(
        &self,
        name: &HeaderName,
        max_len: usize,
    ) -> Result<Result<Option<HeaderValue>, BufferSizeError>, BodyHandleError> {
        let name: &[u8] = name.as_ref();
        let mut buf = BytesMut::with_capacity(max_len);
        let mut nwritten = 0;
        let status = unsafe {
            abi::fastly_http_body::trailer_value_get(
                self.as_u32(),
                name.as_ptr(),
                name.len(),
                buf.as_mut_ptr(),
                buf.capacity(),
                &mut nwritten,
            )
        };
        match status.result().map(|_| nwritten) {
            Ok(nwritten) => {
                assert!(nwritten <= buf.capacity(), "hostcall wrote too many bytes");
                unsafe {
                    buf.set_len(nwritten);
                }
                let buf = buf.freeze();
                // we trust that the hostcall is giving us valid header bytes
                let value = unsafe { HeaderValue::from_maybe_shared_unchecked(buf) };
                Ok(Ok(Some(value)))
            }
            Err(FastlyStatus::AGAIN) => Err(BodyHandleError::TrailersNotReady),
            Err(FastlyStatus::INVAL) => Ok(Ok(None)),
            Err(FastlyStatus::BUFLEN) => Ok(Err(BufferSizeError::header_value(max_len, nwritten))),
            _ => panic!("fastly_http_resp::trailer_value_get returned error"),
        }
    }

    fn get_trailer_values<'a>(
        &'a self,
        name: &'a HeaderName,
        max_len: usize,
    ) -> Result<Box<dyn Iterator<Item = Result<HeaderValue, BufferSizeError>> + 'a>, BodyHandleError>
    {
        Ok(self.get_trailer_values_impl(name, max_len, Some(max_len))?)
    }
}