stream: fix panic in Merge and Chain size_hint (#2430)

5 files changed, 63 insertions, 18 deletions

diff --git a/tokio/src/stream/chain.rs b/tokio/src/stream/chain.rs
index 5f0324a4..6124c91e 100644
--- a/tokio/src/stream/chain.rs
+++ b/tokio/src/stream/chain.rs
@@ -44,14 +44,6 @@ where
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let (a_lower, a_upper) = self.a.size_hint();
-        let (b_lower, b_upper) = self.b.size_hint();
-
-        let upper = match (a_upper, b_upper) {
-            (Some(a_upper), Some(b_upper)) => Some(a_upper + b_upper),
-            _ => None,
-        };
-
-        (a_lower + b_lower, upper)
+        super::merge_size_hints(self.a.size_hint(), self.b.size_hint())
     }
 }
diff --git a/tokio/src/stream/merge.rs b/tokio/src/stream/merge.rs
index 4850cd40..50ba518c 100644
--- a/tokio/src/stream/merge.rs
+++ b/tokio/src/stream/merge.rs
@@ -52,15 +52,7 @@ where
     }
 
     fn size_hint(&self) -> (usize, Option<usize>) {
-        let (a_lower, a_upper) = self.a.size_hint();
-        let (b_lower, b_upper) = self.b.size_hint();
-
-        let upper = match (a_upper, b_upper) {
-            (Some(a_upper), Some(b_upper)) => Some(a_upper + b_upper),
-            _ => None,
-        };
-
-        (a_lower + b_lower, upper)
+        super::merge_size_hints(self.a.size_hint(), self.b.size_hint())
     }
 }
 
diff --git a/tokio/src/stream/mod.rs b/tokio/src/stream/mod.rs
index 307ead5f..a59bdfcd 100644
--- a/tokio/src/stream/mod.rs
+++ b/tokio/src/stream/mod.rs
@@ -817,3 +817,16 @@ pub trait StreamExt: Stream {
 }
 
 impl<St: ?Sized> StreamExt for St where St: Stream {}
+
+/// Merge the size hints from two streams.
+fn merge_size_hints(
+    (left_low, left_high): (usize, Option<usize>),
+    (right_low, right_hign): (usize, Option<usize>),
+) -> (usize, Option<usize>) {
+    let low = left_low.saturating_add(right_low);
+    let high = match (left_high, right_hign) {
+        (Some(h1), Some(h2)) => h1.checked_add(h2),
+        _ => None,
+    };
+    (low, high)
+}
diff --git a/tokio/tests/stream_chain.rs b/tokio/tests/stream_chain.rs
index 0e14618b..98461a8c 100644
--- a/tokio/tests/stream_chain.rs
+++ b/tokio/tests/stream_chain.rs
@@ -69,3 +69,27 @@ async fn pending_first() {
     assert_eq!(stream.size_hint(), (0, None));
     assert_eq!(None, assert_ready!(stream.poll_next()));
 }
+
+#[test]
+fn size_overflow() {
+    struct Monster;
+
+    impl tokio::stream::Stream for Monster {
+        type Item = ();
+        fn poll_next(
+            self: std::pin::Pin<&mut Self>,
+            _cx: &mut std::task::Context<'_>,
+        ) -> std::task::Poll<Option<()>> {
+            panic!()
+        }
+
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            (usize::max_value(), Some(usize::max_value()))
+        }
+    }
+
+    let m1 = Monster;
+    let m2 = Monster;
+    let m = m1.chain(m2);
+    assert_eq!(m.size_hint(), (usize::max_value(), None));
+}
diff --git a/tokio/tests/stream_merge.rs b/tokio/tests/stream_merge.rs
index f0168d72..45ecdcb6 100644
--- a/tokio/tests/stream_merge.rs
+++ b/tokio/tests/stream_merge.rs
@@ -52,3 +52,27 @@ async fn merge_async_streams() {
     assert!(rx.is_woken());
     assert_eq!(None, assert_ready!(rx.poll_next()));
 }
+
+#[test]
+fn size_overflow() {
+    struct Monster;
+
+    impl tokio::stream::Stream for Monster {
+        type Item = ();
+        fn poll_next(
+            self: std::pin::Pin<&mut Self>,
+            _cx: &mut std::task::Context<'_>,
+        ) -> std::task::Poll<Option<()>> {
+            panic!()
+        }
+
+        fn size_hint(&self) -> (usize, Option<usize>) {
+            (usize::max_value(), Some(usize::max_value()))
+        }
+    }
+
+    let m1 = Monster;
+    let m2 = Monster;
+    let m = m1.merge(m2);
+    assert_eq!(m.size_hint(), (usize::max_value(), None));
+}