websocket not working

src/websoket_connection.rs

@@ -21,21 +21,16 @@ struct DataBlock {
 
     message_type: FrameType,
 
-    is_masked: bool,
-
-    length: u64,
-
-    mask_key: Option<u32>,
-
     data: Vec<u8>,
 }
 
-struct DataFrame {
-    frame_type: FrameType,
-    data: Box<[u8]>,
+pub struct DataFrame {
+    pub frame_type: FrameType,
+    pub data: Box<[u8]>,
 }
 
-enum FrameType {
+#[derive(PartialEq, Eq)]
+pub enum FrameType {
     Continuation,
     TextFrame,
     BinaryFrame,
@@ -52,12 +47,28 @@ impl WebsocketConnection {
     }
 
     pub async fn read_next_message(&mut self) -> io::Result<DataFrame> {
-        let mut data;
-        let frame_type;
-        {
-            let first_line = self.parse_single_block().await?;
-            data = first_line.data;
-            frame_type = first_line.message_type;
+        let first_line = self.parse_single_block().await?;
+
+        println!("Block read");
+
+        let mut data = first_line.data;
+        let frame_type = first_line.message_type;
+
+        if !first_line.is_final {
+            let mut current_line = self.parse_single_block().await?;
+
+            while !current_line.is_final {
+                if current_line.message_type != FrameType::Continuation {
+                    return Err(io::Error::new(
+                        io::ErrorKind::InvalidInput,
+                        "That is not how websocket works!!!",
+                    ));
+                }
+                data.extend_from_slice(&current_line.data);
+                current_line = self.parse_single_block().await?;
+            }
+
+            data.extend_from_slice(&current_line.data);
         }
 
         Ok(DataFrame {
@@ -66,6 +77,13 @@ impl WebsocketConnection {
         })
     }
 
+    fn unmask_block(data: &mut [u8], mask: u32) {
+        let mask_bytes = mask.to_be_bytes();
+        for (i, e) in data.iter_mut().enumerate() {
+            *e ^= mask_bytes[i % 4];
+        }
+    }
+
     async fn parse_single_block(&mut self) -> io::Result<DataBlock> {
         let mut first_line: [u8; 2] = [0; 2];
         self.stream.read_exact(&mut first_line).await?;
@@ -100,24 +118,24 @@ impl WebsocketConnection {
         };
 
         let masking_key = if mask {
-            Some(self.stream.read_u32().await?)
+            self.stream.read_u32().await?
         } else {
-            None
+            0
         };
 
-        let mut message_data = Vec::<u8>::with_capacity(length as usize);
-
+        let mut message_data = vec![0u8; length as usize];
         self.stream.read_exact(&mut message_data).await?;
 
+        if mask {
+            Self::unmask_block(&mut message_data, masking_key);
+        }
+
         Ok(DataBlock {
             is_final,
             e1: extension_bit_1,
             e2: extension_bit_2,
             e3: extension_bit_3,
             message_type,
-            is_masked: mask,
-            length,
-            mask_key: masking_key,
             data: message_data,
         })
     }
@@ -141,6 +159,10 @@ impl WebsocketConnection {
                 RequestHeader::Connection(con) => {
                     if con == Connection::Upgrade {
                         connection = true;
+                    } else if let Connection::Other(c) = con
+                        && c.contains("Upgrade")
+                    {
+                        connection = true;
                     }
                 }
                 RequestHeader::Other { name, value } => {
@@ -166,7 +188,7 @@ impl WebsocketConnection {
                 .with_code(crate::response::ResponseCode::SwitchingProtocols)
                 .with_header(crate::response::ResponseHeader::Upgrade(Upgrade {
                     protocol: Protocol::Websocket,
-                    version: "".into(),
+                    version: None,
                 }))
                 .with_header(crate::response::ResponseHeader::Connection(
                     Connection::Upgrade,
@@ -178,6 +200,8 @@ impl WebsocketConnection {
                 .respond(&mut stream)
                 .await?;
 
+            stream.flush().await?;
+
             Ok(Self { stream })
         } else {
             Response::new()