blekin/crates/ericrfb/src/handshake.rs

use std::io::{BufReader, BufWriter, Read, Write};
use std::net::TcpStream;

use thiserror::Error;

use crate::proto::{self, read_exact, read_i32_be, read_modified_utf8, read_u8, read_u16_be};

// ---------------------------------------------------------------------------
// Errors
// ---------------------------------------------------------------------------

#[derive(Debug, Error)]
pub enum HandshakeError {
    #[error("protocol error: {0}")]
    Protocol(#[from] proto::ProtoError),
    #[error("i/o error: {0}")]
    Io(#[from] std::io::Error),
    #[error("auth rejected: {0}")]
    AuthRejected(String),
    #[error("invalid server banner")]
    InvalidBanner,
}

/// Map error status codes from aw.a(int), line 350.
fn auth_error_message(code: i32) -> &'static str {
    match code {
        1 => "no permission",
        2 => "exclusive access active",
        3 => "manually rejected",
        4 => "server password disabled",
        5 => "loopback connection is senseless",
        6 => "authentication failed",
        7 => "access to this kvm port denied",
        _ => "unknown error",
    }
}

// ---------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct Config {
    pub host: String,
    pub port: u16,
    pub applet_id: String,
    pub protocol_version: String,
    pub port_id: u8,
    pub shared: bool,
}

impl Config {
    pub fn new(host: impl Into<String>, port: u16, applet_id: impl Into<String>) -> Self {
        Self {
            host: host.into(),
            port,
            applet_id: applet_id.into(),
            protocol_version: "01.11".into(),
            port_id: 0,
            shared: true,
        }
    }
}

// ---------------------------------------------------------------------------
// Pixel format (from aw.i(), line 519)
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct PixelFormat {
    pub exclusive: bool,
    pub color_depth: u16,
    pub label: String,
}

fn read_pixel_format(r: &mut impl Read) -> Result<PixelFormat, HandshakeError> {
    let flag = read_u8(r)?;
    let color_depth = read_u16_be(r)?;
    let label_len = read_u16_be(r)? as usize;
    let label_bytes = read_exact(r, label_len)?;
    let label = String::from_utf8_lossy(&label_bytes).into_owned();
    Ok(PixelFormat {
        exclusive: flag == 1,
        color_depth,
        label,
    })
}

// ---------------------------------------------------------------------------
// ServerInit (from aw.k(), line 435)
// ---------------------------------------------------------------------------

#[derive(Debug, Clone)]
pub struct ServerInit {
    pub supports_resize: bool,
    pub width: u16,
    pub height: u16,
    pub bits_per_pixel: u8,
    pub depth: u8,
    pub big_endian: bool,
    pub true_color: bool,
    pub red_max: u16,
    pub green_max: u16,
    pub blue_max: u16,
    pub red_shift: u8,
    pub green_shift: u8,
    pub blue_shift: u8,
}

fn read_server_init(r: &mut impl Read) -> Result<ServerInit, HandshakeError> {
    let supports_resize = read_u8(r)? != 0;
    let width = read_u16_be(r)?;
    let height = read_u16_be(r)?;
    let bits_per_pixel = read_u8(r)?;
    let depth = read_u8(r)?;
    let big_endian = read_u8(r)? != 0;
    let true_color = read_u8(r)? != 0;
    let red_max = read_u16_be(r)?;
    let green_max = read_u16_be(r)?;
    let blue_max = read_u16_be(r)?;
    let red_shift = read_u8(r)?;
    let green_shift = read_u8(r)?;
    let blue_shift = read_u8(r)?;
    let _pad = read_exact(r, 3)?;
    Ok(ServerInit {
        supports_resize,
        width,
        height,
        bits_per_pixel,
        depth,
        big_endian,
        true_color,
        red_max,
        green_max,
        blue_max,
        red_shift,
        green_shift,
        blue_shift,
    })
}

// ---------------------------------------------------------------------------
// Session — returned after successful handshake
// ---------------------------------------------------------------------------

#[derive(Debug)]
pub struct Session {
    pub server_version: (u8, u8),
    pub server_name: String,
    pub pixel_format: PixelFormat,
    pub server_init: ServerInit,
    pub reader: BufReader<TcpStream>,
    pub writer: BufWriter<TcpStream>,
}

impl Session {
    pub fn width(&self) -> u16 {
        self.server_init.width
    }

    pub fn height(&self) -> u16 {
        self.server_init.height
    }
}

// ---------------------------------------------------------------------------
// Handshake — aw.g(), line 226, steps 1–11
// ---------------------------------------------------------------------------

pub fn connect(cfg: &Config) -> Result<Session, HandshakeError> {
    let addr = format!("{}:{}", cfg.host, cfg.port);
    let stream = TcpStream::connect(&addr)?;
    let read_stream = stream.try_clone()?;
    let mut r = BufReader::with_capacity(32768, read_stream);
    let mut w = BufWriter::new(stream);

    // Step 1: C→S  75 bytes auth string, zero-padded, ISO-8859-1
    let auth_str = format!("e-RIC AUTH={}", cfg.applet_id);
    let auth_bytes = auth_str.as_bytes();
    let mut auth_buf = [0u8; 75];
    let copy_len = auth_bytes.len().min(75);
    auth_buf[..copy_len].copy_from_slice(&auth_bytes[..copy_len]);
    w.write_all(&auth_buf)?;
    w.flush()?;

    // Step 2: S→C  1 byte status
    let status = read_u8(&mut r)?;
    if status == 3 {
        let error_code = read_i32_be(&mut r)?;
        return Err(HandshakeError::AuthRejected(
            auth_error_message(error_code).into(),
        ));
    }

    // Step 3: S→C  15 bytes banner "-RIC RFB MM.NN\n"
    // The status byte (101 = 'e') is the first byte of "e-RIC RFB MM.NN\n".
    let banner = read_exact(&mut r, 15)?;
    if banner[0] != b'-'
        || banner[1] != b'R'
        || banner[2] != b'I'
        || banner[3] != b'C'
        || banner[4] != b' '
        || banner[5] != b'R'
        || banner[6] != b'F'
        || banner[7] != b'B'
        || banner[8] != b' '
        || banner[14] != b'\n'
    {
        return Err(HandshakeError::InvalidBanner);
    }
    let major = (banner[9] - b'0') * 10 + (banner[10] - b'0');
    let minor = (banner[12] - b'0') * 10 + (banner[13] - b'0');

    // Step 4: S→C  1 byte sync
    let _sync1 = read_u8(&mut r)?;

    // Step 5: S→C  server name (1 pad + modified-UTF-8 string)
    let _pad = read_u8(&mut r)?;
    let server_name = read_modified_utf8(&mut r)?;

    // Step 6: S→C  1 byte sync
    let _sync2 = read_u8(&mut r)?;

    // Step 7: S→C  pixel format struct (variable length)
    let pixel_format = read_pixel_format(&mut r)?;

    // Step 8: C→S  16 bytes client version "e-RIC RFB 01.11\n"
    let version_str = format!("e-RIC RFB {}\n", cfg.protocol_version);
    let version_bytes = version_str.as_bytes();
    let mut version_buf = [0u8; 16];
    let copy_len = version_bytes.len().min(16);
    version_buf[..copy_len].copy_from_slice(&version_bytes[..copy_len]);
    w.write_all(&version_buf)?;
    w.flush()?;

    // Step 9: C→S  2 bytes [shared_flag, port_id]
    w.write_all(&[if cfg.shared { 1 } else { 0 }, cfg.port_id])?;
    w.flush()?;

    // Step 10: S→C  1 byte sync
    let _sync3 = read_u8(&mut r)?;

    // Step 11: S→C  19 bytes ServerInit
    let server_init = read_server_init(&mut r)?;

    Ok(Session {
        server_version: (major, minor),
        server_name,
        pixel_format,
        server_init,
        reader: r,
        writer: w,
    })
}