From ed3763680588771cfb59b59403cc9830820cf72e Mon Sep 17 00:00:00 2001
From: stammw <begue.jc@gmail.com>
Date: Mon, 31 Dec 2018 08:22:11 +0100
Subject: [PATCH] QPACK: rewrite prefixed string using codec traits

---
 quinn-h3/Cargo.toml              |   3 +-
 quinn-h3/src/lib.rs              |   4 +
 quinn-h3/src/qpack/prefix_int.rs |   6 +-
 quinn-h3/src/qpack/string/mod.rs | 136 +++++++++++++++++++++++++++++--
 4 files changed, 140 insertions(+), 9 deletions(-)

diff --git a/quinn-h3/Cargo.toml b/quinn-h3/Cargo.toml
index 438bb48c1..b010b7ca9 100644
--- a/quinn-h3/Cargo.toml
+++ b/quinn-h3/Cargo.toml
@@ -26,4 +26,5 @@ lazy_static = "1"
 quinn-proto = { path = "../quinn-proto", version = "0.1.0" }
 
 [dev-dependencies]
-proptest = "0.6.0"
\ No newline at end of file
+proptest = "0.6.0"
+assert_matches = "1.1"
\ No newline at end of file
diff --git a/quinn-h3/src/lib.rs b/quinn-h3/src/lib.rs
index bca80ba52..f0dffbd58 100644
--- a/quinn-h3/src/lib.rs
+++ b/quinn-h3/src/lib.rs
@@ -5,5 +5,9 @@ extern crate lazy_static;
 #[macro_use]
 extern crate proptest;
 
+#[cfg(test)]
+#[macro_use]
+extern crate assert_matches;
+
 extern crate quinn_proto;
 mod qpack;
diff --git a/quinn-h3/src/qpack/prefix_int.rs b/quinn-h3/src/qpack/prefix_int.rs
index c6abda712..35db24c5b 100644
--- a/quinn-h3/src/qpack/prefix_int.rs
+++ b/quinn-h3/src/qpack/prefix_int.rs
@@ -2,7 +2,7 @@ use bytes::{Buf, BufMut};
 
 use quinn_proto::coding::{self, BufExt, BufMutExt};
 
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum Error {
     Overflow,
     UnexpectedEnd,
@@ -130,9 +130,9 @@ mod test {
     #[test]
     #[should_panic]
     fn size_too_big_of_size() {
-        let mut buf = vec![];
+        let buf = vec![];
         let mut read = Cursor::new(&buf);
-        super::decode(9, &mut read);
+        super::decode(9, &mut read).unwrap();
     }
 
     #[cfg(target_pointer_width = "64")]
diff --git a/quinn-h3/src/qpack/string/mod.rs b/quinn-h3/src/qpack/string/mod.rs
index cbb5ad8fb..9082b946f 100644
--- a/quinn-h3/src/qpack/string/mod.rs
+++ b/quinn-h3/src/qpack/string/mod.rs
@@ -1,12 +1,138 @@
 // This is only here because qpack is new and quinn no uses it yet.
 // TODO remove allow dead code
-#![allow(unused_imports)]
+#![allow(unused_imports, dead_code)]
 
-pub mod bitwin;
+mod bitwin;
+mod decode;
+mod encode;
+
+use std::borrow::Cow;
 pub use self::bitwin::BitWindow;
 
-pub mod decode;
 pub use self::decode::{DecodeIter, Error as HuffmanDecodingError, HpackStringDecode};
-
-pub mod encode;
 pub use self::encode::{Error as HuffmanEncodingError, HpackStringEncode};
+
+use crate::qpack::prefix_int::{self, Error as IntegerError};
+use bytes::{Buf, BufMut};
+use quinn_proto::coding::{self, BufExt, BufMutExt};
+
+#[derive(Debug, PartialEq)]
+pub enum Error {
+    UnexpectedEnd,
+    Integer(IntegerError),
+    HuffmanDecoding(HuffmanDecodingError),
+    HuffmanEncoding(HuffmanEncodingError),
+}
+
+pub fn decode<B: Buf>(size: u8, buf: &mut B) -> Result<Vec<u8>, Error> {
+    let (flags, len) = prefix_int::decode(size - 1, buf)?;
+    if buf.remaining() < len {
+        return Err(Error::UnexpectedEnd);
+    }
+
+    let payload = buf.take(len);
+    let value = if flags & 1 == 0 {
+        payload.collect()
+    } else {
+        let mut decoded = Vec::new();
+        for byte in payload.collect::<Vec<u8>>().hpack_decode() {
+            decoded.push(byte?);
+        }
+        decoded
+    };
+    Ok(value)
+}
+
+pub fn encode<B: BufMut>(size: u8, flags: u8, value: &[u8], buf: &mut B) -> Result<(), Error> {
+    let encoded = Vec::from(value).hpack_encode()?;
+    prefix_int::encode(size - 1, flags << 1 | 1, encoded.len(), buf);
+    for byte in encoded {
+        buf.write(byte);
+    }
+    Ok(())
+}
+
+impl From<HuffmanEncodingError> for Error {
+    fn from(error: HuffmanEncodingError) -> Self {
+        Error::HuffmanEncoding(error)
+    }
+}
+
+impl From<IntegerError> for Error {
+    fn from(error: IntegerError) -> Self {
+        Error::Integer(error)
+    }
+}
+
+impl From<HuffmanDecodingError> for Error {
+    fn from(error: HuffmanDecodingError) -> Self {
+        Error::HuffmanDecoding(error)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use std::io::Cursor;
+
+    #[test]
+    fn codec_6() {
+        let mut buf = Vec::new();
+        encode(6, 0b01, b"name without ref", &mut buf).unwrap();
+        let mut read = Cursor::new(&buf);
+        assert_eq!(
+            &buf,
+            &[
+                0b0110_1100,
+                168,
+                116,
+                149,
+                79,
+                6,
+                76,
+                231,
+                181,
+                42,
+                88,
+                89,
+                127
+            ]
+        );
+        assert_eq!(decode(6, &mut read).unwrap(), b"name without ref");
+    }
+
+    #[test]
+    fn codec_8() {
+        let mut buf = Vec::new();
+        encode(8, 0b01, b"name with ref", &mut buf).unwrap();
+        let mut read = Cursor::new(&buf);
+        assert_eq!(
+            &buf,
+            &[0b100_01010, 168, 116, 149, 79, 6, 76, 234, 88, 89, 127]
+        );
+        assert_eq!(decode(8, &mut read).unwrap(), b"name with ref");
+    }
+
+    #[test]
+    fn codec_8_empty() {
+        let mut buf = Vec::new();
+        encode(8, 0b01, b"", &mut buf).unwrap();
+        let mut read = Cursor::new(&buf);
+        assert_eq!(&buf, &[0b100_00000]);
+        assert_eq!(decode(8, &mut read).unwrap(), b"");
+    }
+
+    #[test]
+    fn decode_non_huffman() {
+        let buf = vec![0b0100_0011, b'b', b'a', b'r'];
+        let mut read = Cursor::new(&buf);
+        assert_eq!(decode(6, &mut read).unwrap(), b"bar");
+    }
+
+    #[test]
+    fn decode_too_short() {
+        let buf = vec![0b0100_0011, b'b', b'a'];
+        let mut read = Cursor::new(&buf);
+        assert_matches!(decode(6, &mut read), Err(Error::UnexpectedEnd));
+    }
+}