1 files changed, 39 insertions, 12 deletions

diff --git a/src/net.rs b/src/net.rs
index 145256a..dfcbb56 100644
--- a/src/net.rs
+++ b/src/net.rs
@@ -4,18 +4,26 @@ use std::io::{Error};
 use std::result::Result;
 use std::net::TcpStream;
 use std::vec::Vec;
+use std::convert::TryInto;
+use std::io::Cursor;
+
+use flate2::Compress;
+use flate2::Decompress;
+use flate2::Compression;
+use flate2::FlushCompress;
+use flate2::FlushDecompress;
+use flate2::read::ZlibDecoder;
 
 extern crate log;
 // use log::{info, warn, debug};
 
-
 use crate::protocol::message;
 use crate::protocol::error::ErrorKind;
 
 pub struct Client {
     tcp_stream: TcpStream,
-    pub address: &'static str,
-    pub port: u32,
+    encoder: Compress,
+    decoder: Decompress,
     pub tls: bool,
     pub compression: bool,
 }
@@ -25,24 +33,25 @@ impl Client {
         use crate::protocol::message::handshake::{HandshakeDeserialize, HandshakeSerialize, HandshakeQRead, VariantMap};
         use crate::protocol::message::handshake::{ClientInitAck, ClientLogin, ClientLoginAck, SessionInit};
 
-        self.tcp_stream.write(&client.serialize().unwrap()).unwrap();
+        self.write(&client.serialize().unwrap()).unwrap();
 
         let mut buf: Vec<u8> = [0; 2048].to_vec();
-        let len = VariantMap::read(&mut self.tcp_stream, &mut buf).unwrap();
+        let len = VariantMap::read(self, &mut buf).unwrap();
         buf.truncate(len);
         let res = ClientInitAck::parse(&buf).unwrap();
         println!("res: {:?}", res);
 
         let login = ClientLogin {user: user.to_string(), password: pass.to_string()};
-        self.tcp_stream.write(&login.serialize().unwrap()).unwrap();
+        self.write(&login.serialize().unwrap()).unwrap();
+        println!("res: {:?}", res);
 
         let mut buf: Vec<u8> = [0; 2048].to_vec();
-        let len = VariantMap::read(&mut self.tcp_stream, &mut buf).unwrap();
+        let len = VariantMap::read(self, &mut buf).unwrap();
         buf.truncate(len);
         let _res = ClientLoginAck::parse(&buf).unwrap();
 
         let mut buf: Vec<u8> = [0; 2048].to_vec();
-        let len = VariantMap::read(&mut self.tcp_stream, &mut buf).unwrap();
+        let len = VariantMap::read(self, &mut buf).unwrap();
         buf.truncate(len);
         let res = SessionInit::parse(&buf).unwrap();
 
@@ -50,10 +59,27 @@ impl Client {
     }
 }
 
+impl std::io::Read for Client {
+    fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
+        return self.tcp_stream.read(buf);
+    }
+}
+
+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);
+    }
+
+    fn flush(&mut self) -> Result<(), Error> {
+       return self.tcp_stream.flush();
+    }
+}
+
 pub fn connect(address: &'static str, port: u32, tls: bool, compression: bool) -> Result<Client, Error> {
     use crate::protocol::primitive::deserialize::Deserialize;
 
-    //let mut s = BufWriter::new(TcpStream::connect(format!("{}:{}", address, port)).unwrap());
     let mut s = TcpStream::connect(format!("{}:{}", address, port)).unwrap();
 
     // Set Features
@@ -90,14 +116,15 @@ pub fn connect(address: &'static str, port: u32, tls: bool, compression: bool) -
     }
 
     let mut buf = [0; 4];
-    s.read_exact(&mut buf)?;
+    s.read(&mut buf)?;
     let (_, val) = ConnAck::parse(&buf).unwrap();
     println!("Received: {:?}", val);
 
+//    let sock = ZlibDecoder::new_with_buf(s, [0; 1].to_vec());
     let server: Client = Client {
         tcp_stream: s,
-        address: address,
-        port: port,
+        encoder: Compress::new(Compression::best(), true),
+        decoder: Decompress::new(true),
         tls: tls,
         compression: compression,
     };