From 912f2f0bf3d1585b2ab42f8a4de46aacbc938042 Mon Sep 17 00:00:00 2001
From: Lorenz Stechauner <lorenz.stechauner@necronda.net>
Date: Fri, 21 May 2021 23:07:20 +0200
Subject: [PATCH] Add webocket frame header

---
 src/websocket/handler.rs | 26 ++++++++++++++----
 src/websocket/mod.rs     | 58 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+), 5 deletions(-)

diff --git a/src/websocket/handler.rs b/src/websocket/handler.rs
index db7049c..b6d3c93 100644
--- a/src/websocket/handler.rs
+++ b/src/websocket/handler.rs
@@ -4,6 +4,7 @@ use crate::http;
 use base64;
 use crypto;
 use crypto::digest::Digest;
+use crate::websocket::FrameHeader;
 
 pub fn connection_handler(
     client: &mut http::HttpStream,
@@ -99,11 +100,26 @@ pub fn connection_handler(
     res.send(&mut client.stream).unwrap();
 
     loop {
-        let mut buf = [0u8; 8192];
-        let res = client.stream.read(&mut buf).unwrap();
-        if res == 0 {
-            break;
+        let header = FrameHeader::from(&mut client.stream).unwrap();
+        if header.mask {
+            println!("Mask: {:X}", header.masking_key.unwrap());
         }
-        println!("Msg: {}/{}", res, String::from_utf8_lossy(&buf));
+
+        let mut buf = [0u8; 8192];
+        client.stream.read_exact(&mut buf[..header.payload_len() as usize]).unwrap();
+
+        if header.mask {
+            let key: [u8; 4] = [
+                (header.masking_key.unwrap() >> 24) as u8,
+                ((header.masking_key.unwrap() >> 16) & 0xFF) as u8,
+                ((header.masking_key.unwrap() >> 8) & 0xFF) as u8,
+                (header.masking_key.unwrap() & 0xFF) as u8,
+            ];
+            for (pos, byte) in buf.iter_mut().enumerate() {
+                *byte ^= key[pos % 4];
+            }
+        }
+
+        println!("Msg: {}", String::from_utf8_lossy(&buf[..header.payload_len() as usize]));
     }
 }
diff --git a/src/websocket/mod.rs b/src/websocket/mod.rs
index 7298495..864d9a0 100644
--- a/src/websocket/mod.rs
+++ b/src/websocket/mod.rs
@@ -1,3 +1,61 @@
 mod handler;
 
 pub use handler::*;
+use crate::http::Stream;
+use crate::error::Error;
+
+pub struct FrameHeader {
+    fin: bool,
+    rsv1: bool,
+    rsv2: bool,
+    rsv3: bool,
+    opcode: u8,
+    mask: bool,
+    payload_len: u8,
+    ex_payload_len: Option<u64>,
+    masking_key: Option<u32>,
+}
+
+impl FrameHeader {
+    pub fn from(socket: &mut Stream) -> Result<Self, Error> {
+        let mut data = [0u8; 2];
+        socket.read_exact(&mut data)?;
+
+        let mut header = FrameHeader {
+            fin: data[0] & 0x80 != 0,
+            rsv1: data[0] & 0x40 != 0,
+            rsv2: data[0] & 0x20 != 0,
+            rsv3: data[0] & 0x10 != 0,
+            opcode: data[0] & 0x0F,
+            mask: data[1] & 0x80 != 0,
+            payload_len: data[1] & 0x7F,
+            ex_payload_len: None,
+            masking_key: None,
+        };
+
+        if header.payload_len == 126 {
+            let mut data = [0u8; 2];
+            socket.read_exact(&mut data)?;
+            header.ex_payload_len = Some(u16::from_be_bytes(data) as u64);
+        } else if header.payload_len == 127 {
+            let mut data = [0u8; 8];
+            socket.read_exact(&mut data)?;
+            header.ex_payload_len = Some(u64::from_be_bytes(data));
+        }
+
+        if header.mask {
+            let mut data =[0u8; 4];
+            socket.read_exact(&mut data)?;
+            header.masking_key = Some(u32::from_be_bytes(data));
+        }
+
+        Ok(header)
+    }
+
+    pub fn payload_len(&self) -> u64 {
+        match self.ex_payload_len {
+            Some(val) => val,
+            None => self.payload_len as u64,
+        }
+    }
+}