blekin/crates/ericrfb/src/session.rs

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<TcpStream>,
    pub writer: BufWriter<TcpStream>,
    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<Self, SessionError> {
        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<Option<Event>, 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(())
    }
}