This commit is contained in:
Jovi Hsu 2021-11-10 17:35:35 +08:00
parent 999edd10e7
commit a0b1157c04
2 changed files with 88 additions and 13 deletions

View File

@ -1,3 +1,4 @@
use std::slice::Iter;
pub enum ProtocalHandlerOutput { pub enum ProtocalHandlerOutput {
Ok, Ok,
WsBuf(Vec<u8>), WsBuf(Vec<u8>),
@ -28,3 +29,68 @@ pub trait ProtocalImpl {
fn new() -> Self; fn new() -> Self;
fn handle(&mut self, input: &[u8]) -> ProtocalHandlerOutput; 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()
}
}

View File

@ -1,5 +1,3 @@
use yew::services::ConsoleService; use yew::services::ConsoleService;
use super::common::*; use super::common::*;
@ -7,6 +5,7 @@ use super::common::*;
const VNC_RFB33: &[u8; 12] = b"RFB 003.003\n"; const VNC_RFB33: &[u8; 12] = b"RFB 003.003\n";
const VNC_RFB37: &[u8; 12] = b"RFB 003.007\n"; const VNC_RFB37: &[u8; 12] = b"RFB 003.007\n";
const VNC_RFB38: &[u8; 12] = b"RFB 003.008\n"; const VNC_RFB38: &[u8; 12] = b"RFB 003.008\n";
const VNC_VER_UNSUPPORTED: &str = "unsupported version";
enum VncState { enum VncState {
Handshake, Handshake,
@ -24,23 +23,27 @@ pub enum SecurityType {
pub struct VncHandler { pub struct VncHandler {
state: VncState, state: VncState,
output: Vec<u8>, // output: Vec<u8>,
} }
impl ProtocalImpl for VncHandler { impl ProtocalImpl for VncHandler {
fn new() -> Self { fn new() -> Self {
VncHandler { VncHandler {
state: VncState::Handshake, state: VncState::Handshake,
output: Vec::new(), // output: Vec::new(),
} }
} }
fn handle(&mut self, input: &[u8]) -> ProtocalHandlerOutput { fn handle(&mut self, input: &[u8]) -> ProtocalHandlerOutput {
match self.state { match self.state {
VncState::Handshake => { VncState::Handshake => {
self.handle_handshake(input); let support_version = self.handle_handshake(input);
self.state = VncState::Authentication; 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 => { VncState::Authentication => {
ConsoleService::log(&format!("{:?}", input)); ConsoleService::log(&format!("{:?}", input));
@ -53,14 +56,20 @@ impl ProtocalImpl for VncHandler {
// private methods // private methods
impl VncHandler { impl VncHandler {
fn handle_handshake(&mut self, rfbversion: &[u8]) { fn handle_handshake(&self, rfbversion: &[u8]) -> Result<&'static [u8], &'static str> {
match rfbversion { match rfbversion {
b"RFB 003.008\n" => { b"RFB 003.003\n" => Ok(VNC_RFB33),
for b in VNC_RFB33 { b"RFB 003.007\n" => Ok(VNC_RFB37),
self.output.push(*b); b"RFB 003.008\n" => Ok(VNC_RFB38),
_ => Err(VNC_VER_UNSUPPORTED),
} }
} }
_ => panic!("unsupported version"), 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),
} }
} }
} }