fastly 0.2.0-alpha3

use crate::abi;
use crate::Error;
use fastly_shared::BodyWriteEnd;
use lazy_static::lazy_static;
use std::convert::TryFrom;
use std::io::{BufReader, BufWriter, Read, Write};
use std::sync::Mutex;

/// 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.
#[derive(Debug, Eq, Hash, PartialEq)]
#[repr(transparent)]
pub struct BodyHandle {
    handle: u32,
}

lazy_static! {
    pub(crate) static ref GOT_DOWNSTREAM: Mutex<bool> = Mutex::new(false);
}

impl BodyHandle {
    pub const INVALID: Self = Self {
        handle: fastly_shared::INVALID_BODY_HANDLE,
    };

    pub fn is_invalid(&self) -> bool {
        self.handle == Self::INVALID.handle
    }

    /// Get an owned `BodyHandle` from a borrowed one.
    ///
    /// This should only be used when calling the raw ABI directly.
    pub(crate) unsafe fn handle(&self) -> Self {
        Self {
            handle: self.handle,
        }
    }

    /// 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 downstream_request() -> Result<Self, Error> {
        let mut got = GOT_DOWNSTREAM.lock().unwrap();
        if *got {
            return Err(Error::msg(
                "cannot get more than one handle to the downstream request body per execution",
            ));
        }

        let mut handle = BodyHandle::INVALID;
        let status = unsafe { abi::xqd_req_body_downstream_get(std::ptr::null_mut(), &mut handle) };
        if status.is_err() || handle.is_invalid() {
            Err(Error::msg("xqd_req_body_downstream_get failed"))
        } else {
            *got = true;
            Ok(handle)
        }
    }

    /// Get a handle to a new, empty body.
    pub fn new() -> Result<BodyHandle, Error> {
        let mut handle = BodyHandle::INVALID;
        let status = unsafe { abi::xqd_body_new(&mut handle) };
        if status.is_err() || handle.is_invalid() {
            Err(Error::msg("xqd_body_new failed"))
        } else {
            Ok(handle)
        }
    }

    /// 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) -> Result<(), Error> {
        let status = unsafe { abi::xqd_body_append(self.handle(), other.handle()) };
        if status.is_err() {
            Err(Error::msg("xqd_body_append failed"))
        } else {
            Ok(())
        }
    }

    /// 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) -> Result<Vec<u8>, Error> {
        let mut body = vec![];
        let mut bufread = BufReader::new(self);
        bufread.read_to_end(&mut body)?;
        Ok(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) -> Result<String, Error> {
        let mut body = String::new();
        let mut bufread = BufReader::new(self);
        bufread.read_to_string(&mut body)?;
        Ok(body)
    }
}

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::xqd_body_read(self.handle(), buf.as_mut_ptr(), buf.len(), &mut nread) };
        if status.is_err() {
            Err(Error::new(ErrorKind::Other, "xqd_body_read failed"))
        } else {
            Ok(nread)
        }
    }
}

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::xqd_body_write(
                self.handle(),
                buf.as_ptr(),
                buf.len(),
                BodyWriteEnd::Back,
                &mut nwritten,
            )
        };
        if status.is_err() {
            Err(Error::new(ErrorKind::Other, "xqd_body_write failed"))
        } else {
            Ok(nwritten)
        }
    }

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

impl TryFrom<&str> for BodyHandle {
    type Error = Error;

    fn try_from(s: &str) -> Result<Self, Error> {
        let mut handle = BodyHandle::new()?;
        handle.write_all(s.as_bytes())?;
        Ok(handle)
    }
}

impl TryFrom<String> for BodyHandle {
    type Error = Error;

    fn try_from(s: String) -> Result<Self, Error> {
        BodyHandle::try_from(s.as_str())
    }
}

impl TryFrom<&[u8]> for BodyHandle {
    type Error = Error;

    fn try_from(s: &[u8]) -> Result<Self, Error> {
        let mut handle = BodyHandle::new()?;
        handle.write_all(s)?;
        Ok(handle)
    }
}

impl TryFrom<Vec<u8>> for BodyHandle {
    type Error = Error;

    fn try_from(s: Vec<u8>) -> Result<Self, Error> {
        BodyHandle::try_from(s.as_slice())
    }
}

/// A low-level interface to a streaming HTTP body.
///
/// The interface to this type is very similar to `BodyHandle`, however it is write-only, and can
/// only be created as a result of calling `ResponseHandle::send_downstream_streaming()`.
///
/// This type implements `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.
///
/// `BodyHandle`s can be appended to a streaming body in amortized constant time.
#[derive(Debug, Eq, Hash, PartialEq)]
#[repr(transparent)]
pub struct StreamingBodyHandle {
    handle: u32,
}

impl StreamingBodyHandle {
    pub(crate) unsafe fn from_body_handle(body_handle: &BodyHandle) -> Self {
        Self {
            handle: body_handle.handle,
        }
    }

    unsafe fn as_body_handle(&self) -> BodyHandle {
        BodyHandle {
            handle: self.handle,
        }
    }

    /// 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) -> Result<(), Error> {
        unsafe { self.as_body_handle() }.append(other)
    }
}

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

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

/// An HTTP body that can be read from, written to, or appended to another body.
///
/// The most efficient way to read and write the body is through the `Read`, `BufRead`, and `Write`
/// implementations.
pub struct Body {
    reader: BufReader<BodyHandle>,
    writer: BufWriter<BodyHandle>,
}

impl Body {
    pub fn new() -> Result<Self, Error> {
        Ok(BodyHandle::new()?.into())
    }

    // this is not exported, since misuse can lead to data getting dropped or appearing out of order
    fn handle(&mut self) -> &mut BodyHandle {
        self.reader.get_mut()
    }

    pub fn into_handle(mut self) -> Result<BodyHandle, Error> {
        // we need to make sure that any currently buffered data in either the reader or writer is
        // put (back) in the host body
        let read_buf = self.reader.buffer();
        if !read_buf.is_empty() {
            let mut nwritten = 0;
            let status = unsafe {
                abi::xqd_body_write(
                    self.reader.get_ref().handle(),
                    read_buf.as_ptr(),
                    read_buf.len(),
                    BodyWriteEnd::Front,
                    &mut nwritten,
                )
            };
            if status.is_err() {
                return Err(Error::msg("xqd_body_write_front failed"));
            }
            if nwritten != read_buf.len() {
                return Err(Error::msg("xqd_body_write_front didn't fully write"));
            }
        }
        self.writer.flush()?;
        Ok(self.reader.into_inner())
    }

    /// 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) -> Result<Vec<u8>, Error> {
        self.into_handle()?.into_bytes()
    }

    /// 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) -> Result<String, Error> {
        self.into_handle()?.into_string()
    }

    /// Append another body onto the end of this body.
    pub fn append(&mut self, other: Body) -> Result<(), Error> {
        // flush the write buffer of the destination body, so that we can use the append method on
        // the underlying handles
        self.writer.flush()?;
        self.handle().append(other.into_handle()?)
    }
}

// For these trait implementations we only implement the methods that the underlying buffered
// adaptors implement; the default implementations for the others will behave the same.
//
// The main bit of caution we must use here is that any read should be preceded by flushing the
// write buffer. `BufWriter` doesn't make any calls if its buffer is empty, so this isn't very
// expensive and could prevent unexpected results if a program is trying to read and write from the
// same body.
impl Read for Body {
    fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
        self.writer.flush()?;
        self.reader.read(buf)
    }

    fn read_vectored(&mut self, bufs: &mut [std::io::IoSliceMut]) -> std::io::Result<usize> {
        self.writer.flush()?;
        self.reader.read_vectored(bufs)
    }
}

impl std::io::BufRead for Body {
    fn fill_buf(&mut self) -> std::io::Result<&[u8]> {
        self.writer.flush()?;
        self.reader.fill_buf()
    }

    fn consume(&mut self, amt: usize) {
        self.reader.consume(amt)
    }
}

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

    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
        self.writer.write_vectored(bufs)
    }

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

impl From<BodyHandle> for Body {
    fn from(handle: BodyHandle) -> Self {
        // we clone the handle here in order to have an owned type for the reader and writer, but
        // this means we have to be careful that we don't make the aliasing observable from the
        // public interface
        let handle2 = unsafe { handle.handle() };
        Self {
            reader: BufReader::new(handle),
            writer: BufWriter::new(handle2),
        }
    }
}

impl TryFrom<&str> for Body {
    type Error = Error;

    fn try_from(s: &str) -> Result<Self, Error> {
        Ok(BodyHandle::try_from(s)?.into())
    }
}

impl TryFrom<String> for Body {
    type Error = Error;

    fn try_from(s: String) -> Result<Self, Error> {
        Ok(BodyHandle::try_from(s)?.into())
    }
}

impl TryFrom<&[u8]> for Body {
    type Error = Error;

    fn try_from(s: &[u8]) -> Result<Self, Error> {
        Ok(BodyHandle::try_from(s)?.into())
    }
}

impl TryFrom<Vec<u8>> for Body {
    type Error = Error;

    fn try_from(s: Vec<u8>) -> Result<Self, Error> {
        Ok(BodyHandle::try_from(s)?.into())
    }
}

/// A streaming HTTP body that can be written to, or appended to from another body.
///
/// The interface to this type is very similar to `Body`, however it is write-only, and can only be
/// created as a result of calling `ResponseExt::send_downstream_streaming()`.
///
/// The most efficient way to write the body is through the `Write` implementation. Writes are
/// buffered, but calling `Write::flush()` will flush the buffer and cause a new chunk to be written
/// to the downstream client.
pub struct StreamingBody {
    writer: BufWriter<StreamingBodyHandle>,
}

impl StreamingBody {
    // this is not exported, since misuse can lead to data getting dropped or appearing out of order
    fn handle(&mut self) -> &mut StreamingBodyHandle {
        self.writer.get_mut()
    }

    /// Append a body onto the end of this streaming body.
    pub fn append(&mut self, other: Body) -> Result<(), Error> {
        // flush the write buffer of the destination body, so that we can use the append method on
        // the underlying handles
        self.writer.flush()?;
        self.handle().append(other.into_handle()?)
    }
}

impl From<StreamingBodyHandle> for StreamingBody {
    fn from(handle: StreamingBodyHandle) -> Self {
        Self {
            writer: BufWriter::new(handle),
        }
    }
}

// This trait implementation is much simpler than those of `Body`, since we don't have to manage
// multiple buffers. It's just a passthrough to the methods defined on `BufWriter`.
impl Write for StreamingBody {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.writer.write(buf)
    }

    fn write_vectored(&mut self, bufs: &[std::io::IoSlice<'_>]) -> std::io::Result<usize> {
        self.writer.write_vectored(bufs)
    }

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