From a0b1157c04ab548c58a0221d8af8cf44b115e099 Mon Sep 17 00:00:00 2001
From: Jovi Hsu <jv.hsu@outlook.com>
Date: Wed, 10 Nov 2021 17:35:35 +0800
Subject: [PATCH] dummy

---
 frontend/src/protocal/common.rs | 66 +++++++++++++++++++++++++++++++++
 frontend/src/protocal/vnc.rs    | 35 ++++++++++-------
 2 files changed, 88 insertions(+), 13 deletions(-)

diff --git a/frontend/src/protocal/common.rs b/frontend/src/protocal/common.rs
index 7619419..ff76c2d 100644
--- a/frontend/src/protocal/common.rs
+++ b/frontend/src/protocal/common.rs
@@ -1,3 +1,4 @@
+use std::slice::Iter;
 pub enum ProtocalHandlerOutput {
     Ok,
     WsBuf(Vec<u8>),
@@ -28,3 +29,68 @@ pub trait ProtocalImpl {
     fn new() -> Self;
     fn handle(&mut self, input: &[u8]) -> ProtocalHandlerOutput;
 }
+
+pub struct StreamReader<'a> {
+    inner: Iter<'a, u8>,
+}
+
+impl<'a> StreamReader<'a> {
+    pub fn new(data: &'a [u8]) -> Self {
+        Self {
+            inner: data.into_iter(),
+        }
+    }
+
+    pub fn read_u8(&mut self) -> Option<u8> {
+        self.inner.next().map(|&b| b)
+    }
+
+    pub fn read_u16(&mut self) -> Option<u16> {
+        let mut buf = [0u8; 2];
+        self.inner.by_ref().take(2).enumerate().for_each(|(i, b)| {
+            buf[i] = *b;
+        });
+        Some(u16::from_be_bytes(buf))
+    }
+
+    pub fn read_u32(&mut self) -> Option<u32> {
+        let mut buf = [0u8; 4];
+        self.inner.by_ref().take(4).enumerate().for_each(|(i, b)| {
+            buf[i] = *b;
+        });
+        Some(u32::from_be_bytes(buf))
+    }
+
+    pub fn read_s8(&mut self) -> Option<i8> {
+        Some(Self::read_u8(self).map(|b| b as i8)?)
+    }
+
+    pub fn read_s16(&mut self) -> Option<i16> {
+        Some(Self::read_u16(self).map(|b| b as i16)?)
+    }
+
+    pub fn read_s32(&mut self) -> Option<i32> {
+        Some(Self::read_u32(self).map(|b| b as i32)?)
+    }
+
+    pub fn read_string_with_len(&mut self, len: usize) -> Option<String> {
+        let mut buf = vec![0u8; len as usize];
+        self.inner
+            .by_ref()
+            .take(len as usize)
+            .enumerate()
+            .for_each(|(i, b)| {
+                buf[i] = *b;
+            });
+        Some(String::from_utf8(buf).unwrap())
+    }
+
+    pub fn read_string(&mut self) -> Option<String> {
+        let len = self.read_u16()? as usize;
+        Some(self.read_string_with_len(len)?)
+    }
+
+    pub fn eof(&mut self) -> bool {
+        self.inner.next().is_none()
+    }
+}
diff --git a/frontend/src/protocal/vnc.rs b/frontend/src/protocal/vnc.rs
index bcbbe5c..83f7b70 100644
--- a/frontend/src/protocal/vnc.rs
+++ b/frontend/src/protocal/vnc.rs
@@ -1,5 +1,3 @@
-
-
 use yew::services::ConsoleService;
 
 use super::common::*;
@@ -7,6 +5,7 @@ use super::common::*;
 const VNC_RFB33: &[u8; 12] = b"RFB 003.003\n";
 const VNC_RFB37: &[u8; 12] = b"RFB 003.007\n";
 const VNC_RFB38: &[u8; 12] = b"RFB 003.008\n";
+const VNC_VER_UNSUPPORTED: &str = "unsupported version";
 
 enum VncState {
     Handshake,
@@ -24,23 +23,27 @@ pub enum SecurityType {
 
 pub struct VncHandler {
     state: VncState,
-    output: Vec<u8>,
+    // output: Vec<u8>,
 }
 
 impl ProtocalImpl for VncHandler {
     fn new() -> Self {
         VncHandler {
             state: VncState::Handshake,
-            output: Vec::new(),
+            // output: Vec::new(),
         }
     }
 
     fn handle(&mut self, input: &[u8]) -> ProtocalHandlerOutput {
         match self.state {
             VncState::Handshake => {
-                self.handle_handshake(input);
+                let support_version = self.handle_handshake(input);
                 self.state = VncState::Authentication;
-                return ProtocalHandlerOutput::WsBuf(self.output.clone());
+                if let Ok(support_version) = support_version {
+                    ProtocalHandlerOutput::WsBuf(support_version.into())
+                } else {
+                    ProtocalHandlerOutput::Err(support_version.err().unwrap().to_string())
+                }
             }
             VncState::Authentication => {
                 ConsoleService::log(&format!("{:?}", input));
@@ -53,14 +56,20 @@ impl ProtocalImpl for VncHandler {
 
 // private methods
 impl VncHandler {
-    fn handle_handshake(&mut self, rfbversion: &[u8]) {
+    fn handle_handshake(&self, rfbversion: &[u8]) -> Result<&'static [u8], &'static str> {
         match rfbversion {
-            b"RFB 003.008\n" => {
-                for b in VNC_RFB33 {
-                    self.output.push(*b);
-                }
-            }
-            _ => panic!("unsupported version"),
+            b"RFB 003.003\n" => Ok(VNC_RFB33),
+            b"RFB 003.007\n" => Ok(VNC_RFB37),
+            b"RFB 003.008\n" => Ok(VNC_RFB38),
+            _ => Err(VNC_VER_UNSUPPORTED),
+        }
+    }
+    fn handle_authenticate(&mut self, auth: &[u8]) -> Result<&'static [u8], &'static str> {
+        match auth {
+            b"NONE\n" => Ok(b"\x00"),
+            b"VNCAUTH\n" => Ok(b"\x01"),
+            b"APPLETALK\n" => Ok(b"\x02"),
+            _ => Err(VNC_VER_UNSUPPORTED),
         }
     }
 }