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 {
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()
}
}

View File

@ -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);
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),
}
}
_ => 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),
}
}
}