use std::io::{BufReader, BufWriter}; use std::net::TcpStream; use crate::codec::{hextile, iip, raw_tile, tight}; use crate::framebuffer::Framebuffer; use crate::handshake::{self, Config, ServerInit}; use crate::msg::{self, ServerMsg}; use crate::proto::{self, RectHeader, read_exact, read_u8}; #[derive(Debug, thiserror::Error)] pub enum SessionError { #[error("handshake: {0}")] Handshake(#[from] handshake::HandshakeError), #[error("protocol: {0}")] Proto(#[from] proto::ProtoError), #[error("io: {0}")] Io(#[from] std::io::Error), #[error("unsupported encoding: {0}")] UnsupportedEncoding(i32), #[error("unsupported message type: {0}")] UnsupportedMessage(u8), } /// Events emitted by the session pump to consumers. #[derive(Debug)] pub enum Event { /// A region of the framebuffer was updated. FramebufferDirty, /// The framebuffer was resized. Resize { width: u16, height: u16 }, /// Bell from server. Bell, /// Server sent a debug string. Debug(String), /// Server sent an RFB command (key, value). RfbCommand(String, String), /// Server name updated. NameUpdate(String), } /// Active protocol session with framebuffer. pub struct ActiveSession { pub framebuffer: Framebuffer, pub server_name: String, pub reader: BufReader, pub writer: BufWriter, pub server_init: ServerInit, zlib: tight::ZlibStreams, tile_cache: iip::TileCache, } impl ActiveSession { /// Connect, handshake, send SetEncodings + initial FBUpdateRequest. pub fn connect(cfg: &Config, encodings: &[i32]) -> Result { let raw = handshake::connect(cfg)?; let w = raw.server_init.width; let h = raw.server_init.height; let mut session = Self { framebuffer: Framebuffer::new(w, h), server_name: raw.server_name, reader: raw.reader, writer: raw.writer, server_init: raw.server_init, zlib: tight::ZlibStreams::new(), tile_cache: iip::TileCache::new(w, h), }; // Tell server to send 8bpp RGB332 pixels msg::write_set_pixel_format_rgb332(&mut session.writer)?; // Send SetEncodings msg::write_set_encodings(&mut session.writer, encodings)?; // Request full non-incremental framebuffer update msg::write_fb_update_request( &mut session.writer, 0, 0, session.framebuffer.width, session.framebuffer.height, false, )?; Ok(session) } /// Request an incremental framebuffer update. pub fn request_update(&mut self) -> Result<(), SessionError> { msg::write_fb_update_request( &mut self.writer, 0, 0, self.framebuffer.width, self.framebuffer.height, true, )?; Ok(()) } /// Process one server message. Returns an event if meaningful to the consumer. pub fn process_one(&mut self) -> Result, SessionError> { let msg_type = read_u8(&mut self.reader)?; let msg = ServerMsg::from(msg_type); match msg { ServerMsg::FramebufferUpdate => { self.handle_fb_update()?; Ok(Some(Event::FramebufferDirty)) } ServerMsg::Bell => Ok(Some(Event::Bell)), ServerMsg::Ping => { let payload = msg::read_ping(&mut self.reader)?; msg::write_ping_response(&mut self.writer, payload)?; Ok(None) } ServerMsg::BandwidthProbe => { msg::write_bandwidth_marker(&mut self.writer, 1)?; msg::read_bandwidth_probe(&mut self.reader)?; msg::write_bandwidth_marker(&mut self.writer, 2)?; Ok(None) } ServerMsg::Ack => { msg::read_ack(&mut self.reader)?; Ok(None) } ServerMsg::DebugString => { let s = msg::read_debug_string(&mut self.reader)?; Ok(Some(Event::Debug(s))) } ServerMsg::RfbCommand => { let (k, v) = msg::read_rfb_command(&mut self.reader)?; Ok(Some(Event::RfbCommand(k, v))) } ServerMsg::ServerCutText => { let _text = msg::read_server_cut_text(&mut self.reader)?; Ok(None) } ServerMsg::ServerNameUpdate => { let name = msg::read_server_name_update(&mut self.reader)?; self.server_name = name.clone(); Ok(Some(Event::NameUpdate(name))) } ServerMsg::LayoutLocale => { let _locale = msg::read_layout_locale(&mut self.reader)?; Ok(None) } ServerMsg::DesktopResize => { // Reads same struct as handshake pixel-format (aw.i, line 519) let _flag = read_u8(&mut self.reader)?; let _depth = proto::read_u16_be(&mut self.reader)?; let label_len = proto::read_u16_be(&mut self.reader)? as usize; let _label = read_exact(&mut self.reader, label_len)?; Ok(None) } ServerMsg::ModeChange => { // Re-read ServerInit (aw.k, line 435) — framebuffer dimensions may change let si = handshake::read_server_init_from(&mut self.reader)?; let old_w = self.framebuffer.width; let old_h = self.framebuffer.height; self.server_init = si; if self.server_init.width != old_w || self.server_init.height != old_h { self.framebuffer .resize(self.server_init.width, self.server_init.height); self.tile_cache .resize(self.server_init.width, self.server_init.height); return Ok(Some(Event::Resize { width: self.server_init.width, height: self.server_init.height, })); } Ok(None) } ServerMsg::RdpEvent => { let _event_type = msg::read_rdp_event(&mut self.reader)?; Ok(None) } ServerMsg::PixelFormatChange => { // aw.e(), line 537: 1 pad + 4×u8 + 8×u16 = 21 bytes let _data = read_exact(&mut self.reader, 21)?; Ok(None) } ServerMsg::SetColourMapEntries => Err(SessionError::UnsupportedMessage(msg_type)), ServerMsg::Unknown(t) => Err(SessionError::UnsupportedMessage(t)), } } fn handle_fb_update(&mut self) -> Result<(), SessionError> { let num_rects = msg::read_fb_update_header(&mut self.reader)?; for _ in 0..num_rects { let hdr = RectHeader::read_from(&mut self.reader)?; match hdr.encoding { 0 => { // Raw: read w*h bytes let size = hdr.w as usize * hdr.h as usize; let data = read_exact(&mut self.reader, size)?; self.framebuffer .apply_raw(hdr.x, hdr.y, hdr.w, hdr.h, &data); } 1 => { // CopyRect: read src_x, src_y (u16 each) let src_x = proto::read_u16_be(&mut self.reader)?; let src_y = proto::read_u16_be(&mut self.reader)?; self.framebuffer .copy_rect(src_x, src_y, hdr.x, hdr.y, hdr.w, hdr.h); } 5 => { hextile::decode_hextile( &mut self.reader, &mut self.framebuffer, hdr.x, hdr.y, hdr.w, hdr.h, )?; } 7 => { // Tight tight::decode_tight( &mut self.reader, &mut self.framebuffer, &mut self.zlib, hdr.x, hdr.y, hdr.w, hdr.h, )?; } 9 => { iip::decode_iip( &mut self.reader, &mut self.framebuffer, &mut self.tile_cache, &mut self.zlib, hdr.x, hdr.y, hdr.w, hdr.h, )?; } 10 => { raw_tile::decode_raw_tile( &mut self.reader, &mut self.framebuffer, hdr.x, hdr.y, hdr.w, hdr.h, )?; } other => { return Err(SessionError::UnsupportedEncoding(other)); } } } Ok(()) } }