update

3 files changed, 54 insertions, 16 deletions

diff --git a/src/net.rs b/src/net.rs
index dfcbb56..a4c74fb 100644
--- a/src/net.rs
+++ b/src/net.rs
@@ -21,7 +21,7 @@ use crate::protocol::message;
 use crate::protocol::error::ErrorKind;
 
 pub struct Client {
-    tcp_stream: TcpStream,
+    tcp_stream: ZlibDecoder<TcpStream>,
     encoder: Compress,
     decoder: Decompress,
     pub tls: bool,
@@ -61,7 +61,21 @@ impl Client {
 
 impl std::io::Read for Client {
     fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
-        return self.tcp_stream.read(buf);
+//        let mut cbuf = buf.to_vec();
+//        let read_bytes = self.tcp_stream.peek(&mut cbuf);
+//        println!("read bytes: {:?}", read_bytes);
+//        println!("cbuf: {:?}", cbuf);
+//        let decompressed_bytes_pre = self.decoder.total_out();
+//        self.decoder.decompress(&cbuf, buf, FlushDecompress::None)?;
+//        let decompressed_bytes = self.decoder.total_out();
+//        let in_bytes = self.decoder.total_in();
+//        println!("in bytes: {:?}", in_bytes);
+//        println!("decompressed bytes: {:?}", decompressed_bytes);
+//        println!("buf: {:?}", buf);
+//        return Ok(((decompressed_bytes - decompressed_bytes_pre)).try_into().unwrap());
+        let res = self.tcp_stream.read(buf);
+        println!("buf: {:?}, total in: {:?}, total out: {:?}", buf, self.tcp_stream.total_in(), self.tcp_stream.total_out());
+        return res;
     }
 }
 
@@ -69,7 +83,7 @@ impl std::io::Write for Client {
     fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
         let mut cbuf = Vec::with_capacity(buf.len());
         self.encoder.compress_vec(buf, &mut cbuf, FlushCompress::Finish)?;
-        return self.tcp_stream.write(&buf);
+        return self.tcp_stream.write(&cbuf);
     }
 
     fn flush(&mut self) -> Result<(), Error> {
@@ -121,8 +135,9 @@ pub fn connect(address: &'static str, port: u32, tls: bool, compression: bool) -
     println!("Received: {:?}", val);
 
 //    let sock = ZlibDecoder::new_with_buf(s, [0; 1].to_vec());
+    let sock = ZlibDecoder::new(s);
     let server: Client = Client {
-        tcp_stream: s,
+        tcp_stream: sock,
         encoder: Compress::new(Compression::best(), true),
         decoder: Decompress::new(true),
         tls: tls,
diff --git a/src/protocol/message/handshake/types.rs b/src/protocol/message/handshake/types.rs
index c809764..3c5d019 100644
--- a/src/protocol/message/handshake/types.rs
+++ b/src/protocol/message/handshake/types.rs
@@ -77,7 +77,7 @@ impl HandshakeQRead for VariantMap {
         let ulen = len as usize;
 
         // Read the 00 00 00 0a VariantType bytes and discard
-        s.read(&mut b[4..ulen])?;
+        s.read(&mut b[4..(ulen + 4)])?;
 
 //        let mut pos = 8;
 //        let len: usize = len as usize;
diff --git a/src/tests/handshake_types.rs b/src/tests/handshake_types.rs
index 701b871..227bdad 100644
--- a/src/tests/handshake_types.rs
+++ b/src/tests/handshake_types.rs
@@ -18,25 +18,48 @@ pub fn serialize_variantmap() {
 pub fn read_variantmap() {
     use std::io::Cursor;
 
-    let test_bytes: Vec<u8> = vec![0, 0, 0, 67, 0, 0, 0, 10, 0, 0, 0, 10, 0,
-        0, 0, 0, 20, 0, 67, 0, 111, 0, 110, 0, 102, 0, 105, 0, 103, 0, 117, 0, 114, 0, 101, 0, 100,
-                                   0, 0, 0, 1, 0, 1,
-                                     0, 0, 0, 10, 0,
-        0, 0, 0, 20, 0, 67, 0, 111, 0, 110, 0, 102, 0, 105, 0, 103, 0, 117, 0, 114, 0, 101, 0, 100,
-                                   0, 0, 0, 1, 0, 1,
-                                   0, 0, 0, 1];
+    let test_bytes: Vec<u8> = vec![
+        // len
+        0, 0, 0, 74,
+        // var
+        0, 0, 0, 10, // 4
+        // var
+        0, 0, 0, 10, 0, // 5
+        // strlen, str
+        0, 0, 0, 20, 0, 67, 0, 111, 0, 110, 0, 102, 0, 105, 0, 103, 0, 117, 0, 114, 0, 101, 0, 100, // 24
+        // bool
+        0, 0, 0, 1, 0, 1, // 6
+        // var
+        0, 0, 0, 10, 0, // 5
+        // strlen, str
+        0, 0, 0, 20, 0, 67, 0, 111, 0, 110, 0, 102, 0, 105, 0, 103, 0, 117, 0, 114, 0, 101, 0, 100, // 24
+        // bool
+        0, 0, 0, 1, 0, 1, //6
+        // extra
+        0, 0, 0, 1];
 
     let mut buf: Vec<u8> = [0; 78].to_vec();
     let len = VariantMap::read(&mut Cursor::new(&test_bytes), &mut buf).unwrap();
 
     assert_eq!(len, 78);
 
-    let result_bytes: Vec<u8> = vec![0, 0, 0, 67, 0, 0, 0, 10, 0, 0, 0, 10, 0,
+    let result_bytes: Vec<u8> = vec![
+        // len
+        0, 0, 0, 74,
+        // var
+        0, 0, 0, 10,
+        // var
+        0, 0, 0, 10, 0,
+        // strlen, str
         0, 0, 0, 20, 0, 67, 0, 111, 0, 110, 0, 102, 0, 105, 0, 103, 0, 117, 0, 114, 0, 101, 0, 100,
-                                   0, 0, 0, 1, 0, 1,
-                                     0, 0, 0, 10, 0,
+        // bool
+        0, 0, 0, 1, 0, 1,
+        // var
+        0, 0, 0, 10, 0,
+        // strlen, str
         0, 0, 0, 20, 0, 67, 0, 111, 0, 110, 0, 102, 0, 105, 0, 103, 0, 117, 0, 114, 0, 101, 0, 100,
-                                   0, 0, 0, 1, 0, 1];
+        // bool
+        0, 0, 0, 1, 0, 1];
     assert_eq!(buf, result_bytes);
 }