Fix issues with skipping certain stream segments.

This commit is contained in:
Kavin 2023-11-28 02:53:47 +00:00
parent 67a519978b
commit 3e8098d105
No known key found for this signature in database
GPG key ID: 6E4598CA5C92C41F

View file

@ -488,7 +488,6 @@ where
buffer: BytesMut, buffer: BytesMut,
found_stream: bool, found_stream: bool,
remaining: usize, remaining: usize,
finished: bool,
} }
impl<S> UmpTransformStream<S> impl<S> UmpTransformStream<S>
@ -501,7 +500,6 @@ where
buffer: BytesMut::new(), buffer: BytesMut::new(),
found_stream: false, found_stream: false,
remaining: 0, remaining: 0,
finished: false,
} }
} }
} }
@ -519,7 +517,7 @@ where
match item { match item {
Some(Ok(bytes)) => { Some(Ok(bytes)) => {
if this.found_stream { if this.found_stream {
if this.remaining > 0 { return if this.remaining > 0 {
let len = std::cmp::min(this.remaining, bytes.len()); let len = std::cmp::min(this.remaining, bytes.len());
this.remaining -= len; this.remaining -= len;
if this.remaining == 0 { if this.remaining == 0 {
@ -527,19 +525,16 @@ where
this.buffer.extend_from_slice(&bytes[len..]); this.buffer.extend_from_slice(&bytes[len..]);
this.found_stream = false; this.found_stream = false;
} }
return Poll::Ready(Some(Ok(bytes.slice(0..len)))); Poll::Ready(Some(Ok(bytes.slice(0..len))))
} else { } else {
this.finished = true; Poll::Ready(None)
};
return Poll::Ready(None);
}
} else { } else {
this.buffer.extend_from_slice(&bytes); this.buffer.extend_from_slice(&bytes);
} }
} }
Some(Err(e)) => return Poll::Ready(Some(Err(e))), Some(Err(e)) => return Poll::Ready(Some(Err(e))),
None => { None => {
this.finished = true;
return Poll::Ready(None); return Poll::Ready(None);
} }
} }
@ -554,7 +549,6 @@ where
let _ = this.buffer.split_to(s1 + s2 + segment_length as usize); let _ = this.buffer.split_to(s1 + s2 + segment_length as usize);
} }
} else { } else {
this.found_stream = true;
this.remaining = segment_length as usize - 1; this.remaining = segment_length as usize - 1;
let _ = this.buffer.split_to(s1 + s2 + 1); let _ = this.buffer.split_to(s1 + s2 + 1);
@ -566,6 +560,7 @@ where
return Poll::Ready(Some(Ok(this.buffer.split_to(len).into()))); return Poll::Ready(Some(Ok(this.buffer.split_to(len).into())));
} else { } else {
this.remaining -= this.buffer.len(); this.remaining -= this.buffer.len();
this.found_stream = true;
return Poll::Ready(Some(Ok(this.buffer.to_vec().into()))); return Poll::Ready(Some(Ok(this.buffer.to_vec().into())));
} }