1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

use bytes::{Buf, Bytes, BytesMut};
pub use fastly_shared::{FastlyStatus, HttpVersion, FASTLY_ABI_VERSION};
pub use fastly_sys::*;

pub(crate) struct MultiValueHostcall<F> {
    fill_buf: F,
    term: u8,
    buf: BytesMut,
    max_buf_size: usize,
    cursor: u32,
    is_done: bool,
}

impl<F> MultiValueHostcall<F> {
    pub(crate) fn new(term: u8, max_buf_size: usize, fill_buf: F) -> Self {
        Self {
            fill_buf,
            term,
            buf: BytesMut::with_capacity(max_buf_size),
            max_buf_size,
            cursor: 0,
            is_done: false,
        }
    }
}

/// Errors related to a [`MultiValueHostcall`](struct.MultiValueHostcall.html).
///
/// Users do not directly interact with this error enum. It is most commonly used to propagate an
/// error to the user informing them that the buffer provided to a multi-value hostcall was not
/// sufficient.
///
/// See [`RequestHandle::get_header_names`][get_header_names] or
/// [`RequestHandle::get_header_values`][get_header_values] for examples of hostcalls that handle
/// this error.
///
/// [get_header_names]: ../request/struct.RequestHandle.html#method.get_header_names
/// [get_header_values]: ../request/struct.RequestHandle.html#method.get_header_values
#[derive(Copy, Clone, Debug, thiserror::Error, PartialEq, Eq)]
pub(crate) enum MultiValueHostcallError {
    /// The provided buffer size was too small.
    ///
    /// This error variant can be used to inform a user that they can try a multi-value hostcall
    /// again, using a larger buffer.
    #[error("MultiValueHostcall buffer too small")]
    BufferTooSmall,
    /// The [`FastlyStatus`][fastly-status] error code returned by the closure `fill_buf`.
    ///
    /// [fastly-status]: fastly_shared/struct.FastlyStatus.html
    #[error("MultiValueHostcall closure returned error: {0:?}")]
    ClosureError(FastlyStatus),
}

impl<F> std::iter::Iterator for MultiValueHostcall<F>
where
    F: Fn(*mut u8, usize, u32, *mut i64, *mut usize) -> FastlyStatus,
{
    type Item = Result<Bytes, MultiValueHostcallError>;

    fn next(&mut self) -> Option<Self::Item> {
        // first fill the buffer, if it's empty
        if self.buf.is_empty() {
            if self.is_done {
                // if there are no more calls to make, and the buffer is empty, we're done
                return None;
            }
            self.buf.reserve(self.max_buf_size);
            let mut ending_cursor = 0;
            let mut nwritten = 0;
            let status = (self.fill_buf)(
                self.buf.as_mut_ptr(),
                self.buf.capacity(),
                self.cursor,
                &mut ending_cursor,
                &mut nwritten,
            );
            if status.is_err() {
                self.is_done = true;
                return Some(Err(match status {
                    FastlyStatus::BUFLEN => MultiValueHostcallError::BufferTooSmall,
                    status => MultiValueHostcallError::ClosureError(status),
                }));
            }
            if nwritten == 0 {
                // if we get no bytes, we're definitely done; this only comes up if there are no
                // values at all, otherwise we see the ending cursor at -1 and stop
                self.is_done = true;
                return None;
            }
            assert!(
                nwritten <= self.buf.capacity(),
                "fill_buf set invalid nwritten: {}, capacity: {}",
                nwritten,
                self.buf.capacity()
            );
            unsafe {
                self.buf.set_len(nwritten);
            }
            if ending_cursor < 0 {
                // no more calls necessary after this one
                self.is_done = true;
            } else {
                assert!(
                    ending_cursor <= u32::MAX as i64 && ending_cursor > self.cursor as i64,
                    "fill_buf set invalid ending_cursor: {}, cursor: {}, nwritten: {}",
                    ending_cursor,
                    self.cursor,
                    nwritten
                );
                // otherwise adjust the cursor for the next fill
                self.cursor = ending_cursor as u32;
            }
        }
        // Find the index of the first terminator byte in the buffer, or panic. A missing
        // terminator violates the protocol of these hostcalls, which must always terminate each
        // element with the terminator byte.
        let first_term_ix = self
            .buf
            .iter()
            .position(|b| b == &self.term)
            .expect("terminator byte was not found");
        // split off the first element from the buffer
        let elt = self.buf.split_to(first_term_ix);
        // drop the terminator byte, which now remains in the buffer
        self.buf.advance(1);
        Some(Ok(elt.freeze()))
    }
}

impl<F> std::iter::FusedIterator for MultiValueHostcall<F> where
    F: Fn(*mut u8, usize, u32, *mut i64, *mut usize) -> FastlyStatus
{
}