Skip to content

Commit

Permalink
fix(#11): 🐛 do not rely on Content-Length header
Browse files Browse the repository at this point in the history
  • Loading branch information
ym-project committed Sep 10, 2023
1 parent 8cd89ee commit cc1a161
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 12 deletions.
4 changes: 2 additions & 2 deletions src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ mod tests {
.to_http_parts();
let msgpack = MsgPackMessage::<Data>::new(&req, &mut payload).await;

assert_eq!(msgpack.err().unwrap(), MsgPackError::Overflow);
assert_eq!(msgpack.ok().unwrap(), Data { payload: true });

// Pass body and don't pass Content-Length header
let data =
Expand All @@ -198,6 +198,6 @@ mod tests {
.to_http_parts();
let msgpack = MsgPackMessage::<Data>::new(&req, &mut payload).await;

assert_eq!(msgpack.err().unwrap(), MsgPackError::Overflow);
assert_eq!(msgpack.ok().unwrap(), Data { payload: true });
}
}
17 changes: 7 additions & 10 deletions src/msgpack_message.rs
Original file line number Diff line number Diff line change
Expand Up @@ -74,26 +74,23 @@ impl<T: DeserializeOwned + 'static> Future for MsgPackMessage<T> {
}

let limit = self.limit;
let length = match self.length.take() {
Some(len) => len,
None => 0,
};

if length > limit {
return Poll::Ready(Err(MsgPackError::Overflow));

if let Some(len) = self.length.take() {
if len > limit {
return Poll::Ready(Err(MsgPackError::Overflow));
}
}

let mut stream = self.stream.take().expect("MsgPackMessage could not be used second time");

self.fut = Some(
async move {
let mut body = BytesMut::with_capacity(length);
let mut body = BytesMut::with_capacity(8192);

while let Some(item) = stream.next().await {
let chunk = item?;
let current_length = body.len() + chunk.len();

if current_length > length || current_length > limit {
if body.len() + chunk.len() > limit {
return Err(MsgPackError::Overflow);
} else {
body.extend_from_slice(&chunk);
Expand Down

0 comments on commit cc1a161

Please sign in to comment.