fastly 0.6.0-beta1

Fastly Compute@Edge API
Documentation
//! HTTP bodies.

use crate::abi;
use fastly_shared::BodyWriteEnd;
use std::io::{BufReader, Read, Write};
use std::sync::atomic::{AtomicBool, Ordering};

/// A low-level interface to HTTP bodies.
///
/// 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`] type.
///
/// These handles can also be used to append bodies to each other in amortized constant time.
///
/// [`Body`]: struct.Body.html
/// [`Read`]: https://doc.rust-lang.org/std/io/trait.Read.html
/// [`Write`]: https://doc.rust-lang.org/std/io/trait.Write.html
#[derive(Debug, Eq, Hash, PartialEq)]
#[repr(transparent)]
pub struct BodyHandle {
    pub(super) 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.
    pub const INVALID: Self = Self {
        handle: fastly_shared::INVALID_BODY_HANDLE,
    };

    /// Return true if the body handle is valid.
    pub fn is_valid(&self) -> bool {
        !self.is_invalid()
    }

    /// Return true if the body handle is invalid.
    pub fn is_invalid(&self) -> bool {
        self.handle == Self::INVALID.handle
    }

    /// Make a handle from its underlying representation.
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    pub(crate) fn from_u32(handle: u32) -> Self {
        Self { handle }
    }

    /// Get the underlying representation of the handle.
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    pub(crate) unsafe fn as_u32(&self) -> u32 {
        self.handle
    }

    /// Get a mutable reference to the underlying `u32` representation of the handle.
    ///
    /// This should only be used when calling the raw ABI directly, and care should be taken not to
    /// reuse or alias handle values.
    pub(crate) fn as_u32_mut(&mut self) -> &mut u32 {
        &mut self.handle
    }

    /// 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.
    // FIXME KTM 2020-05-08: This is, the *worst* name.
    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");
        }
    }

    #[deprecated(since = "0.6.0", note = "renamed to `BodyHandle::from_client()`")]
    pub fn downstream_request() -> Self {
        Self::from_client()
    }

    /// Get a handle to the downstream request body.
    ///
    /// This handle may only be retrieved once per execution, either through this function or
    /// through `fastly::request::downstream_request_and_body_handles()`.
    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 [`BodyHandle`](struct.BodyHandle.html).
    #[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.
    ///
    /// 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.as_u32()) }
            .result()
            .expect("fastly_http_body::append failed")
    }

    /// Read the entirety of the body.
    ///
    /// Note that this will involve copying and buffering the body, and so should only be used for
    /// convenience on small request bodies.
    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.
    ///
    /// Note that this will involve copying and buffering the body, and so should only be used for
    /// convenience on small request bodies.
    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
    }

    /// Write a slice of bytes to the end of this [`BodyHandle`](struct.BodyHandle.html).
    ///
    /// This function returns the number of bytes written.
    ///
    /// ```rust,no_run
    /// # use fastly::handle::BodyHandle;
    /// # let mut body = BodyHandle::new();
    /// body.write_bytes(&[0, 1, 2, 3]);
    /// ```
    pub fn write_bytes(&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::Back,
                &mut nwritten,
            )
        };
        status
            .result()
            .map(|_| nwritten)
            .expect("fastly_http_body::write failed")
    }

    /// Write a string slice to the end of this [`BodyHandle`](struct.BodyHandle.html).
    ///
    /// This function returns the number of bytes written.
    ///
    /// ```rust,no_run
    /// # use fastly::handle::BodyHandle;
    /// # let mut body = BodyHandle::new();
    /// body.write_str("woof woof");
    /// ```
    pub fn write_str(&mut self, string: &str) -> usize {
        self.write_bytes(string.as_bytes())
    }
}

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)
        };
        if status.is_err() {
            Err(Error::new(
                ErrorKind::Other,
                "fastly_http_body::read failed",
            ))
        } else {
            Ok(nread)
        }
    }
}

impl Write for BodyHandle {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        Ok(self.write_bytes(buf))
    }

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

impl From<&str> for BodyHandle {
    fn from(s: &str) -> Self {
        let mut handle = Self::new();
        handle.write_str(s);
        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_bytes(s);
        handle
    }
}

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