fastly 0.2.0-alpha2

use crate::abi::{self, XqdStatus};
use crate::body::{Body, BodyHandle, StreamingBody, StreamingBodyHandle};
use crate::error::{anyhow, Error};
use bytes::{BufMut, BytesMut};
use http::header::{HeaderName, HeaderValue};
use http::{Response, StatusCode, Version};
use lazy_static::lazy_static;
use std::convert::TryFrom;
use std::io::Write;
use std::sync::Mutex;

#[derive(Debug, Eq, Hash, PartialEq)]
#[repr(transparent)]
pub struct ResponseHandle {
    pub(crate) handle: u32,
}

impl ResponseHandle {
    pub const INVALID: Self = ResponseHandle {
        handle: fastly_shared::INVALID_RESPONSE_HANDLE,
    };

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

    /// Get an owned `ResponseHandle` 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,
        }
    }

    pub fn new() -> Result<Self, Error> {
        let mut handle = ResponseHandle::INVALID;
        let status = unsafe { abi::xqd_resp_new(&mut handle) };
        if status.is_err() || handle.is_invalid() {
            Err(Error::msg("xqd_resp_new failed"))
        } else {
            Ok(handle)
        }
    }

    pub fn get_header_names<'a>(
        &'a self,
        max_len: usize,
    ) -> impl Iterator<Item = Result<HeaderName, Error>> + 'a {
        abi::MultiValueHostcall::new(
            b'\0',
            max_len,
            move |buf, buf_size, cursor, ending_cursor, nwritten| unsafe {
                abi::xqd_resp_header_names_get(
                    self.handle(),
                    buf,
                    buf_size,
                    cursor,
                    ending_cursor,
                    nwritten,
                )
            },
        )
        .map(|res| {
            res.and_then(|name_bytes| {
                HeaderName::from_bytes(&name_bytes)
                    .map_err(|e| anyhow!("invalid header name: {}", e))
            })
        })
    }

    pub fn get_header_values<'a>(
        &'a self,
        name: &'a HeaderName,
        max_len: usize,
    ) -> impl Iterator<Item = Result<HeaderValue, Error>> + 'a {
        abi::MultiValueHostcall::new(
            b'\0',
            max_len,
            move |buf, buf_size, cursor, ending_cursor, nwritten| unsafe {
                let name: &[u8] = name.as_ref();
                abi::xqd_resp_header_values_get(
                    self.handle(),
                    name.as_ptr(),
                    name.len(),
                    buf,
                    buf_size,
                    cursor,
                    ending_cursor,
                    nwritten,
                )
            },
        )
        .map(|res| {
            res.map(|value_bytes| unsafe {
                // we trust that the hostcall is giving us valid header bytes
                HeaderValue::from_maybe_shared_unchecked(value_bytes)
            })
        })
    }

    pub fn set_header_values<'a, I>(&mut self, name: &HeaderName, values: I) -> Result<(), Error>
    where
        I: IntoIterator<Item = &'a HeaderValue>,
    {
        // build a buffer of all the values, each terminated by a nul byte
        let mut buf = vec![];
        for value in values {
            buf.put(value.as_bytes());
            buf.put_u8(b'\0');
        }

        let name: &[u8] = name.as_ref();
        let status = unsafe {
            abi::xqd_resp_header_values_set(
                self.handle(),
                name.as_ptr(),
                name.len(),
                buf.as_ptr(),
                buf.len(),
            )
        };

        if status.is_err() {
            Err(Error::msg("xqd_req_header_values_set failed"))
        } else {
            Ok(())
        }
    }

    pub fn get_header_value(
        &self,
        name: &HeaderName,
        max_len: usize,
    ) -> Result<HeaderValue, Error> {
        let name: &[u8] = name.as_ref();
        let mut buf = BytesMut::with_capacity(max_len);
        let mut nwritten = 0;
        let status = unsafe {
            abi::xqd_resp_header_value_get(
                self.handle(),
                name.as_ptr(),
                name.len(),
                buf.as_mut_ptr(),
                buf.capacity(),
                &mut nwritten,
            )
        };
        if status.is_err() {
            return Err(Error::msg("xqd_resp_header_value_get returned error"));
        }
        assert!(nwritten <= buf.capacity(), "hostcall wrote too many bytes");
        unsafe {
            buf.set_len(nwritten);
        }
        Ok(HeaderValue::from_bytes(&buf).map_err(|_| Error::msg("invalid header"))?)
    }

    pub fn insert_header(&mut self, name: &HeaderName, value: &HeaderValue) -> Result<(), Error> {
        let name_bytes: &[u8] = name.as_ref();
        let value_bytes: &[u8] = value.as_ref();
        let status = unsafe {
            abi::xqd_resp_header_insert(
                self.handle(),
                name_bytes.as_ptr(),
                name_bytes.len(),
                value_bytes.as_ptr(),
                value_bytes.len(),
            )
        };
        if status.is_err() {
            Err(Error::msg("xqd_resp_header_insert returned error"))
        } else {
            Ok(())
        }
    }

    pub fn append_header(&mut self, name: &HeaderName, value: &HeaderValue) -> Result<(), Error> {
        let name_bytes: &[u8] = name.as_ref();
        let value_bytes: &[u8] = value.as_ref();
        let status = unsafe {
            abi::xqd_resp_header_append(
                self.handle(),
                name_bytes.as_ptr(),
                name_bytes.len(),
                value_bytes.as_ptr(),
                value_bytes.len(),
            )
        };
        if status.is_err() {
            Err(Error::msg("xqd_resp_header_append returned error"))
        } else {
            Ok(())
        }
    }

    pub fn set_status(&mut self, status: StatusCode) {
        let status = unsafe { abi::xqd_resp_status_set(self.handle(), status.as_u16()) };
        assert_eq!(
            status,
            XqdStatus::OK,
            "setting a StatusCode should always succeed"
        );
    }

    pub fn get_status(&self) -> Result<StatusCode, Error> {
        let mut status = 0;
        let xqd_status = unsafe { abi::xqd_resp_status_get(self.handle(), &mut status) };
        if xqd_status.is_err() {
            Err(Error::msg("xqd_resp_status_get failed"))
        } else {
            StatusCode::from_u16(status).map_err(|e| anyhow!("invalid status code: {}", e))
        }
    }

    pub fn get_version(&self) -> Result<Version, Error> {
        let mut version = 0;
        let status = unsafe { abi::xqd_resp_version_get(self.handle(), &mut version) };
        if status.is_err() {
            Err(Error::msg("xqd_resp_version_get failed"))
        } else {
            abi::HttpVersion::try_from(version)
                .map(Into::into)
                .map_err(Error::msg)
        }
    }

    pub fn set_version(&mut self, v: Version) -> Result<(), Error> {
        let status =
            unsafe { abi::xqd_resp_version_set(self.handle(), abi::HttpVersion::from(v) as u32) };
        if status.is_err() {
            Err(Error::msg("xqd_req_version_get failed"))
        } else {
            Ok(())
        }
    }

    /// Send this response downstream to the client with the given body.
    ///
    /// This will only begin sending bytes once the program has completed execution. If you want to
    /// begin sending before the program is complete, use
    /// `ResponseHandle::send_downstream_streaming()`.
    pub fn send_downstream(self, body: BodyHandle) -> Result<(), Error> {
        let status = unsafe { abi::xqd_resp_send_downstream(self, body, false as u32) };
        if status.is_err() {
            Err(Error::msg("xqd_resp_send_downstream failed"))
        } else {
            Ok(())
        }
    }

    /// Immediately begin sending this response downstream to the client, and return a
    /// `StreamingBodyHandle` that can accept further data to send.
    pub fn send_downstream_streaming(self, body: BodyHandle) -> Result<StreamingBodyHandle, Error> {
        let streaming_body_handle = unsafe { StreamingBodyHandle::from_body_handle(&body) };
        let status = unsafe { abi::xqd_resp_send_downstream(self, body, true as u32) };
        if status.is_err() {
            Err(Error::msg("xqd_resp_send_downstream failed"))
        } else {
            Ok(streaming_body_handle)
        }
    }
}

pub(crate) fn handles_to_response(
    resp_handle: ResponseHandle,
    resp_body_handle: BodyHandle,
) -> Result<Response<Body>, Error> {
    let mut resp = Response::builder()
        .status(resp_handle.get_status()?)
        .version(resp_handle.get_version()?);

    for name in resp_handle.get_header_names(crate::HEADER_NAME_MAX_LEN) {
        let name = name?;
        for value in resp_handle.get_header_values(&name, crate::HEADER_VALUE_MAX_LEN) {
            resp = resp.header(&name, value?);
        }
    }

    Ok(resp.body(resp_body_handle.into())?)
}

pub trait ResponseExt {
    fn send_downstream(self) -> Result<(), Error>;

    fn send_downstream_streaming(self) -> Result<StreamingBody, Error>;

    /// Replace the body of a response with a `Body` with the same contents.
    fn inner_to_body(self) -> Result<Response<Body>, Error>;

    /// Replace the body of a response with the remaining contents of its body.
    ///
    /// Note that this will involve copying and buffering the body, and so should only be used for
    /// convenience on small response bodies.
    fn inner_to_bytes(self) -> Result<Response<Vec<u8>>, Error>;
}

fn send_downstream_impl(
    resp: Response<Body>,
    streaming: bool,
) -> Result<Option<StreamingBody>, Error> {
    lazy_static! {
        static ref SENT: Mutex<bool> = Mutex::new(false);
    }

    let mut sent = SENT.lock().unwrap();
    if *sent {
        return Err(Error::msg(
            "cannot send more than one downstream response per execution",
        ));
    }

    let (parts, body) = resp.into_parts();

    let mut resp_handle = ResponseHandle::new()?;

    for name in parts.headers.keys() {
        resp_handle.set_header_values(name, parts.headers.get_all(name))?;
    }

    resp_handle.set_status(parts.status);
    resp_handle.set_version(parts.version)?;

    let res = if streaming {
        Some(
            resp_handle
                .send_downstream_streaming(body.into_handle()?)?
                .into(),
        )
    } else {
        resp_handle.send_downstream(body.into_handle()?)?;
        None
    };

    *sent = true;
    Ok(res)
}

impl ResponseExt for Response<Body> {
    /// Send this response downstream to the client.
    ///
    /// This will only begin sending bytes once the program has completed execution. If you want to
    /// begin sending before the program is complete, use
    /// `ResponseExt::send_downstream_streaming()`.
    fn send_downstream(self) -> Result<(), Error> {
        let res = send_downstream_impl(self, false)?;
        debug_assert!(res.is_none());
        Ok(())
    }

    /// Immediately begin sending this response downstream to the client, and return a
    /// `StreamingBody` that can accept further data to send.
    fn send_downstream_streaming(self) -> Result<StreamingBody, Error> {
        let res = send_downstream_impl(self, true)?;
        // streaming = true means we always get back a `Some`
        Ok(res.expect("streaming body is present"))
    }

    fn inner_to_body(self) -> Result<Response<Body>, Error> {
        Ok(self)
    }

    fn inner_to_bytes(self) -> Result<Response<Vec<u8>>, Error> {
        let (parts, body) = self.into_parts();
        Ok(Response::from_parts(parts, body.into_bytes()?))
    }
}

impl<T: AsRef<[u8]>> ResponseExt for Response<T> {
    fn send_downstream(self) -> Result<(), Error> {
        let mut body = Body::new()?;
        body.write_all(self.body().as_ref())?;
        self.map(|_| body).send_downstream()
    }

    fn send_downstream_streaming(self) -> Result<StreamingBody, Error> {
        let mut body = Body::new()?;
        body.write_all(self.body().as_ref())?;
        self.map(|_| body).send_downstream_streaming()
    }

    fn inner_to_body(self) -> Result<Response<Body>, Error> {
        let mut body = Body::new()?;
        body.write_all(self.body().as_ref())?;
        Ok(self.map(|_| body))
    }

    fn inner_to_bytes(self) -> Result<Response<Vec<u8>>, Error> {
        Ok(self.map(|b| b.as_ref().to_vec()))
    }
}