From dff7eeed200204fef91d4a3edb84ddefa39d4553 Mon Sep 17 00:00:00 2001
From: stammw <begue.jc@gmail.com>
Date: Sat, 23 Nov 2019 10:03:20 +0100
Subject: [PATCH] H3: fix and optimize new StreamType decoding

---
 quinn-h3/src/streams.rs   | 25 +++++++++++++++----------
 quinn-proto/src/varint.rs |  7 +++++++
 2 files changed, 22 insertions(+), 10 deletions(-)

diff --git a/quinn-h3/src/streams.rs b/quinn-h3/src/streams.rs
index 4a4efe553..2ee3b8422 100644
--- a/quinn-h3/src/streams.rs
+++ b/quinn-h3/src/streams.rs
@@ -36,13 +36,13 @@ impl TryFrom<(StreamType, RecvStream)> for NewUni {
 }
 
 pub struct RecvUni {
-    inner: Option<(RecvStream, Vec<u8>, usize)>,
+    inner: Option<(RecvStream, [u8; VarInt::MAX_SIZE], usize, usize)>,
 }
 
 impl RecvUni {
     pub fn new(recv: RecvStream) -> Self {
         Self {
-            inner: Some((recv, vec![0u8; VarInt::MAX.size()], 0)),
+            inner: Some((recv, [0u8; VarInt::MAX_SIZE], 1, 0)),
         }
     }
 }
@@ -54,24 +54,29 @@ impl Future for RecvUni {
         loop {
             match self.inner {
                 None => panic!("polled after resolved"),
-                Some((ref mut recv, ref mut buf, ref mut len)) => {
-                    match ready!(Pin::new(recv).poll_read(cx, &mut buf[..=*len]))? {
+                Some((ref mut recv, ref mut buf, ref mut expected, ref mut len)) => {
+                    match ready!(Pin::new(recv).poll_read(cx, &mut buf[*len..*expected]))? {
                         0 => {
                             return Poll::Ready(Err(Error::Peer(
                                 "Uni stream closed before type received".into(),
                             )))
                         }
-                        _ => {
-                            *len += 1;
-                            let mut cur = io::Cursor::new(&buf);
-                            if let Ok(ty) = StreamType::decode(&mut cur) {
+                        read => {
+                            *len += read;
+                            if *len == 1 {
+                                *expected = VarInt::encoded_size(buf[0]);
+                            }
+                            if len == expected {
+                                let mut cur = io::Cursor::new(&buf);
+                                let ty = StreamType::decode(&mut cur)
+                                    .map_err(|_| Error::Internal("stream type decode"))?;
                                 match mem::replace(&mut self.inner, None) {
-                                    Some((recv, _, _)) => {
+                                    Some((recv, _, _, _)) => {
                                         return Poll::Ready(NewUni::try_from((ty, recv)))
                                     }
                                     _ => unreachable!(),
                                 };
-                            };
+                            }
                         }
                     }
                 }
diff --git a/quinn-proto/src/varint.rs b/quinn-proto/src/varint.rs
index eeab89114..405b7ae80 100644
--- a/quinn-proto/src/varint.rs
+++ b/quinn-proto/src/varint.rs
@@ -16,6 +16,8 @@ pub struct VarInt(pub(crate) u64);
 impl VarInt {
     /// The largest representable value
     pub const MAX: VarInt = VarInt((1 << 62) - 1);
+    /// The largest encoded value length
+    pub const MAX_SIZE: usize = 8;
 
     /// Construct a `VarInt` infallibly
     pub const fn from_u32(x: u32) -> Self {
@@ -60,6 +62,11 @@ impl VarInt {
             unreachable!("malformed VarInt");
         }
     }
+
+    /// Length of an encoded value from its first byte
+    pub fn encoded_size(first: u8) -> usize {
+        2usize.pow((first >> 6) as u32)
+    }
 }
 
 impl From<VarInt> for u64 {