diff --git a/src/rust/cmn.rs b/src/rust/cmn.rs index 4d2158139a0..62040844a56 100644 --- a/src/rust/cmn.rs +++ b/src/rust/cmn.rs @@ -129,6 +129,7 @@ const BOUNDARY: &'static str = "MDuXWGyeE33QFXGchb2VFWc4Z7945d"; pub struct MultiPartReader<'a> { raw_parts: Vec<(Headers, &'a mut Read)>, current_part: Option<(Cursor>, &'a mut Read)>, + last_part_boundary: Option>>, } impl<'a> MultiPartReader<'a> { @@ -163,16 +164,33 @@ impl<'a> MultiPartReader<'a> { vec![(Attr::Ext("boundary".to_string()), Value::Ext(BOUNDARY.to_string()))], ) } + + /// Returns true if we are totally used + fn is_depleted(&self) -> bool { + self.raw_parts.len() == 0 && self.current_part.is_none() + } + + /// Returns true if we are handling our last part + fn is_last_part(&self) -> bool { + self.raw_parts.len() == 0 && self.current_part.is_some() + } } impl<'a> Read for MultiPartReader<'a> { fn read(&mut self, buf: &mut [u8]) -> io::Result { - if self.raw_parts.len() == 0 && self.current_part.is_none() { + if self.last_part_boundary.is_some() { + let br = self.last_part_boundary.as_mut().unwrap().read(buf).unwrap_or(0); + if br < buf.len() { + self.last_part_boundary = None; + } + return Ok(br) + } else if self.is_depleted() { return Ok(0) } else if self.raw_parts.len() > 0 && self.current_part.is_none() { let (headers, reader) = self.raw_parts.remove(0); let mut c = Cursor::new(Vec::::new()); - write!(&mut c, "{}{}", headers, LINE_ENDING).unwrap(); + write!(&mut c, "{}--{}{}{}{}", LINE_ENDING, BOUNDARY, LINE_ENDING, + headers, LINE_ENDING).unwrap(); c.seek(SeekFrom::Start(0)).unwrap(); self.current_part = Some((c, reader)); } @@ -186,10 +204,23 @@ impl<'a> Read for MultiPartReader<'a> { match rr { Ok(bytes_read) => { if bytes_read == 0 { + if self.is_last_part() { + // before clearing the last part, we will add the boundary that + // will be written last + self.last_part_boundary = Some(Cursor::new( + format!("{}--{}", LINE_ENDING, BOUNDARY).into_bytes())) + } // We are depleted - this can trigger the next part to come in self.current_part = None; } - Ok(hb + bytes_read) + let mut total_bytes_read = hb + bytes_read; + while total_bytes_read < buf.len() && !self.is_depleted() { + match self.read(&mut buf[total_bytes_read ..]) { + Ok(br) => total_bytes_read += br, + Err(err) => return Err(err), + } + } + Ok(total_bytes_read) } Err(err) => { // fail permanently diff --git a/src/rust/lib.rs b/src/rust/lib.rs index a0ed54e18cb..df9c36e3755 100644 --- a/src/rust/lib.rs +++ b/src/rust/lib.rs @@ -29,8 +29,27 @@ mod tests { .add_part(&mut r2, 25, &"application/plain".parse().unwrap()); let mut res = String::new(); - assert_eq!(mpr.read_to_string(&mut res), Ok(57)); + let r = mpr.read_to_string(&mut res); + assert_eq!(res.len(), r.clone().unwrap()); println!("{}", res); - assert!(false); + + let expected = +"\r\n--MDuXWGyeE33QFXGchb2VFWc4Z7945d\r\n\ +Content-Length: 50\r\n\ +Content-Type: application/json\r\n\ +\r\n\ +foo\r\n\ +--MDuXWGyeE33QFXGchb2VFWc4Z7945d\r\n\ +Content-Length: 25\r\n\ +Content-Type: application/plain\r\n\ +\r\n\ +bar\r\n\ +--MDuXWGyeE33QFXGchb2VFWc4Z7945d"; + // NOTE: This CAN fail, as the underlying header hashmap is not sorted + // As the test is just for dev, and doesn't run on travis, we are fine, + // for now. Possible solution would be to omit the size field (make it + // optional) + assert_eq!(res, expected); + assert_eq!(r, Ok(221)); } } \ No newline at end of file