diff --git a/quinn-proto/src/connection/mod.rs b/quinn-proto/src/connection/mod.rs index 67b00c1ff..d5bc84322 100644 --- a/quinn-proto/src/connection/mod.rs +++ b/quinn-proto/src/connection/mod.rs @@ -904,8 +904,12 @@ where } /// Read the next ordered chunk from the given recv stream - pub fn read_chunk(&mut self, id: StreamId) -> Result, ReadError> { - let result = self.streams.read_chunk(id); + pub fn read_chunk( + &mut self, + id: StreamId, + max_length: usize, + ) -> Result, ReadError> { + let result = self.streams.read_chunk(id, max_length); self.post_read(id, &result); Ok(result?.map(|x| x.result)) } diff --git a/quinn-proto/src/connection/streams.rs b/quinn-proto/src/connection/streams.rs index 7d9670796..1cae38502 100644 --- a/quinn-proto/src/connection/streams.rs +++ b/quinn-proto/src/connection/streams.rs @@ -203,8 +203,8 @@ impl Streams { self.try_read(id, |rs| rs.read_unordered()) } - pub(crate) fn read_chunk(&mut self, id: StreamId) -> ReadResult { - self.try_read(id, |rs| rs.read_chunk()) + pub(crate) fn read_chunk(&mut self, id: StreamId, max_length: usize) -> ReadResult { + self.try_read(id, |rs| rs.read_chunk(max_length)) } pub(crate) fn read_chunks( @@ -1139,8 +1139,8 @@ impl Recv { } } - fn read_chunk(&mut self) -> StreamReadResult { - match self.assembler.read_chunk(usize::MAX)? { + fn read_chunk(&mut self, max_length: usize) -> StreamReadResult { + match self.assembler.read_chunk(max_length)? { Some(bytes) => Ok(Some(bytes)), None => self.read_blocked().map(|()| None), } diff --git a/quinn/src/streams.rs b/quinn/src/streams.rs index 915530c31..d5543fada 100644 --- a/quinn/src/streams.rs +++ b/quinn/src/streams.rs @@ -389,13 +389,20 @@ where /// /// Slightly more efficient than `read` due to not copying. Chunk boundaries /// do not correspond to peer writes, and hence cannot be used as framing. - pub fn read_chunk(&mut self) -> ReadChunk<'_, S> { - ReadChunk { stream: self } + pub fn read_chunk(&mut self, max_length: usize) -> ReadChunk<'_, S> { + ReadChunk { + stream: self, + max_length, + } } /// Foundation of [`read_chunk()`]: RecvStream::read_chunk - fn poll_read_chunk(&mut self, cx: &mut Context) -> Poll, ReadError>> { - self.poll_read_generic(cx, |conn, stream| conn.inner.read_chunk(stream)) + fn poll_read_chunk( + &mut self, + cx: &mut Context, + max_length: usize, + ) -> Poll, ReadError>> { + self.poll_read_generic(cx, |conn, stream| conn.inner.read_chunk(stream, max_length)) } /// Read the next segments of data @@ -812,6 +819,7 @@ where S: proto::crypto::Session, { stream: &'a mut RecvStream, + max_length: usize, } impl<'a, S> Future for ReadChunk<'a, S> @@ -820,7 +828,8 @@ where { type Output = Result, ReadError>; fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { - self.stream.poll_read_chunk(cx) + let max_length = self.max_length; + self.stream.poll_read_chunk(cx, max_length) } }