Skip to content

Commit

Permalink
clean up multipart and field stream trait impl (#2462)
Browse files Browse the repository at this point in the history
  • Loading branch information
fakeshadow committed Nov 25, 2021
1 parent 52bbbd1 commit 89c6d62
Showing 1 changed file with 41 additions and 34 deletions.
75 changes: 41 additions & 34 deletions actix-multipart/src/server.rs
Expand Up @@ -32,7 +32,7 @@ const MAX_HEADERS: usize = 32;
pub struct Multipart {
safety: Safety,
error: Option<MultipartError>,
inner: Option<Rc<RefCell<InnerMultipart>>>,
inner: Option<InnerMultipart>,
}

enum InnerMultipartItem {
Expand Down Expand Up @@ -96,12 +96,12 @@ impl Multipart {
Multipart {
error: None,
safety: Safety::new(),
inner: Some(Rc::new(RefCell::new(InnerMultipart {
inner: Some(InnerMultipart {
boundary,
payload: PayloadRef::new(PayloadBuffer::new(stream)),
state: InnerState::FirstBoundary,
item: InnerMultipartItem::None,
}))),
}),
}
}

Expand All @@ -118,20 +118,27 @@ impl Multipart {
impl Stream for Multipart {
type Item = Result<Field, MultipartError>;

fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
if let Some(err) = self.error.take() {
Poll::Ready(Some(Err(err)))
} else if self.safety.current() {
let this = self.get_mut();
let mut inner = this.inner.as_mut().unwrap().borrow_mut();
if let Some(mut payload) = inner.payload.get_mut(&this.safety) {
payload.poll_stream(cx)?;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let this = self.get_mut();

match this.inner.as_mut() {
Some(inner) => {
if let Some(mut buffer) = inner.payload.get_mut(&this.safety) {
// check safety and poll read payload to buffer.
buffer.poll_stream(cx)?;
} else if !this.safety.is_clean() {
// safety violation
return Poll::Ready(Some(Err(MultipartError::NotConsumed)));
} else {
return Poll::Pending;
}

inner.poll(&this.safety, cx)
}
inner.poll(&this.safety, cx)
} else if !self.safety.is_clean() {
Poll::Ready(Some(Err(MultipartError::NotConsumed)))
} else {
Poll::Pending
None => Poll::Ready(Some(Err(this
.error
.take()
.expect("Multipart polled after finish")))),
}
}
}
Expand All @@ -152,17 +159,15 @@ impl InnerMultipart {
Ok(httparse::Status::Complete((_, hdrs))) => {
// convert headers
let mut headers = HeaderMap::with_capacity(hdrs.len());

for h in hdrs {
if let Ok(name) = HeaderName::try_from(h.name) {
if let Ok(value) = HeaderValue::try_from(h.value) {
headers.append(name, value);
} else {
return Err(ParseError::Header.into());
}
} else {
return Err(ParseError::Header.into());
}
let name =
HeaderName::try_from(h.name).map_err(|_| ParseError::Header)?;
let value = HeaderValue::try_from(h.value)
.map_err(|_| ParseError::Header)?;
headers.append(name, value);
}

Ok(Some(headers))
}
Ok(httparse::Status::Partial) => Err(ParseError::Header.into()),
Expand Down Expand Up @@ -458,17 +463,19 @@ impl Stream for Field {
type Item = Result<Bytes, MultipartError>;

fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
if self.safety.current() {
let mut inner = self.inner.borrow_mut();
if let Some(mut payload) = inner.payload.as_ref().unwrap().get_mut(&self.safety) {
payload.poll_stream(cx)?;
}
inner.poll(&self.safety)
} else if !self.safety.is_clean() {
Poll::Ready(Some(Err(MultipartError::NotConsumed)))
let this = self.get_mut();
let mut inner = this.inner.borrow_mut();
if let Some(mut buffer) = inner.payload.as_ref().unwrap().get_mut(&this.safety) {
// check safety and poll read payload to buffer.
buffer.poll_stream(cx)?;
} else if !this.safety.is_clean() {
// safety violation
return Poll::Ready(Some(Err(MultipartError::NotConsumed)));
} else {
Poll::Pending
return Poll::Pending;
}

inner.poll(&this.safety)
}
}

Expand Down

0 comments on commit 89c6d62

Please sign in to comment.