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
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
use std::io::{self, BufRead, Read, Write, Cursor, BufReader};
use std::path::Path;
use std::fs::File;
use std::time::Duration;
use std::mem::transmute;

use super::data_stream::{DataStream, StreamReader, kill_stream};

use ext::ReadExt;

use http::hyper::h1::HttpReader;
use http::hyper::buffer;
use http::hyper::h1::HttpReader::*;
use http::hyper::net::{HttpStream, NetworkStream};

pub type BodyReader<'a, 'b> =
    self::HttpReader<&'a mut self::buffer::BufReader<&'b mut NetworkStream>>;

/// Type representing the data in the body of an incoming request.
///
/// This type is the only means by which the body of a request can be retrieved.
/// This type is not usually used directly. Instead, types that implement
/// [FromData](/rocket/data/trait.FromData.html) are used via code generation by
/// specifying the `data = "<param>"` route parameter as follows:
///
/// ```rust,ignore
/// #[post("/submit", data = "<var>")]
/// fn submit(var: T) -> ... { ... }
/// ```
///
/// Above, `T` can be any type that implements `FromData`. Note that `Data`
/// itself implements `FromData`.
///
/// # Reading Data
///
/// Data may be read from a `Data` object by calling either the
/// [open](#method.open) or [peek](#method.peek) methods.
///
/// The `open` method consumes the `Data` object and returns the raw data
/// stream. The `Data` object is consumed for safety reasons: consuming the
/// object ensures that holding a `Data` object means that all of the data is
/// available for reading.
///
/// The `peek` method returns a slice containing at most 4096 bytes of buffered
/// body data. This enables partially or fully reading from a `Data` object
/// without consuming the `Data` object.
pub struct Data {
    buffer: Vec<u8>,
    is_done: bool,
    // TODO: This sucks as it depends on a TCPStream. Oh, hyper.
    stream: StreamReader,
    // Ideally we wouldn't have these, but Hyper forces us to.
    position: usize,
    capacity: usize,
}

impl Data {
    /// Returns the raw data stream.
    ///
    /// The stream contains all of the data in the body of the request,
    /// including that in the `peek` buffer. The method consumes the `Data`
    /// instance. This ensures that a `Data` type _always_ represents _all_ of
    /// the data in a request.
    pub fn open(mut self) -> impl BufRead {
        // Swap out the buffer and stream for empty ones so we can move.
        let mut buffer = vec![];
        let mut stream = EmptyReader(self.stream.get_ref().clone());
        ::std::mem::swap(&mut buffer, &mut self.buffer);
        ::std::mem::swap(&mut stream, &mut self.stream);

        // Setup the underlying reader at the correct pointers.
        let mut cursor = Cursor::new(buffer);
        cursor.set_position(self.position as u64);
        let buffered = cursor.take((self.capacity - self.position) as u64);

        // Create the actual DataSteam.
        DataStream {
            network: stream.get_ref().clone(),
            stream: buffered.chain(BufReader::new(stream)),
        }
    }

    #[doc(hidden)]
    pub fn from_hyp(mut h_body: BodyReader) -> Result<Data, &'static str> {
        // FIXME: This is asolutely terrible, thanks to Hyper.

        // Retrieve the underlying HTTPStream from Hyper.
        let mut stream = match h_body.get_ref().get_ref()
                                     .downcast_ref::<HttpStream>() {
            Some(s) => {
                let owned_stream = s.clone();
                let buf_len = h_body.get_ref().get_buf().len() as u64;
                match h_body {
                    SizedReader(_, n) => SizedReader(owned_stream, n - buf_len),
                    EofReader(_) => EofReader(owned_stream),
                    EmptyReader(_) => EmptyReader(owned_stream),
                    ChunkedReader(_, n) =>
                        ChunkedReader(owned_stream, n.map(|k| k - buf_len)),
                }
            },
            None => return Err("Stream is not an HTTP stream!"),
        };

        // Set the read timeout to 5 seconds.
        stream.get_mut().set_read_timeout(Some(Duration::from_secs(5))).unwrap();

        // Create the Data object from hyper's buffer.
        let (vec, pos, cap) = h_body.get_mut().take_buf();
        Ok(Data::new(vec, pos, cap, stream))
    }

    /// Retrieve the `peek` buffer.
    ///
    /// The peek buffer contains at most 4096 bytes of the body of the request.
    /// The actual size of the returned buffer varies by web request. The
    /// [peek_complete](#method.peek_complete) can be used to determine if this
    /// buffer contains _all_ of the data in the body of the request.
    #[inline(always)]
    pub fn peek(&self) -> &[u8] {
        &self.buffer[self.position..self.capacity]
    }

    /// Returns true if the `peek` buffer contains all of the data in the body
    /// of the request. Returns `false` if it does not or if it is not known if
    /// it does.
    #[inline(always)]
    pub fn peek_complete(&self) -> bool {
        self.is_done
    }

    /// A helper method to write the body of the request to any `Write` type.
    ///
    /// This method is identical to `io::copy(&mut data.open(), writer)`.
    #[inline(always)]
    pub fn stream_to<W: Write>(self, writer: &mut W) -> io::Result<u64> {
        io::copy(&mut self.open(), writer)
    }

    /// A helper method to write the body of the request to a file at the path
    /// determined by `path`.
    ///
    /// This method is identical to
    /// `io::copy(&mut self.open(), &mut File::create(path)?)`.
    #[inline(always)]
    pub fn stream_to_file<P: AsRef<Path>>(self, path: P) -> io::Result<u64> {
        io::copy(&mut self.open(), &mut File::create(path)?)
    }

    // Creates a new data object with an internal buffer `buf`, where the cursor
    // in the buffer is at `pos` and the buffer has `cap` valid bytes. The
    // remainder of the data bytes can be read from `stream`.
    #[doc(hidden)]
    pub fn new(mut buf: Vec<u8>,
               pos: usize,
               mut cap: usize,
               mut stream: StreamReader)
               -> Data {
        // Make sure the buffer is large enough for the bytes we want to peek.
        const PEEK_BYTES: usize = 4096;
        if buf.len() < PEEK_BYTES {
            trace_!("Resizing peek buffer from {} to {}.", buf.len(), PEEK_BYTES);
            buf.resize(PEEK_BYTES, 0);
        }

        // Fill the buffer with as many bytes as possible. If we read less than
        // that buffer's length, we know we reached the EOF. Otherwise, it's
        // unclear, so we just say we didn't reach EOF.
        trace!("Init buffer cap: {}", cap);
        let eof = match stream.read_max(&mut buf[cap..]) {
            Ok(n) => {
                trace_!("Filled peek buf with {} bytes.", n);
                cap += n;
                cap < buf.len()
            }
            Err(e) => {
                error_!("Failed to read into peek buffer: {:?}.", e);
                false
            },
        };

        trace_!("Peek buffer size: {}, remaining: {}", buf.len(), buf.len() - cap);
        Data {
            buffer: buf,
            stream: stream,
            is_done: eof,
            position: pos,
            capacity: cap,
        }
    }
}

impl Drop for Data {
    fn drop(&mut self) {
        // This is okay since the network stream expects to be shared mutably.
        unsafe {
            let stream: &mut StreamReader = transmute(self.stream.by_ref());
            kill_stream(stream, self.stream.get_mut());
        }
    }
}