refactor: multipart tweaks

actix-multipart/src/payload.rs

@@ -1,6 +1,6 @@
 use std::{
     cell::{RefCell, RefMut},
-    cmp,
+    cmp, mem,
     pin::Pin,
     rc::Rc,
     task::{Context, Poll},
@@ -12,7 +12,7 @@ use actix_web::{
 };
 use futures_core::stream::{LocalBoxStream, Stream};
 
-use crate::{error::MultipartError, safety::Safety};
+use crate::{error::Error, safety::Safety};
 
 pub(crate) struct PayloadRef {
     payload: Rc<RefCell<PayloadBuffer>>,
@@ -21,7 +21,7 @@ pub(crate) struct PayloadRef {
 impl PayloadRef {
     pub(crate) fn new(payload: PayloadBuffer) -> PayloadRef {
         PayloadRef {
-            payload: Rc::new(payload.into()),
+            payload: Rc::new(RefCell::new(payload)),
         }
     }
 
@@ -44,28 +44,33 @@ impl Clone for PayloadRef {
 
 /// Payload buffer.
 pub(crate) struct PayloadBuffer {
-    pub(crate) eof: bool,
-    pub(crate) buf: BytesMut,
     pub(crate) stream: LocalBoxStream<'static, Result<Bytes, PayloadError>>,
+    pub(crate) buf: BytesMut,
+    /// EOF flag. If true, no more payload reads will be attempted.
+    pub(crate) eof: bool,
 }
 
 impl PayloadBuffer {
-    /// Constructs new `PayloadBuffer` instance.
+    /// Constructs new payload buffer.
     pub(crate) fn new<S>(stream: S) -> Self
     where
         S: Stream<Item = Result<Bytes, PayloadError>> + 'static,
     {
         PayloadBuffer {
-            eof: false,
-            buf: BytesMut::new(),
             stream: Box::pin(stream),
+            buf: BytesMut::with_capacity(1_024), // pre-allocate 1KiB
+            eof: false,
         }
     }
 
     pub(crate) fn poll_stream(&mut self, cx: &mut Context<'_>) -> Result<(), PayloadError> {
         loop {
             match Pin::new(&mut self.stream).poll_next(cx) {
-                Poll::Ready(Some(Ok(data))) => self.buf.extend_from_slice(&data),
+                Poll::Ready(Some(Ok(data))) => {
+                    self.buf.extend_from_slice(&data);
+                    // try to read more data
+                    continue;
+                }
                 Poll::Ready(Some(Err(err))) => return Err(err),
                 Poll::Ready(None) => {
                     self.eof = true;
@@ -76,7 +81,7 @@ impl PayloadBuffer {
         }
     }
 
-    /// Read exact number of bytes.
+    /// Reads exact number of bytes.
     #[cfg(test)]
     pub(crate) fn read_exact(&mut self, size: usize) -> Option<Bytes> {
         if size <= self.buf.len() {
@@ -86,46 +91,57 @@ impl PayloadBuffer {
         }
     }
 
-    pub(crate) fn read_max(&mut self, size: u64) -> Result<Option<Bytes>, MultipartError> {
+    pub(crate) fn read_max(&mut self, size: u64) -> Result<Option<Bytes>, Error> {
         if !self.buf.is_empty() {
             let size = cmp::min(self.buf.len() as u64, size) as usize;
             Ok(Some(self.buf.split_to(size).freeze()))
         } else if self.eof {
-            Err(MultipartError::Incomplete)
+            Err(Error::Incomplete)
         } else {
             Ok(None)
         }
     }
 
-    /// Read until specified ending.
-    pub(crate) fn read_until(&mut self, line: &[u8]) -> Result<Option<Bytes>, MultipartError> {
-        let res = memchr::memmem::find(&self.buf, line)
-            .map(|idx| self.buf.split_to(idx + line.len()).freeze());
+    /// Reads until specified ending.
+    ///
+    /// Returns:
+    ///
+    /// - `Ok(Some(chunk))` - `needle` is found, with chunk ending after needle
+    /// - `Err(Incomplete)` - `needle` is not found and we're at EOF
+    /// - `Ok(None)` - `needle` is not found otherwise
+    pub(crate) fn read_until(&mut self, needle: &[u8]) -> Result<Option<Bytes>, Error> {
+        match memchr::memmem::find(&self.buf, needle) {
+            // buffer exhausted and EOF without finding needle
+            None if self.eof => Err(Error::Incomplete),
 
-        if res.is_none() && self.eof {
-            Err(MultipartError::Incomplete)
-        } else {
-            Ok(res)
+            // needle not yet found
+            None => Ok(None),
+
+            // needle found, split chunk out of buf
+            Some(idx) => Ok(Some(self.buf.split_to(idx + needle.len()).freeze())),
         }
     }
 
-    /// Read bytes until new line delimiter.
-    pub(crate) fn readline(&mut self) -> Result<Option<Bytes>, MultipartError> {
+    /// Reads bytes until new line delimiter.
+    #[inline]
+    pub(crate) fn readline(&mut self) -> Result<Option<Bytes>, Error> {
         self.read_until(b"\n")
     }
 
-    /// Read bytes until new line delimiter or EOF.
-    pub(crate) fn readline_or_eof(&mut self) -> Result<Option<Bytes>, MultipartError> {
+    /// Reads bytes until new line delimiter or until EOF.
+    #[inline]
+    pub(crate) fn readline_or_eof(&mut self) -> Result<Option<Bytes>, Error> {
         match self.readline() {
-            Err(MultipartError::Incomplete) if self.eof => Ok(Some(self.buf.split().freeze())),
+            Err(Error::Incomplete) if self.eof => Ok(Some(self.buf.split().freeze())),
             line => line,
         }
     }
 
-    /// Put unprocessed data back to the buffer.
+    /// Puts unprocessed data back to the buffer.
     pub(crate) fn unprocessed(&mut self, data: Bytes) {
-        let buf = BytesMut::from(data.as_ref());
-        let buf = std::mem::replace(&mut self.buf, buf);
+        // TODO: use BytesMut::from when it's released, see https://github.com/tokio-rs/bytes/pull/710
+        let buf = BytesMut::from(&data[..]);
+        let buf = mem::replace(&mut self.buf, buf);
         self.buf.extend_from_slice(&buf);
     }
 }