From 332999de335eaae07755f54f51cba0fcf9d68eb0 Mon Sep 17 00:00:00 2001 From: Jovi Hsu Date: Wed, 12 Oct 2022 09:04:23 +0000 Subject: [PATCH] Add tight encoding support for web vnc --- README.md | 5 +- webvnc/Cargo.toml | 1 + webvnc/src/canvas.rs | 18 +- webvnc/src/lib.rs | 1 + webvnc/src/utils.rs | 35 ++ webvnc/src/vnc/decoder/mod.rs | 7 + webvnc/src/vnc/decoder/raw.rs | 51 +++ webvnc/src/vnc/decoder/tight.rs | 416 ++++++++++++++++++++ webvnc/src/vnc/decoder/zlib.rs | 77 ++++ webvnc/src/vnc/{zlib.rs => decoder/zrle.rs} | 84 ++-- webvnc/src/vnc/mod.rs | 3 +- webvnc/src/vnc/vnc_impl.rs | 220 ++++++----- 12 files changed, 765 insertions(+), 153 deletions(-) create mode 100644 webvnc/src/vnc/decoder/mod.rs create mode 100644 webvnc/src/vnc/decoder/raw.rs create mode 100644 webvnc/src/vnc/decoder/tight.rs create mode 100644 webvnc/src/vnc/decoder/zlib.rs rename webvnc/src/vnc/{zlib.rs => decoder/zrle.rs} (84%) diff --git a/README.md b/README.md index d8e8ea4..e32458e 100644 --- a/README.md +++ b/README.md @@ -17,8 +17,9 @@ ## Milestones * VNC Clients: - - Raw encoding support (Done) - - ZRLE encoding support (Done) (Default) + - Raw encoding support + - Tight encoding support (Default) + - ZRLE encoding support - Other encoding support (WIP) * SSH Clients: diff --git a/webvnc/Cargo.toml b/webvnc/Cargo.toml index a143172..1048a70 100644 --- a/webvnc/Cargo.toml +++ b/webvnc/Cargo.toml @@ -34,6 +34,7 @@ features = [ "FileReader", "HtmlButtonElement", "HtmlCanvasElement", + "HtmlImageElement", "ImageData", "Location", "KeyboardEvent", diff --git a/webvnc/src/canvas.rs b/webvnc/src/canvas.rs index 0c1ae45..a1b1c29 100644 --- a/webvnc/src/canvas.rs +++ b/webvnc/src/canvas.rs @@ -7,7 +7,8 @@ use crate::{ use wasm_bindgen::prelude::*; use wasm_bindgen::{Clamped, JsCast}; use web_sys::{ - CanvasRenderingContext2d, HtmlButtonElement, HtmlCanvasElement, KeyboardEvent, MouseEvent, + CanvasRenderingContext2d, HtmlButtonElement, HtmlCanvasElement, HtmlImageElement, + KeyboardEvent, MouseEvent, }; struct Canvas { canvas: HtmlCanvasElement, @@ -206,6 +207,21 @@ impl Canvas { let data = data.unwrap(); let _ = self.ctx.put_image_data(&data, ri.x as f64, ri.y as f64); } + ImageType::Jpeg => { + let image = HtmlImageElement::new().unwrap(); + let base64 = crate::utils::base64_encode(&ri.data); + image.set_src(&format!( + "data:image/jpeg;base64,{}", + std::str::from_utf8(&base64).unwrap() + )); + let _ = self.ctx.draw_image_with_html_image_element_and_dw_and_dh( + &image, + ri.x as f64, + ri.y as f64, + ri.width as f64, + ri.height as f64, + ); + } } } diff --git a/webvnc/src/lib.rs b/webvnc/src/lib.rs index fd2d7be..4dec500 100644 --- a/webvnc/src/lib.rs +++ b/webvnc/src/lib.rs @@ -114,6 +114,7 @@ fn vnc_out_handler(ws: &WebSocket, vnc: &Vnc, canvas: &CanvasUtils) { } fn vnc_close_handle(vnc: &Vnc, canvas: &CanvasUtils, msg: &str) { + console_log!("Websocket disconnect with message {}", msg); vnc.close(); unsafe { REFRESHER.take(); diff --git a/webvnc/src/utils.rs b/webvnc/src/utils.rs index b1d7929..a1c4122 100644 --- a/webvnc/src/utils.rs +++ b/webvnc/src/utils.rs @@ -8,3 +8,38 @@ pub fn set_panic_hook() { #[cfg(feature = "console_error_panic_hook")] console_error_panic_hook::set_once(); } + +const BASIS_64: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + +pub fn base64_encode(input: &[u8]) -> Vec { + let mut i = 0; + let len = input.len(); + let mut out = Vec::with_capacity(((len + 2) / 3 * 4) + 1); + + while i < len - 2 { + out.push(BASIS_64[(input[i] as usize >> 2) & 0x3F]); + out.push( + BASIS_64[((input[i] as usize & 0x3) << 4) | ((input[i + 1] as usize & 0xF0) >> 4)], + ); + out.push( + BASIS_64[((input[i + 1] as usize & 0xF) << 2) | ((input[i + 2] as usize & 0xC0) >> 6)], + ); + out.push(BASIS_64[input[i + 2] as usize & 0x3F]); + i += 3; + } + + if i < len { + out.push(BASIS_64[(input[i] as usize >> 2) & 0x3F]); + if i == (len - 1) { + out.push(BASIS_64[((input[i] as usize & 0x3) << 4)]); + out.push(0x3d); // = + } else { + out.push( + BASIS_64[((input[i] as usize & 0x3) << 4) | ((input[i + 1] as usize & 0xF0) >> 4)], + ); + out.push(BASIS_64[((input[i + 1] as usize & 0xF) << 2)]); + } + out.push(0x3d); // = + } + out +} diff --git a/webvnc/src/vnc/decoder/mod.rs b/webvnc/src/vnc/decoder/mod.rs new file mode 100644 index 0000000..ef76d17 --- /dev/null +++ b/webvnc/src/vnc/decoder/mod.rs @@ -0,0 +1,7 @@ +mod raw; +mod tight; +mod zlib; +mod zrle; +pub use raw::Decoder as RawDecoder; +pub use tight::Decoder as TightDecoder; +pub use zrle::Decoder as ZrleDecoder; diff --git a/webvnc/src/vnc/decoder/raw.rs b/webvnc/src/vnc/decoder/raw.rs new file mode 100644 index 0000000..b7243e1 --- /dev/null +++ b/webvnc/src/vnc/decoder/raw.rs @@ -0,0 +1,51 @@ +use super::super::vnc_impl::{PixelFormat, VncRect}; +use super::super::{ImageData, ImageType, StreamReader}; +use std::io::Result; + +#[derive(Default)] +pub struct Decoder {} + +impl Decoder { + pub fn new() -> Self { + Self {} + } + + pub fn decode( + &mut self, + format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + out_wait: &mut usize, + ) -> Result>> { + let len = rect.width as usize * rect.height as usize * format.bits_per_pixel as usize / 8; + if len > input.remain { + *out_wait = len; + return Ok(None); + } + + let mut image_data = Vec::with_capacity(len); + input.read_exact_vec(&mut image_data, len); + let mut y = 0; + let mut x = 0; + + while y < rect.height { + while x < rect.width { + let idx = (y as usize * rect.width as usize + x as usize) * 4; + image_data.swap(idx, idx + 2); + x += 1; + } + x = 0; + y += 1; + } + Ok(Some(vec![ + (ImageData { + type_: ImageType::Raw, + x: rect.x, + y: rect.y, + width: rect.width, + height: rect.height, + data: image_data, + }), + ])) + } +} diff --git a/webvnc/src/vnc/decoder/tight.rs b/webvnc/src/vnc/decoder/tight.rs new file mode 100644 index 0000000..5d8898e --- /dev/null +++ b/webvnc/src/vnc/decoder/tight.rs @@ -0,0 +1,416 @@ +use super::super::vnc_impl::{PixelFormat, VncRect}; +use super::super::{ImageData, ImageType, StreamReader}; +use super::zlib::ZlibReader; +use std::io::{Error, ErrorKind, Read, Result}; + +const MAX_PALETTE: usize = 256; + +#[derive(Default)] +pub struct Decoder { + zlibs: [Option; 4], + ctrl: Option, + filter: Option, + len: usize, + num_colors: usize, + palette: Vec, +} + +impl Decoder { + pub fn new() -> Self { + let mut new = Self { + palette: Vec::with_capacity(MAX_PALETTE * 4), + ..Default::default() + }; + for i in 0..4 { + let decompressor = flate2::Decompress::new(true); + new.zlibs[i] = Some(decompressor); + } + new + } + + pub fn decode( + &mut self, + format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + out_wait: &mut usize, + ) -> Result>> { + if self.ctrl.is_none() { + if input.remain < 1 { + *out_wait = 1; + return Ok(None); + } + let ctrl = input.read_u8(); + for i in 0..4 { + if (ctrl >> i) & 1 == 1 { + self.zlibs[i].as_mut().unwrap().reset(true); + } + } + + // Figure out filter + self.ctrl = Some(ctrl >> 4); + } + + let mut ret = Vec::new(); + match self.ctrl.as_ref().unwrap() { + 8 => { + // fill Rect + self.fill_rect(format, rect, input, out_wait, &mut ret)? + } + 9 => { + // jpeg Rect + self.jpeg_rect(format, rect, input, out_wait, &mut ret)? + } + 10 => { + // png Rect + return Err(Error::new( + ErrorKind::InvalidData, + "PNG received in standard Tight rect", + )); + } + x if x & 0x8 == 0 => { + // basic Rect + self.basic_rect(format, rect, input, out_wait, &mut ret)? + } + _ => { + let err_msg = format!( + "Illegal tight compression received ({})", + self.ctrl.as_ref().unwrap() + ); + return Err(Error::new(ErrorKind::InvalidData, err_msg)); + } + } + + if !ret.is_empty() { + self.ctrl = None; + Ok(Some(ret)) + } else { + Ok(None) + } + } + + fn read_data(&mut self, input: &mut StreamReader, wait: &mut usize) -> Option> { + if self.len == 0 { + if input.remain < 3 { + *wait = 3; + return None; + } + + let mut byte = input.read_u8() as usize; + self.len = byte & 0x7f; + if byte & 0x80 == 0x80 { + byte = input.read_u8() as usize; + self.len |= (byte & 0x7f) << 7; + + if byte & 0x80 == 0x80 { + byte = input.read_u8() as usize; + self.len |= (byte as usize) << 14; + } + } + } + + if input.remain < self.len { + *wait = self.len; + None + } else { + let mut data = Vec::with_capacity(self.len); + input.read_exact_vec(&mut data, self.len); + self.len = 0; + Some(data) + } + } + + fn fill_rect( + &mut self, + _format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + wait: &mut usize, + out: &mut Vec, + ) -> Result<()> { + if input.remain < 3 { + *wait = 3; + } else { + let mut rgb = [0; 3]; + input.read_exact(&mut rgb, 3); + rgb.swap(0, 2); + out.push(ImageData { + type_: ImageType::Fill, + x: rect.x, + y: rect.y, + width: rect.width, + height: rect.height, + data: rgb.to_vec(), + }) + } + Ok(()) + } + + fn jpeg_rect( + &mut self, + _format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + wait: &mut usize, + out: &mut Vec, + ) -> Result<()> { + if let Some(d) = self.read_data(input, wait) { + out.push(ImageData { + type_: ImageType::Jpeg, + x: rect.x, + y: rect.y, + width: rect.width, + height: rect.height, + data: d, + }); + } + Ok(()) + } + + fn basic_rect( + &mut self, + format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + wait: &mut usize, + out: &mut Vec, + ) -> Result<()> { + if self.filter.is_none() { + let filter; + if self.ctrl.as_ref().unwrap() & 0x4 == 4 { + if input.remain < 1 { + *wait = 1; + return Ok(()); + } + filter = input.read_u8(); + } else { + // Implicit CopyFilter + filter = 0; + } + + self.filter = Some(filter); + } + + let stream_id = self.ctrl.as_ref().unwrap() & 0x3; + match self.filter.as_ref().unwrap() { + 0 => { + // copy filter + self.copy_filter(stream_id, format, rect, input, wait, out)?; + } + 1 => { + // palette + self.palette_filter(stream_id, format, rect, input, wait, out)? + } + 2 => { + // gradient + return Err(Error::new( + ErrorKind::InvalidData, + "Gradient filter not implemented", + )); + } + _ => { + let err_msg = format!( + "Illegal tight filter received (filter: {})", + self.filter.as_ref().unwrap() + ); + return Err(Error::new(ErrorKind::InvalidData, err_msg)); + } + } + + if !out.is_empty() { + self.filter = None; + } + Ok(()) + } + + fn copy_filter( + &mut self, + stream: u8, + _format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + wait: &mut usize, + out: &mut Vec, + ) -> Result<()> { + let uncompressed_size = rect.width as usize * rect.height as usize * 3; + if uncompressed_size == 0 { + return Ok(()); + }; + + let mut data; + if uncompressed_size < 12 { + if input.remain < uncompressed_size { + *wait = uncompressed_size; + return Ok(()); + } + data = Vec::with_capacity(uncompressed_size); + input.read_exact_vec(&mut data, uncompressed_size); + } else if let Some(d) = self.read_data(input, wait) { + let mut reader = ZlibReader::new(self.zlibs[stream as usize].take().unwrap(), &d); + data = Vec::with_capacity(uncompressed_size); + unsafe { data.set_len(uncompressed_size) } + reader.read_exact(&mut data)?; + self.zlibs[stream as usize] = Some(reader.into_inner()?); + } else { + return Ok(()); + } + let mut image = Vec::with_capacity(uncompressed_size / 3 * 4); + let mut j = 0; + while j < uncompressed_size { + image.extend_from_slice(&data[j..j + 3]); + image.push(255); + j += 3; + } + + out.push(ImageData { + type_: ImageType::Raw, + x: rect.x, + y: rect.y, + width: rect.width, + height: rect.height, + data: image, + }); + + Ok(()) + } + + fn palette_filter( + &mut self, + stream: u8, + _format: &PixelFormat, + rect: &VncRect, + input: &mut StreamReader, + wait: &mut usize, + out: &mut Vec, + ) -> Result<()> { + if self.num_colors == 0 { + if input.remain < 1 { + *wait = 1; + return Ok(()); + } + self.num_colors = input.read_u8() as usize + 1; + } + + let palette_size = self.num_colors * 3; + + if self.palette.is_empty() { + if input.remain < palette_size { + *wait = palette_size; + return Ok(()); + } + input.read_exact_vec(&mut self.palette, palette_size); + } + + let bpp = if self.num_colors <= 2 { 1 } else { 8 }; + let row_size = (rect.width as usize * bpp + 7) / 8; + let uncompressed_size = rect.height as usize * row_size; + + if uncompressed_size == 0 { + return Ok(()); + } + + let mut data; + if uncompressed_size < 12 { + if input.remain < uncompressed_size { + *wait = uncompressed_size; + return Ok(()); + } + data = Vec::with_capacity(uncompressed_size); + input.read_exact_vec(&mut data, uncompressed_size); + } else if let Some(d) = self.read_data(input, wait) { + let mut reader = ZlibReader::new(self.zlibs[stream as usize].take().unwrap(), &d); + data = Vec::with_capacity(uncompressed_size); + unsafe { data.set_len(uncompressed_size) } + reader.read_exact(&mut data)?; + self.zlibs[stream as usize] = Some(reader.into_inner()?); + } else { + return Ok(()); + } + + if self.num_colors == 2 { + self.mono_rect(data, rect, out)? + } else { + self.palette_rect(data, rect, out)? + } + + self.num_colors = 0; + self.palette.clear(); + + Ok(()) + } + + fn mono_rect(&mut self, data: Vec, rect: &VncRect, out: &mut Vec) -> Result<()> { + // Convert indexed (palette based) image data to RGB + // TODO: reduce number of calculations inside loop + let total = rect.width as usize * rect.height as usize * 4; + let mut image = Vec::with_capacity(total); + unsafe { image.set_len(total) } + + let w = (rect.width as usize + 7) / 8; + let w1 = rect.width as usize / 8; + + for y in 0..rect.height as usize { + let mut dp; + let mut sp; + for x in 0..w1 { + for b in (0..=7).rev() { + dp = (y * rect.width as usize + x * 8 + 7 - b) * 4; + sp = (data[y * w + x] as usize >> b & 1) * 3; + image[dp] = self.palette[sp]; + image[dp + 1] = self.palette[sp + 1]; + image[dp + 2] = self.palette[sp + 2]; + image[dp + 3] = 255; + } + } + let x = w1; + let mut b = 7; + while b >= 8 - rect.width as usize % 8 { + dp = (y * rect.width as usize + x * 8 + 7 - b) * 4; + sp = (data[y * w + x] as usize >> b & 1) * 3; + image[dp] = self.palette[sp]; + image[dp + 1] = self.palette[sp + 1]; + image[dp + 2] = self.palette[sp + 2]; + image[dp + 3] = 255; + b -= 1; + } + } + out.push(ImageData { + type_: ImageType::Raw, + x: rect.x, + y: rect.y, + width: rect.width, + height: rect.height, + data: image, + }); + Ok(()) + } + + fn palette_rect( + &mut self, + data: Vec, + rect: &VncRect, + out: &mut Vec, + ) -> Result<()> { + // Convert indexed (palette based) image data to RGB + let total = rect.width as usize * rect.height as usize * 4; + let mut image = Vec::with_capacity(total); + let mut i = 0; + let mut j = 0; + while i < total { + let sp = data[j] as usize * 3; + image.extend_from_slice(&self.palette[sp..sp + 3]); + image.push(255); + i += 4; + j += 1; + } + out.push(ImageData { + type_: ImageType::Raw, + x: rect.x, + y: rect.y, + width: rect.width, + height: rect.height, + data: image, + }); + Ok(()) + } +} diff --git a/webvnc/src/vnc/decoder/zlib.rs b/webvnc/src/vnc/decoder/zlib.rs new file mode 100644 index 0000000..1b14220 --- /dev/null +++ b/webvnc/src/vnc/decoder/zlib.rs @@ -0,0 +1,77 @@ +/* +Copyright (c) 2016 whitequark + +Permission is hereby granted, free of charge, to any +person obtaining a copy of this software and associated +documentation files (the "Software"), to deal in the +Software without restriction, including without +limitation the rights to use, copy, modify, merge, +publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software +is furnished to do so, subject to the following +conditions: + +The above copyright notice and this permission notice +shall be included in all copies or substantial portions +of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF +ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED +TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A +PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT +SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY +CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION +OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR +IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +DEALINGS IN THE SOFTWARE. +*/ +use std::io::{Read, Result}; + +pub struct ZlibReader<'a> { + decompressor: flate2::Decompress, + input: &'a [u8], +} + +impl<'a> ZlibReader<'a> { + pub fn new(decompressor: flate2::Decompress, input: &'a [u8]) -> ZlibReader<'a> { + ZlibReader { + decompressor, + input, + } + } + + pub fn into_inner(self) -> Result { + if self.input.is_empty() { + Ok(self.decompressor) + } else { + Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "leftover zlib byte data", + )) + } + // Ok(self.decompressor) + } +} + +impl<'a> Read for ZlibReader<'a> { + fn read(&mut self, output: &mut [u8]) -> std::io::Result { + let in_before = self.decompressor.total_in(); + let out_before = self.decompressor.total_out(); + let result = + self.decompressor + .decompress(self.input, output, flate2::FlushDecompress::None); + let consumed = (self.decompressor.total_in() - in_before) as usize; + let produced = (self.decompressor.total_out() - out_before) as usize; + + self.input = &self.input[consumed..]; + match result { + Ok(flate2::Status::Ok) => Ok(produced), + Ok(flate2::Status::BufError) => Ok(0), + Err(error) => Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error)), + Ok(flate2::Status::StreamEnd) => Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "zlib stream end", + )), + } + } +} diff --git a/webvnc/src/vnc/zlib.rs b/webvnc/src/vnc/decoder/zrle.rs similarity index 84% rename from webvnc/src/vnc/zlib.rs rename to webvnc/src/vnc/decoder/zrle.rs index a03f7b5..756664d 100644 --- a/webvnc/src/vnc/zlib.rs +++ b/webvnc/src/vnc/decoder/zrle.rs @@ -26,59 +26,12 @@ IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. */ -use super::vnc_impl::*; -use super::{ImageData, ImageType}; +use super::super::vnc_impl::*; +use super::super::{ImageData, ImageType, StreamReader}; +use super::zlib::ZlibReader; use byteorder::ReadBytesExt; use std::io::{Read, Result}; -struct ZlibReader<'a> { - decompressor: flate2::Decompress, - input: &'a [u8], -} - -impl<'a> ZlibReader<'a> { - fn new(decompressor: flate2::Decompress, input: &'a [u8]) -> ZlibReader<'a> { - ZlibReader { - decompressor, - input, - } - } - - fn into_inner(self) -> Result { - if self.input.is_empty() { - Ok(self.decompressor) - } else { - Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "leftover ZRLE byte data", - )) - } - } -} - -impl<'a> Read for ZlibReader<'a> { - fn read(&mut self, output: &mut [u8]) -> std::io::Result { - let in_before = self.decompressor.total_in(); - let out_before = self.decompressor.total_out(); - let result = - self.decompressor - .decompress(self.input, output, flate2::FlushDecompress::None); - let consumed = (self.decompressor.total_in() - in_before) as usize; - let produced = (self.decompressor.total_out() - out_before) as usize; - - self.input = &self.input[consumed..]; - match result { - Ok(flate2::Status::Ok) => Ok(produced), - Ok(flate2::Status::BufError) => Ok(0), - Err(error) => Err(std::io::Error::new(std::io::ErrorKind::InvalidData, error)), - Ok(flate2::Status::StreamEnd) => Err(std::io::Error::new( - std::io::ErrorKind::InvalidData, - "ZRLE stream end", - )), - } - } -} - struct BitReader { reader: T, buffer: u8, @@ -150,6 +103,7 @@ impl Read for BitReader { } pub struct Decoder { + padding_len: u32, decompressor: Option, } @@ -157,6 +111,7 @@ impl Decoder { pub fn new() -> Decoder { Decoder { decompressor: Some(flate2::Decompress::new(/*zlib_header*/ true)), + padding_len: 0, } } @@ -164,8 +119,9 @@ impl Decoder { &mut self, format: &PixelFormat, rect: &VncRect, - input: &[u8], - ) -> Result> { + input: &mut StreamReader, + out_wait: &mut usize, + ) -> Result>> { fn read_run_length(reader: &mut dyn Read) -> Result { let mut run_length_part; let mut run_length = 1; @@ -198,7 +154,23 @@ impl Decoder { pixels.extend_from_slice(&palette[start..start + bpp]) } + if self.padding_len == 0 { + if input.remain < 4 { + *out_wait = 4; + return Ok(None); + } else { + self.padding_len = input.read_u32(); + } + } + + if input.remain < self.padding_len as usize { + *out_wait = self.padding_len as usize; + return Ok(None); + } + let mut ret = Vec::new(); + let mut image_data = Vec::with_capacity(self.padding_len as usize); + input.read_exact_vec(&mut image_data, self.padding_len as usize); let bpp = format.bits_per_pixel as usize / 8; let pixel_mask = (format.red_max as u32) << format.red_shift @@ -219,7 +191,10 @@ impl Decoder { }; let mut palette = Vec::with_capacity(128 * bpp); - let mut reader = BitReader::new(ZlibReader::new(self.decompressor.take().unwrap(), input)); + let mut reader = BitReader::new(ZlibReader::new( + self.decompressor.take().unwrap(), + &image_data, + )); let mut y = 0; while y < rect.height { @@ -338,6 +313,7 @@ impl Decoder { } self.decompressor = Some(reader.into_inner()?.into_inner()?); - Ok(ret) + self.padding_len = 0; + Ok(Some(ret)) } } diff --git a/webvnc/src/vnc/mod.rs b/webvnc/src/vnc/mod.rs index a7b2d61..9901a4e 100644 --- a/webvnc/src/vnc/mod.rs +++ b/webvnc/src/vnc/mod.rs @@ -1,8 +1,8 @@ +mod decoder; mod des; mod vnc_impl; mod x11cursor; mod x11keyboard; -mod zlib; pub enum MouseEventType { Down, @@ -26,6 +26,7 @@ pub enum ImageType { Raw, Copy, Fill, + Jpeg, } pub enum VncOutput { diff --git a/webvnc/src/vnc/vnc_impl.rs b/webvnc/src/vnc/vnc_impl.rs index 3c25433..64b12b3 100644 --- a/webvnc/src/vnc/vnc_impl.rs +++ b/webvnc/src/vnc/vnc_impl.rs @@ -1,5 +1,5 @@ use super::*; -use super::{des, x11cursor::MouseUtils, x11keyboard, zlib, MouseEventType}; +use super::{decoder, des, x11cursor::MouseUtils, x11keyboard, MouseEventType}; use crate::{console_log, log}; const VNC_RFB33: &[u8; 12] = b"RFB 003.003\n"; @@ -31,6 +31,7 @@ pub enum VncEncoding { CopyRect = 1, RRE = 2, Hextile = 5, + Tight = 7, TRLE = 15, ZRLE = 16, CursorPseudo = -239, @@ -84,7 +85,9 @@ pub struct Vnc { padding_rect: Option, outbuf: Vec, outs: Vec, - zlib_decoder: zlib::Decoder, + raw: Option, + zrle: Option, + tight: Option, } impl Vnc { @@ -92,6 +95,7 @@ impl Vnc { Self { state: VncState::Init, supported_encodings: vec![ + VncEncoding::Tight, VncEncoding::ZRLE, VncEncoding::Raw, VncEncoding::CopyRect, @@ -115,11 +119,16 @@ impl Vnc { padding_rect: None, outbuf: Vec::with_capacity(128), outs: Vec::with_capacity(10), - zlib_decoder: zlib::Decoder::new(), + raw: None, + zrle: None, + tight: None, } } pub fn do_input(&mut self, input: Vec) { + if let VncState::Disconnected = self.state { + return; + } // console_log!( // "VNC input {}, left {}, require {}", // input.len(), @@ -130,6 +139,9 @@ impl Vnc { while self.reader.remain() >= self.require { self.handle_input(); // console_log!("left {}, require {}", self.reader.remain(), self.require); + if let VncState::Disconnected = self.state { + break; + } } } @@ -231,46 +243,18 @@ impl Vnc { #[allow(dead_code)] impl Vnc { - fn read_u8(&mut self) -> u8 { - self.reader.read_u8() - } - - fn read_u16(&mut self) -> u16 { - self.reader.read_u16() - } - - fn read_u32(&mut self) -> u32 { - self.reader.read_u32() - } - - fn read_u64(&mut self) -> u64 { - self.reader.read_u64() - } - - fn read_exact_vec(&mut self, out_vec: &mut Vec, len: usize) { - self.reader.read_exact_vec(out_vec, len); - } - - fn read_exact(&mut self, buf: &mut [u8], len: usize) { - self.reader.read_exact(buf, len) - } - - fn read_string(&mut self, len: usize) -> String { - self.reader.read_string(len) - } - fn read_string_l8(&mut self) -> String { - let len = self.read_u8() as usize; + let len = self.reader.read_u8() as usize; self.reader.read_string(len) } fn read_string_l16(&mut self) -> String { - let len = self.read_u16() as usize; + let len = self.reader.read_u16() as usize; self.reader.read_string(len) } fn read_string_l32(&mut self) -> String { - let len = self.read_u32() as usize; + let len = self.reader.read_u32() as usize; self.reader.read_string(len) } } @@ -414,7 +398,7 @@ impl Vnc { fn do_handshake(&mut self) { let mut rfbversion: [u8; 12] = [0; 12]; - self.read_exact(&mut rfbversion, 12); + self.reader.read_exact(&mut rfbversion, 12); let support_version = match &rfbversion { b"RFB 003.003\n" => Ok(VNC_RFB33), b"RFB 003.007\n" => Ok(VNC_RFB33), @@ -435,7 +419,7 @@ impl Vnc { fn do_authenticate(&mut self) { // console_log!("VNC: do_authenticate {}", self.reader.remain()); if self.security_type == SecurityType::Invalid { - let auth_type = self.read_u32(); + let auth_type = self.reader.read_u32(); match auth_type { 1 => { self.security_type = SecurityType::None; @@ -450,14 +434,14 @@ impl Vnc { } } else { let mut challenge = [0u8; 16]; - self.read_exact(&mut challenge, 16); + self.reader.read_exact(&mut challenge, 16); self.challenge = challenge; self.outs.push(VncOutput::RequirePassword); } } fn handle_auth_result(&mut self) { - let response = self.read_u32(); + let response = self.reader.read_u32(); console_log!("Auth resp {}", response); match response { 0 => self.send_client_initilize(), @@ -476,10 +460,10 @@ impl Vnc { // 4 CARD32 name-length // name-length CARD8 array name-string fn handle_server_init(&mut self) { - self.width = self.read_u16(); - self.height = self.read_u16(); + self.width = self.reader.read_u16(); + self.height = self.reader.read_u16(); let mut pfb: [u8; 16] = [0u8; 16]; - self.read_exact(&mut pfb, 16); + self.reader.read_exact(&mut pfb, 16); // This pixel format will be used unless the client requests a different format using the SetPixelFormat message self.pf = (&pfb).into(); console_log!( @@ -501,7 +485,7 @@ impl Vnc { ServerMessage::ServerCutText => self.read_cut_text(), ServerMessage::FramebufferUpdate => self.read_rect(), ServerMessage::None => { - let msg_type = self.read_u8(); + let msg_type = self.reader.read_u8(); match msg_type { 0 => self.handle_framebuffer_update(), @@ -520,8 +504,8 @@ impl Vnc { // 2 CARD16 number-of-rectangles // This is followed by number-of-rectanglesrectangles of pixel data. fn handle_framebuffer_update(&mut self) { - let _padding = self.read_u8(); - self.num_rect_left = self.read_u16(); + let _padding = self.reader.read_u8(); + self.num_rect_left = self.reader.read_u16(); // console_log!("VNC: {} rects", self.num_rect_left); self.require = 12; // the length of the first rectangle hdr self.msg_handling = ServerMessage::FramebufferUpdate; @@ -541,16 +525,17 @@ impl Vnc { fn read_rect(&mut self) { if self.padding_rect.is_none() { // a brand new rectangle - let x = self.read_u16(); - let y = self.read_u16(); - let width = self.read_u16(); - let height = self.read_u16(); - let encoding_type = self.read_u32().into(); + let x = self.reader.read_u16(); + let y = self.reader.read_u16(); + let width = self.reader.read_u16(); + let height = self.reader.read_u16(); + let encoding_type = self.reader.read_u32().into(); match encoding_type { VncEncoding::Raw => self.handle_raw_encoding(x, y, width, height), VncEncoding::CopyRect => self.handle_copy_rect_encoding(x, y, width, height), VncEncoding::RRE => self.handle_rre_encoding(x, y, width, height), VncEncoding::Hextile => self.handle_hextile_encoding(x, y, width, height), + VncEncoding::Tight => self.handle_tight_encoding(x, y, width, height), VncEncoding::TRLE => self.handle_trle_encoding(x, y, width, height), VncEncoding::ZRLE => self.handle_zrle_encoding(x, y, width, height), VncEncoding::CursorPseudo => { @@ -561,42 +546,36 @@ impl Vnc { } } } else { - if let VncEncoding::ZRLE = self.padding_rect.as_ref().unwrap().encoding_type { - if self.require == 4 { - self.require = self.read_u32() as usize; - return; - } - } - // we now have an entire rectangle - let rect = self.padding_rect.take().unwrap(); - let mut image_data: Vec = Vec::with_capacity(self.require); - self.read_exact_vec(&mut image_data, self.require); + let rect = self.padding_rect.as_ref().unwrap(); match rect.encoding_type { VncEncoding::Raw => { // raw - let mut y = 0; - let mut x = 0; + let mut decoder = if self.raw.is_none() { + decoder::RawDecoder::new() + } else { + self.raw.take().unwrap() + }; - while y < rect.height { - while x < rect.width { - let idx = (y as usize * rect.width as usize + x as usize) * 4; - image_data.swap(idx, idx + 2); - x += 1; + match decoder.decode(&self.pf, rect, &mut self.reader, &mut self.require) { + Ok(images) => { + if let Some(images) = images { + for i in images { + self.outs.push(VncOutput::RenderImage(i)); + } + self.read_rect_end(); + } + } + Err(e) => { + self.disconnect_with_err(&format!("{:?}", e)); } - x = 0; - y += 1; } - self.outs.push(VncOutput::RenderImage(ImageData { - type_: ImageType::Raw, - x: rect.x, - y: rect.y, - width: rect.width, - height: rect.height, - data: image_data, - })); + + self.raw = Some(decoder); } VncEncoding::CopyRect => { // copy rectangle + let mut image_data: Vec = Vec::with_capacity(self.require); + self.reader.read_exact_vec(&mut image_data, self.require); self.outs.push(VncOutput::RenderImage(ImageData { type_: ImageType::Copy, x: rect.x, @@ -605,29 +584,69 @@ impl Vnc { height: rect.height, data: image_data, })); + self.read_rect_end(); } - VncEncoding::ZRLE => { - // ZRLE - match self.zlib_decoder.decode(&self.pf, &rect, &image_data) { + VncEncoding::Tight => { + // Tight + let mut decoder = if self.tight.is_none() { + decoder::TightDecoder::new() + } else { + self.tight.take().unwrap() + }; + + match decoder.decode(&self.pf, rect, &mut self.reader, &mut self.require) { Ok(images) => { - for i in images { - self.outs.push(VncOutput::RenderImage(i)); + if let Some(images) = images { + for i in images { + self.outs.push(VncOutput::RenderImage(i)); + } + self.read_rect_end(); } } Err(e) => { self.disconnect_with_err(&format!("{:?}", e)); } } + + self.tight = Some(decoder); + } + VncEncoding::ZRLE => { + // ZRLE + let mut decoder = if self.zrle.is_none() { + decoder::ZrleDecoder::new() + } else { + self.zrle.take().unwrap() + }; + + match decoder.decode(&self.pf, rect, &mut self.reader, &mut self.require) { + Ok(images) => { + if let Some(images) = images { + for i in images { + self.outs.push(VncOutput::RenderImage(i)); + } + self.read_rect_end(); + } + } + Err(e) => { + self.disconnect_with_err(&format!("{:?}", e)); + } + } + + self.zrle = Some(decoder); } _ => unimplemented!("Unknow encoding {}", rect.encoding_type as u32), } - self.num_rect_left -= 1; - if 0 == self.num_rect_left { - self.msg_handling = ServerMessage::None; - self.require = 1; // any message from sever will be handled - } else { - self.require = 12; // the length of the next rectangle hdr - } + } + } + + fn read_rect_end(&mut self) { + self.padding_rect.take(); + self.num_rect_left -= 1; + if 0 == self.num_rect_left { + self.msg_handling = ServerMessage::None; + self.require = 1; // any message from sever will be handled + } else { + self.require = 12; // the length of the next rectangle hdr } } @@ -644,9 +663,9 @@ impl Vnc { // 2 CARD16 first-colour // 2 CARD16 number-of-colours fn handle_set_colour_map(&mut self) { - let _padding = self.read_u8(); - let _first_colour = self.read_u16(); - self.require = self.read_u16() as usize * 6; + let _padding = self.reader.read_u8(); + let _first_colour = self.reader.read_u16(); + self.require = self.reader.read_u16() as usize * 6; self.msg_handling = ServerMessage::SetColourMap; } @@ -664,7 +683,7 @@ impl Vnc { // just consume the data let mut v = Vec::with_capacity(self.require); - self.read_exact_vec(&mut v, self.require); + self.reader.read_exact_vec(&mut v, self.require); self.require = 1; self.msg_handling = ServerMessage::None; } @@ -684,15 +703,15 @@ impl Vnc { // length CARD8 array text fn handle_server_cut_text(&mut self) { for _ in 0..3 { - self.read_u8(); + self.reader.read_u8(); } - self.require = self.read_u32() as usize; + self.require = self.reader.read_u32() as usize; self.msg_handling = ServerMessage::ServerCutText; console_log!("VNC: ServerCutText {} bytes", self.require); } fn read_cut_text(&mut self) { - let text = self.read_string(self.require); + let text = self.reader.read_string(self.require); self.require = 1; self.msg_handling = ServerMessage::None; self.outs.push(VncOutput::SetClipboard(text)); @@ -734,6 +753,17 @@ impl Vnc { unimplemented!() } + fn handle_tight_encoding(&mut self, x: u16, y: u16, width: u16, height: u16) { + self.require = 1; + self.padding_rect = Some(VncRect { + x, + y, + width, + height, + encoding_type: VncEncoding::Tight, + }) + } + fn handle_trle_encoding(&mut self, _x: u16, _y: u16, _width: u16, _height: u16) { unimplemented!() }