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)
}
}