// Copyright (c) 2022 Huawei Technologies Co.,Ltd. All rights r&eserved.
//
// StratoVirt is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

use crate::{
    auth::AuthState,
    pixman::{bytes_per_pixel, get_image_height, get_image_width, PixelFormat},
    round_up_div,
    server::{vnc_refresh_notify, VncServer},
    utils::BuffPool,
    vnc::{
        framebuffer_upadate, set_area_dirty, write_pixel, DisplayMouse, BIT_PER_BYTE,
        DIRTY_PIXELS_NUM, DIRTY_WIDTH_BITS, MAX_WINDOW_HEIGHT, MAX_WINDOW_WIDTH, VNC_RECT_INFO,
    },
    VncError,
};
use anyhow::{anyhow, Result};
use log::error;
use rustls::ServerConnection;
use sscanf::scanf;
use std::{
    cmp,
    io::{Read, Write},
    net::{Shutdown, TcpStream},
    os::unix::prelude::{AsRawFd, RawFd},
    sync::{Arc, Mutex},
};
use util::{
    bitmap::Bitmap,
    loop_context::{read_fd, EventNotifier, EventNotifierHelper, NotifierOperation},
};
use vmm_sys_util::epoll::EventSet;
use vmm_sys_util::eventfd::EventFd;

pub const APP_NAME: &str = "stratovirt";
const MAX_RECVBUF_LEN: usize = 1024;
const MAX_SEND_LEN: usize = 64 * 1024;
const NUM_OF_COLORMAP: u16 = 256;

// VNC encodings types.
pub const ENCODING_RAW: i32 = 0;
pub const ENCODING_HEXTILE: i32 = 5;
const ENCODING_ZLIB: i32 = 6;
const ENCODING_TIGHT: i32 = 7;
const ENCODING_ZRLE: i32 = 16;
const ENCODING_ZYWRLE: i32 = 17;
const ENCODING_DESKTOPRESIZE: i32 = -223;
pub const ENCODING_RICH_CURSOR: i32 = -239;
const ENCODING_POINTER_TYPE_CHANGE: i32 = -257;
const ENCODING_LED_STATE: i32 = -261;
const ENCODING_DESKTOP_RESIZE_EXT: i32 = -308;
pub const ENCODING_ALPHA_CURSOR: i32 = -314;
const ENCODING_WMVI: i32 = 1464686185;

/// Image display feature.
pub enum VncFeatures {
    VncFeatureResize,
    VncFeatureResizeExt,
    VncFeatureHextile,
    VncFeaturePointerTypeChange,
    VncFeatureWmvi,
    VncFeatureTight,
    VncFeatureZlib,
    VncFeatureRichCursor,
    VncFeatureAlphaCursor,
    VncFeatureTightPng,
    VncFeatureZrle,
    VncFeatureZywrle,
    VncFeatureLedState,
    VncFeatureXvp,
    VncFeatureClipboardExt,
}

/// Client to server message in Remote Framebuffer Protocol.
pub enum ClientMsg {
    SetPixelFormat = 0,
    SetEncodings = 2,
    FramebufferUpdateRequest = 3,
    KeyEvent = 4,
    PointerEvent = 5,
    ClientCutText = 6,
    InvalidMsg,
}

/// Server to client message in Remote Framebuffer Protocol.
pub enum ServerMsg {
    FramebufferUpdate = 0,
    SetColourMapEntries = 1,
}

impl From<u8> for ClientMsg {
    fn from(v: u8) -> Self {
        match v {
            0 => ClientMsg::SetPixelFormat,
            2 => ClientMsg::SetEncodings,
            3 => ClientMsg::FramebufferUpdateRequest,
            4 => ClientMsg::KeyEvent,
            5 => ClientMsg::PointerEvent,
            6 => ClientMsg::ClientCutText,
            _ => ClientMsg::InvalidMsg,
        }
    }
}

/// RFB protocol version.
#[derive(Clone)]
pub struct VncVersion {
    pub major: u16,
    pub minor: u16,
}

impl VncVersion {
    pub fn new(major: u16, minor: u16) -> Self {
        VncVersion { major, minor }
    }
}

impl Default for VncVersion {
    fn default() -> Self {
        Self::new(0, 0)
    }
}

#[derive(PartialEq)]
pub enum UpdateState {
    No,
    Incremental,
    Force,
}

/// Dirty area of image
#[derive(Clone)]
pub struct Rectangle {
    pub x: i32,
    pub y: i32,
    pub w: i32,
    pub h: i32,
}

impl Rectangle {
    pub fn new(x: i32, y: i32, w: i32, h: i32) -> Self {
        Rectangle { x, y, w, h }
    }
}

/// Display Output mode information of client.
#[derive(Clone)]
pub struct DisplayMode {
    /// Encoding type.
    pub enc: i32,
    /// Data storage type for client.
    pub client_be: bool,
    /// The pixel need to convert.
    pub convert: bool,
    /// Image pixel format in pixman.
    pub pf: PixelFormat,
}

impl DisplayMode {
    pub fn new(enc: i32, client_be: bool, convert: bool, pf: PixelFormat) -> Self {
        DisplayMode {
            enc,
            client_be,
            convert,
            pf,
        }
    }
}

impl Default for DisplayMode {
    fn default() -> Self {
        Self::new(0, false, false, PixelFormat::default())
    }
}

unsafe impl Send for RectInfo {}
pub struct RectInfo {
    /// Vnc client state.
    pub client: Arc<ClientState>,
    /// Dirty area of image.
    pub rects: Vec<Rectangle>,
}

impl RectInfo {
    pub fn new(client: &Arc<ClientState>, rects: Vec<Rectangle>) -> Self {
        RectInfo {
            client: client.clone(),
            rects,
        }
    }
}

impl Clone for RectInfo {
    fn clone(&self) -> Self {
        let mut rects = Vec::new();
        for rect in &self.rects {
            rects.push(rect.clone());
        }
        Self {
            client: self.client.clone(),
            rects,
        }
    }

    fn clone_from(&mut self, source: &Self) {
        *self = source.clone()
    }
}

/// The connection state of vnc client.
pub struct ConnState {
    /// Connection status.
    pub dis_conn: bool,
    /// State flags whether the image needs to be updated for the client.
    update_state: UpdateState,
    /// RFB protocol version.
    pub version: VncVersion,
    /// Vnc display feature.
    feature: i32,
}

impl Default for ConnState {
    fn default() -> Self {
        ConnState {
            dis_conn: false,
            update_state: UpdateState::No,
            version: VncVersion::default(),
            feature: 0,
        }
    }
}

impl ConnState {
    pub fn is_disconnect(&mut self) -> bool {
        self.dis_conn
    }

    pub fn has_feature(&self, feature: VncFeatures) -> bool {
        self.feature & (1 << feature as usize) != 0
    }
}

/// Struct to record the state with the vnc client.
pub struct ClientState {
    /// Write event fd.
    write_fd: Arc<Mutex<EventFd>>,
    /// TcpStream receive buffer.
    pub in_buffer: Arc<Mutex<BuffPool>>,
    /// TcpStream write buffer.
    pub out_buffer: Arc<Mutex<BuffPool>>,
    /// Output mode information of client display.
    pub client_dpm: Arc<Mutex<DisplayMode>>,
    /// The connection state of vnc client.
    pub conn_state: Arc<Mutex<ConnState>>,
    /// Identify the image update area.
    pub dirty_bitmap: Arc<Mutex<Bitmap<u64>>>,
}

impl Default for ClientState {
    fn default() -> Self {
        ClientState {
            write_fd: Arc::new(Mutex::new(EventFd::new(libc::EFD_NONBLOCK).unwrap())),
            in_buffer: Arc::new(Mutex::new(BuffPool::new())),
            out_buffer: Arc::new(Mutex::new(BuffPool::new())),
            client_dpm: Arc::new(Mutex::new(DisplayMode::default())),
            conn_state: Arc::new(Mutex::new(ConnState::default())),
            dirty_bitmap: Arc::new(Mutex::new(Bitmap::<u64>::new(
                MAX_WINDOW_HEIGHT as usize
                    * round_up_div(DIRTY_WIDTH_BITS as u64, u64::BITS as u64) as usize,
            ))),
        }
    }
}

/// Handle the message with vnc client.
pub struct ClientIoHandler {
    /// TcpStream connected with client.
    pub stream: TcpStream,
    /// Tls server connection.
    pub tls_conn: Option<rustls::ServerConnection>,
    /// Message handler.
    pub msg_handler: fn(&mut ClientIoHandler) -> Result<()>,
    /// Size of buff in next handle.
    pub expect: usize,
    /// State with vnc client.
    pub client: Arc<ClientState>,
    /// Configure for vnc server.
    pub server: Arc<VncServer>,
    /// Disconnect event fd.
    pub disconn_evt: EventFd,
    /// Tcp listening address.
    pub addr: String,
}

impl ClientIoHandler {
    pub fn new(
        stream: TcpStream,
        client: Arc<ClientState>,
        server: Arc<VncServer>,
        addr: String,
    ) -> Self {
        ClientIoHandler {
            stream,
            tls_conn: None,
            msg_handler: ClientIoHandler::handle_version,
            expect: 12,
            client,
            server,
            disconn_evt: EventFd::new(libc::EFD_NONBLOCK).unwrap(),
            addr,
        }
    }
}

impl ClientIoHandler {
    fn client_handle_read(&mut self) -> Result<(), anyhow::Error> {
        // Read message from tcpstream.
        if let Err(e) = self.read_msg() {
            return Err(e);
        }

        let client = self.client.clone();
        while client.in_buffer.lock().unwrap().len() >= self.expect {
            if let Err(e) = (self.msg_handler)(self) {
                return Err(e);
            }

            if self.client.conn_state.lock().unwrap().dis_conn {
                return Err(anyhow!(VncError::Disconnection));
            }

            if self.expect == 0 {
                break;
            }
        }

        Ok(())
    }

    fn client_handle_write(&mut self) {
        let client = self.client.clone();
        let mut locked_buffer = client.out_buffer.lock().unwrap();
        let len = locked_buffer.len();
        let buf = locked_buffer.read_front(len);
        self.write_msg(buf);
        locked_buffer.remov_front(len);
        drop(locked_buffer);
    }

    /// Read buf from stream, return the size.
    fn read_msg(&mut self) -> Result<usize> {
        let mut buf = Vec::new();
        let mut len: usize = 0;
        if self.tls_conn.is_none() {
            match self.read_plain_msg(&mut buf) {
                Ok(n) => len = n,
                Err(e) => return Err(e),
            }
        }
        if let Some(tc) = &self.tls_conn {
            if tc.is_handshaking() {
                return Ok(0_usize);
            }
        }
        if self.tls_conn.is_some() {
            match self.read_tls_msg(&mut buf) {
                Ok(n) => len = n,
                Err(e) => return Err(e),
            }
        }
        self.client
            .in_buffer
            .lock()
            .unwrap()
            .read(&mut buf[..len].to_vec());

        Ok(len)
    }

    // Read from vencrypt channel.
    fn read_tls_msg(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
        let mut len = 0_usize;
        if self.tls_conn.is_none() {
            return Ok(0_usize);
        }
        let tc: &mut ServerConnection;
        match &mut self.tls_conn {
            Some(sc) => tc = sc,
            None => return Ok(0_usize),
        }

        if let Err(e) = tc.read_tls(&mut self.stream) {
            error!("tls_conn read error {:?}", e);
            return Err(anyhow!(VncError::ReadMessageFailed(format!(
                "tls_conn read error {:?}",
                e
            ))));
        }

        if let Ok(io_state) = tc.process_new_packets() {
            if io_state.plaintext_bytes_to_read() > 0 {
                len = io_state.plaintext_bytes_to_read();
                buf.resize(len, 0u8);
                if let Err(e) = tc.reader().read_exact(buf) {
                    error!("tls_conn read error {:?}", e);
                    buf.clear();
                    return Err(anyhow!(VncError::ReadMessageFailed(format!(
                        "tls_conn read error {:?}",
                        e
                    ))));
                }
            }
        }
        Ok(len)
    }

    /// Read plain txt.
    fn read_plain_msg(&mut self, buf: &mut Vec<u8>) -> Result<usize> {
        let mut len = 0_usize;
        buf.resize(MAX_RECVBUF_LEN, 0u8);
        match self.stream.read(buf) {
            Ok(ret) => len = ret,
            Err(e) => error!("read msg error: {:?}", e),
        }

        Ok(len)
    }

    /// Write buf to stream
    /// Choose different channel according to whether or not to encrypt
    pub fn write_msg(&mut self, buf: &[u8]) {
        if self.tls_conn.is_none() {
            self.write_plain_msg(buf)
        } else {
            self.write_tls_msg(buf)
        }
    }

    // Send vencrypt message.
    fn write_tls_msg(&mut self, buf: &[u8]) {
        let buf_size = buf.len();
        let mut offset = 0;

        let tc: &mut ServerConnection;
        match &mut self.tls_conn {
            Some(ts) => tc = ts,
            None => {
                return;
            }
        }

        while offset < buf_size {
            let next = cmp::min(buf_size, offset + MAX_SEND_LEN);
            let tmp_buf = &buf[offset..next].to_vec();
            if let Err(e) = tc.writer().write_all(tmp_buf) {
                error!("write msg error: {:?}", e);
                return;
            }
            if let Err(_e) = vnc_write_tls_message(tc, &mut self.stream) {
                self.client.conn_state.lock().unwrap().dis_conn = true;
                return;
            }
            offset = next;
        }
    }

    /// Send plain txt.
    fn write_plain_msg(&mut self, buf: &[u8]) {
        let buf_size = buf.len();
        let mut offset = 0;
        loop {
            let tmp_buf = &buf[offset..];
            match self.stream.write(tmp_buf) {
                Ok(ret) => {
                    offset += ret;
                }
                Err(e) => {
                    if e.kind() == std::io::ErrorKind::WouldBlock {
                        self.stream.flush().unwrap();
                        continue;
                    } else {
                        error!("write msg error: {:?}", e);
                        self.client.conn_state.lock().unwrap().dis_conn = true;
                        return;
                    }
                }
            }
            self.stream.flush().unwrap();
            if offset >= buf_size {
                break;
            }
        }
    }

    /// Exchange RFB protocol version with client.
    fn handle_version(&mut self) -> Result<()> {
        let buf = self.read_incoming_msg();
        let res = String::from_utf8_lossy(&buf);
        let ver_str = &res[0..12].to_string();
        let ver;
        match scanf!(ver_str, "RFB {usize:/\\d\\{3\\}/}.{usize:/\\d\\{3\\}/}\n") {
            Ok(v) => {
                ver = v;
            }
            Err(e) => {
                let msg = format!("Unsupport RFB version: {}", e);
                error!("{}", msg);
                return Err(anyhow!(VncError::UnsupportRFBProtocolVersion));
            }
        }

        let mut version = VncVersion::new(ver.0 as u16, ver.1 as u16);
        if version.major != 3 || ![3, 4, 5, 7, 8].contains(&version.minor) {
            let mut buf = Vec::new();
            buf.append(&mut (AuthState::Invalid as u32).to_be_bytes().to_vec());
            self.write_msg(&buf);
            return Err(anyhow!(VncError::UnsupportRFBProtocolVersion));
        }

        if [4, 5].contains(&version.minor) {
            version.minor = 3;
        }
        self.client.conn_state.lock().unwrap().version = version;
        let auth = self.server.security_type.lock().unwrap().auth;

        if self.client.conn_state.lock().unwrap().version.minor == 3 {
            error!("Waiting for handle minor=3 ...");
            match auth {
                AuthState::No => {
                    let mut buf = Vec::new();
                    buf.append(&mut (AuthState::No as u32).to_be_bytes().to_vec());
                    self.write_msg(&buf);
                    self.update_event_handler(1, ClientIoHandler::handle_client_init);
                }
                _ => {
                    self.auth_failed("Unsupported auth method");
                    return Err(anyhow!(VncError::AuthFailed(String::from(
                        "Unsupported auth method"
                    ))));
                }
            }
        } else {
            let mut buf = [0u8; 2];
            buf[0] = 1; // Number of security types.
            buf[1] = auth as u8;
            self.write_msg(&buf);
            self.update_event_handler(1, ClientIoHandler::handle_auth);
        }
        Ok(())
    }

    /// Initialize the connection of vnc client.
    pub fn handle_client_init(&mut self) -> Result<()> {
        let mut buf = Vec::new();
        // Send server framebuffer info.
        let locked_surface = self.server.vnc_surface.lock().unwrap();
        let width = get_image_width(locked_surface.server_image);
        let height = get_image_height(locked_surface.server_image);
        drop(locked_surface);
        if !(0..=MAX_WINDOW_WIDTH as i32).contains(&width)
            || !(0..=MAX_WINDOW_HEIGHT as i32).contains(&height)
        {
            error!("Invalid Image Size!");
            return Err(anyhow!(VncError::InvalidImageSize));
        }
        buf.append(&mut (width as u16).to_be_bytes().to_vec());
        buf.append(&mut (height as u16).to_be_bytes().to_vec());
        let client = self.client.clone();
        pixel_format_message(&client, &mut buf);

        buf.append(&mut (APP_NAME.to_string().len() as u32).to_be_bytes().to_vec());
        buf.append(&mut APP_NAME.to_string().as_bytes().to_vec());
        self.write_msg(&buf);
        self.update_event_handler(1, ClientIoHandler::handle_protocol_msg);
        Ok(())
    }

    /// Authentication
    fn handle_auth(&mut self) -> Result<()> {
        let buf = self.read_incoming_msg();
        let auth = self.server.security_type.lock().unwrap().auth;
        let version = self.client.conn_state.lock().unwrap().version.clone();

        if buf[0] != auth as u8 {
            self.auth_failed("Authentication failed");
            error!("handle_auth");
            return Err(anyhow!(VncError::AuthFailed(String::from("handle_auth"))));
        }

        match auth {
            AuthState::No => {
                if version.minor >= 8 {
                    let buf = [0u8; 4];
                    self.write_msg(&buf);
                }
                self.update_event_handler(1, ClientIoHandler::handle_client_init);
            }
            AuthState::Vencrypt => {
                // Send VeNCrypt version 0.2.
                let mut buf = [0u8; 2];
                buf[0] = 0_u8;
                buf[1] = 2_u8;

                self.write_msg(&buf);
                self.update_event_handler(2, ClientIoHandler::client_vencrypt_init);
            }
            _ => {
                self.auth_failed("Unhandled auth method");
                error!("handle_auth");
                return Err(anyhow!(VncError::AuthFailed(String::from("handle_auth"))));
            }
        }
        Ok(())
    }

    /// Process the data sent by the client
    pub fn handle_protocol_msg(&mut self) -> Result<()> {
        // According to RFB protocol, first byte identifies the event type.
        let buf = self.read_incoming_msg();
        match ClientMsg::from(buf[0]) {
            ClientMsg::SetPixelFormat => {
                return self.set_pixel_format();
            }
            ClientMsg::SetEncodings => {
                return self.set_encodings();
            }
            ClientMsg::FramebufferUpdateRequest => {
                self.update_frame_buff();
            }
            ClientMsg::KeyEvent => {
                self.key_envent();
            }
            ClientMsg::PointerEvent => {
                self.point_event();
            }
            ClientMsg::ClientCutText => {
                self.client_cut_event();
            }
            _ => {
                self.update_event_handler(1, ClientIoHandler::handle_protocol_msg);
            }
        }
        Ok(())
    }

    /// Tell the client that the specified pixel values should be
    /// mapped to the given RGB intensities.
    fn send_color_map(&mut self) {
        let mut buf: Vec<u8> = Vec::new();
        buf.append(
            &mut (ServerMsg::SetColourMapEntries as u8)
                .to_be_bytes()
                .to_vec(),
        );
        buf.append(&mut (0_u8).to_be_bytes().to_vec());
        // First color.
        buf.append(&mut (0_u16).to_be_bytes().to_vec());
        // Number of colors.
        buf.append(&mut (NUM_OF_COLORMAP as u16).to_be_bytes().to_vec());

        let pf = self.client.client_dpm.lock().unwrap().pf.clone();
        for i in 0..NUM_OF_COLORMAP as u16 {
            let r = ((i >> pf.red.shift) & pf.red.max as u16) << (16 - pf.red.bits);
            let g = ((i >> pf.green.shift) & pf.green.max as u16) << (16 - pf.green.bits);
            let b = ((i >> pf.blue.shift) & pf.blue.max as u16) << (16 - pf.blue.bits);
            buf.append(&mut (r as u16).to_be_bytes().to_vec());
            buf.append(&mut (g as u16).to_be_bytes().to_vec());
            buf.append(&mut (b as u16).to_be_bytes().to_vec());
        }

        self.write_msg(&buf);
    }

    /// Set image format.
    fn set_pixel_format(&mut self) -> Result<()> {
        if self.expect == 1 {
            self.expect = 20;
            return Ok(());
        }

        let buf = self.read_incoming_msg();
        let mut bit_per_pixel: u8 = buf[4];
        let big_endian_flag = buf[6];
        let true_color_flag: u8 = buf[7];
        let mut red_max: u16 = u16::from_be_bytes([buf[8], buf[9]]);
        let mut green_max: u16 = u16::from_be_bytes([buf[10], buf[11]]);
        let mut blue_max: u16 = u16::from_be_bytes([buf[12], buf[13]]);
        let mut red_shift: u8 = buf[14];
        let mut green_shift: u8 = buf[15];
        let mut blue_shift: u8 = buf[16];
        if true_color_flag == 0 {
            bit_per_pixel = 8;
            red_max = 7;
            green_max = 7;
            blue_max = 3;
            red_shift = 0;
            green_shift = 3;
            blue_shift = 6;
        }

        // Verify the validity of pixel format.
        // bit_per_pixel: Bits occupied by each pixel.
        if ![8, 16, 32].contains(&bit_per_pixel) {
            self.client.conn_state.lock().unwrap().dis_conn = true;
            error!("Worng format of bits_per_pixel");
            return Err(anyhow!(VncError::ProtocolMessageFailed(String::from(
                "set pixel format"
            ))));
        }

        let mut locked_dpm = self.client.client_dpm.lock().unwrap();
        locked_dpm.pf.red.set_color_info(red_shift, red_max);
        locked_dpm.pf.green.set_color_info(green_shift, green_max);
        locked_dpm.pf.blue.set_color_info(blue_shift, blue_max);
        locked_dpm.pf.pixel_bits = bit_per_pixel;
        locked_dpm.pf.pixel_bytes = bit_per_pixel / BIT_PER_BYTE as u8;
        // Standard pixel format, depth is equal to 24.
        locked_dpm.pf.depth = if bit_per_pixel == 32 {
            24
        } else {
            bit_per_pixel
        };
        locked_dpm.client_be = big_endian_flag != 0;

        if !locked_dpm.pf.is_default_pixel_format() {
            locked_dpm.convert = true;
        }
        drop(locked_dpm);
        if true_color_flag == 0 {
            self.send_color_map();
        }

        VNC_RECT_INFO.lock().unwrap().clear();
        self.update_event_handler(1, ClientIoHandler::handle_protocol_msg);
        Ok(())
    }

    /// Set encoding.
    fn set_encodings(&mut self) -> Result<()> {
        let client = self.client.clone();
        let server = self.server.clone();
        let buf = self.read_incoming_msg();
        if self.expect == 1 {
            self.expect = 4;
            return Ok(());
        }

        let mut num_encoding: u16;
        if self.expect == 4 {
            num_encoding = u16::from_be_bytes([buf[2], buf[3]]);
            if num_encoding > 0 {
                self.expect = 4 + (num_encoding as usize) * 4;
                return Ok(());
            }
        } else {
            num_encoding = u16::from_be_bytes([buf[2], buf[3]]);
        }

        let mut locked_dpm = self.client.client_dpm.lock().unwrap();
        let mut locked_state = self.client.conn_state.lock().unwrap();
        while num_encoding > 0 {
            let offset = (4 * num_encoding) as usize;
            let enc = i32::from_be_bytes([
                buf[offset],
                buf[offset + 1],
                buf[offset + 2],
                buf[offset + 3],
            ]);
            match enc {
                ENCODING_RAW => {
                    locked_dpm.enc = enc;
                }
                ENCODING_HEXTILE => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureHextile as usize;
                    locked_dpm.enc = enc;
                }
                ENCODING_TIGHT => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureTight as usize;
                    locked_dpm.enc = enc;
                }
                ENCODING_ZLIB => {
                    // ZRLE compress better than ZLIB, so prioritize ZRLE.
                    if locked_state.feature & (1 << VncFeatures::VncFeatureZrle as usize) == 0 {
                        locked_state.feature |= 1 << VncFeatures::VncFeatureZlib as usize;
                        locked_dpm.enc = enc;
                    }
                }
                ENCODING_ZRLE => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureZrle as usize;
                    locked_dpm.enc = enc;
                }
                ENCODING_ZYWRLE => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureZywrle as usize;
                    locked_dpm.enc = enc;
                }
                ENCODING_DESKTOPRESIZE => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureResize as usize;
                }
                ENCODING_DESKTOP_RESIZE_EXT => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureResizeExt as usize;
                }
                ENCODING_POINTER_TYPE_CHANGE => {
                    locked_state.feature |= 1 << VncFeatures::VncFeaturePointerTypeChange as usize;
                }
                ENCODING_RICH_CURSOR => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureRichCursor as usize;
                }
                ENCODING_ALPHA_CURSOR => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureAlphaCursor as usize;
                }
                ENCODING_WMVI => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureWmvi as usize;
                }
                ENCODING_LED_STATE => {
                    locked_state.feature |= 1 << VncFeatures::VncFeatureLedState as usize;
                }
                _ => {}
            }

            num_encoding -= 1;
        }

        drop(locked_dpm);
        drop(locked_state);
        let mut buf: Vec<u8> = Vec::new();
        // VNC desktop resize.
        desktop_resize(&client, &server, &mut buf);
        // VNC display cursor define.
        display_cursor_define(&client, &server, &mut buf);
        client.out_buffer.lock().unwrap().read(&mut buf);
        vnc_flush_notify(&client);
        self.update_event_handler(1, ClientIoHandler::handle_protocol_msg);
        Ok(())
    }

    /// Update image for client.
    fn update_frame_buff(&mut self) {
        if self.expect == 1 {
            self.expect = 10;
            return;
        }
        let buf = self.read_incoming_msg();
        let locked_surface = self.server.vnc_surface.lock().unwrap();
        let width = get_image_width(locked_surface.server_image);
        let height = get_image_height(locked_surface.server_image);
        drop(locked_surface);
        let client = self.client.clone();
        let mut locked_state = client.conn_state.lock().unwrap();
        if buf[1] != 0 {
            if locked_state.update_state != UpdateState::Force {
                locked_state.update_state = UpdateState::Incremental;
            }
        } else {
            locked_state.update_state = UpdateState::Force;
            let x = u16::from_be_bytes([buf[2], buf[3]]) as i32;
            let y = u16::from_be_bytes([buf[4], buf[5]]) as i32;
            let w = u16::from_be_bytes([buf[6], buf[7]]) as i32;
            let h = u16::from_be_bytes([buf[8], buf[9]]) as i32;
            set_area_dirty(
                &mut client.dirty_bitmap.lock().unwrap(),
                x,
                y,
                w,
                h,
                width,
                height,
            );
        }
        drop(locked_state);
        let server = self.server.clone();
        vnc_refresh_notify(&server);
        self.update_event_handler(1, ClientIoHandler::handle_protocol_msg);
    }

    /// Invalid authentication, send 1 to reject.
    fn auth_failed(&mut self, msg: &str) {
        let auth_rej: u8 = 1;
        let mut buf: Vec<u8> = vec![1u8];
        buf.append(&mut (auth_rej as u32).to_be_bytes().to_vec());
        // If the RFB protocol version is above 3.8, an error reason will be returned.
        if self.client.conn_state.lock().unwrap().version.minor >= 8 {
            let err_msg = msg;
            buf.append(&mut (err_msg.len() as u32).to_be_bytes().to_vec());
            buf.append(&mut err_msg.as_bytes().to_vec());
        }
        self.write_msg(&buf);
    }

    /// Read the data from the receiver buffer.
    pub fn read_incoming_msg(&mut self) -> Vec<u8> {
        let mut locked_in_buffer = self.client.in_buffer.lock().unwrap();
        let buf = locked_in_buffer.read_front(self.expect).to_vec();
        buf
    }

    /// Action token after the event.
    ///
    /// # Arguments
    ///
    /// * `expect` - the size of bytes of next callback function.
    /// * `msg_handler` - callback function of the next event.
    pub fn update_event_handler(
        &mut self,
        expect: usize,
        msg_handler: fn(&mut ClientIoHandler) -> Result<()>,
    ) {
        self.client
            .in_buffer
            .lock()
            .unwrap()
            .remov_front(self.expect);
        self.expect = expect;
        self.msg_handler = msg_handler;
    }

    fn disconn_evt_handler(&mut self) -> Vec<EventNotifier> {
        let notifiers = vec![
            EventNotifier::new(
                NotifierOperation::Delete,
                self.stream.as_raw_fd(),
                None,
                EventSet::IN,
                Vec::new(),
            ),
            EventNotifier::new(
                NotifierOperation::Delete,
                self.client.write_fd.lock().unwrap().as_raw_fd(),
                None,
                EventSet::IN,
                Vec::new(),
            ),
            EventNotifier::new(
                NotifierOperation::Delete,
                self.disconn_evt.as_raw_fd(),
                None,
                EventSet::IN,
                Vec::new(),
            ),
        ];

        notifiers
    }
}

/// Internal notifiers for Client message.
impl EventNotifierHelper for ClientIoHandler {
    fn internal_notifiers(client_io_handler: Arc<Mutex<ClientIoHandler>>) -> Vec<EventNotifier> {
        let mut notifiers: Vec<EventNotifier> = Vec::new();

        // Register event notifier for read.
        let client_io = client_io_handler.clone();
        let handler: Box<dyn Fn(EventSet, RawFd) -> Option<Vec<EventNotifier>>> =
            Box::new(move |event, _fd: RawFd| {
                let mut locked_client_io = client_io.lock().unwrap();
                if event & EventSet::ERROR == EventSet::ERROR
                    || event & EventSet::HANG_UP == EventSet::HANG_UP
                    || event & EventSet::READ_HANG_UP == EventSet::READ_HANG_UP
                {
                    locked_client_io.client.conn_state.lock().unwrap().dis_conn = true;
                } else if event & EventSet::IN == EventSet::IN {
                    if let Err(_e) = locked_client_io.client_handle_read() {
                        locked_client_io.client.conn_state.lock().unwrap().dis_conn = true;
                    }
                }

                // Do disconnection event.
                if locked_client_io
                    .client
                    .conn_state
                    .lock()
                    .unwrap()
                    .is_disconnect()
                {
                    locked_client_io.disconn_evt.write(1).unwrap();
                }
                drop(locked_client_io);

                None
            });
        let client_io = client_io_handler.clone();
        notifiers.push(EventNotifier::new(
            NotifierOperation::AddShared,
            client_io.lock().unwrap().stream.as_raw_fd(),
            None,
            EventSet::IN | EventSet::READ_HANG_UP,
            vec![Arc::new(Mutex::new(handler))],
        ));

        // Register event notifier for write.
        let client_io = client_io_handler.clone();
        let client = client_io.lock().unwrap().client.clone();
        let handler: Box<dyn Fn(EventSet, RawFd) -> Option<Vec<EventNotifier>>> =
            Box::new(move |_event, fd| {
                read_fd(fd);
                let mut locked_client_io = client_io.lock().unwrap();
                locked_client_io.client_handle_write();

                // do disconnection event.
                if locked_client_io.client.conn_state.lock().unwrap().dis_conn {
                    locked_client_io.disconn_evt.write(1).unwrap();
                }
                drop(locked_client_io);

                None
            });
        notifiers.push(EventNotifier::new(
            NotifierOperation::AddShared,
            client.write_fd.lock().unwrap().as_raw_fd(),
            None,
            EventSet::IN,
            vec![Arc::new(Mutex::new(handler))],
        ));

        // Register event for disconnect.
        let client_io = client_io_handler.clone();
        let handler: Box<dyn Fn(EventSet, RawFd) -> Option<Vec<EventNotifier>>> =
            Box::new(move |_event, fd| {
                read_fd(fd);
                // Drop client info from vnc server.
                let mut locked_client_io = client_io.lock().unwrap();
                let addr = locked_client_io.addr.clone();
                let server = locked_client_io.server.clone();
                let notifiers = locked_client_io.disconn_evt_handler();
                // Shutdown stream.
                if let Err(e) = locked_client_io.stream.shutdown(Shutdown::Both) {
                    error!("Shutdown stream failed: {}", e);
                }
                drop(locked_client_io);
                server.client_handlers.lock().unwrap().remove(&addr);

                Some(notifiers)
            });

        notifiers.push(EventNotifier::new(
            NotifierOperation::AddShared,
            client_io_handler.lock().unwrap().disconn_evt.as_raw_fd(),
            None,
            EventSet::IN,
            vec![Arc::new(Mutex::new(handler))],
        ));

        notifiers
    }
}

/// Generate the data that needs to be sent.
/// Add to send queue
pub fn get_rects(client: &Arc<ClientState>, server: &Arc<VncServer>, dirty_num: i32) -> i32 {
    if !is_need_update(client) || dirty_num == 0 {
        return 0;
    }

    let mut num_rects = 0;
    let mut x: u64;
    let mut y: u64 = 0;
    let mut h: u64;
    let mut x2: u64;
    let mut rects = Vec::new();
    let bpl = client.dirty_bitmap.lock().unwrap().vol() / MAX_WINDOW_HEIGHT as usize;
    let locked_surface = server.vnc_surface.lock().unwrap();
    let mut locked_dirty = client.dirty_bitmap.lock().unwrap();

    let height = get_image_height(locked_surface.server_image) as u64;
    let width = get_image_width(locked_surface.server_image) as u64;
    loop {
        // Find the first non-zero bit in dirty bitmap.
        let offset = locked_dirty.find_next_bit(y as usize * bpl).unwrap() as u64;
        if offset >= height as u64 * bpl as u64 {
            break;
        }

        x = offset % bpl as u64;
        y = offset / bpl as u64;
        // Find value in one line to the end.
        x2 = locked_dirty.find_next_zero(offset as usize).unwrap() as u64 % bpl as u64;
        let mut i = y;
        while i < height {
            if !locked_dirty.contain((i * bpl as u64 + x) as usize).unwrap() {
                break;
            }
            for j in x..x2 {
                locked_dirty.clear((i * bpl as u64 + j) as usize).unwrap();
            }
            i += 1;
        }

        h = i - y;
        x2 = cmp::min(x2, width / DIRTY_PIXELS_NUM as u64);
        if x2 > x as u64 {
            rects.push(Rectangle::new(
                (x * DIRTY_PIXELS_NUM as u64) as i32,
                y as i32,
                ((x2 - x) * DIRTY_PIXELS_NUM as u64) as i32,
                h as i32,
            ));
            num_rects += 1;
        }

        if x == 0 && x2 == width / DIRTY_PIXELS_NUM as u64 {
            y += h;
            if y == height {
                break;
            }
        }
    }

    VNC_RECT_INFO
        .lock()
        .unwrap()
        .push(RectInfo::new(client, rects));

    drop(locked_dirty);
    drop(locked_surface);

    let mut locked_state = client.conn_state.lock().unwrap();
    locked_state.update_state = UpdateState::No;
    drop(locked_state);

    num_rects
}

fn vnc_write_tls_message(tc: &mut ServerConnection, stream: &mut TcpStream) -> Result<()> {
    while tc.wants_write() {
        match tc.write_tls(stream) {
            Ok(_) => {}
            Err(e) => {
                if e.kind() == std::io::ErrorKind::WouldBlock {
                    stream.flush().unwrap();
                    continue;
                } else {
                    error!("write msg error: {:?}", e);
                    return Err(anyhow!(VncError::Disconnection));
                }
            }
        }
    }

    Ok(())
}

/// Whether the client's image data needs to be updated.
pub fn is_need_update(client: &Arc<ClientState>) -> bool {
    match client.conn_state.lock().unwrap().update_state {
        UpdateState::No => false,
        UpdateState::Incremental => {
            // throttle_output_offset
            true
        }
        UpdateState::Force => {
            // force_update_offset
            true
        }
    }
}

/// Set pixformat for client.
pub fn pixel_format_message(client: &Arc<ClientState>, buf: &mut Vec<u8>) {
    let mut locked_dpm = client.client_dpm.lock().unwrap();
    locked_dpm.pf.init_pixelformat();
    let big_endian: u8 = if cfg!(target_endian = "big") {
        1_u8
    } else {
        0_u8
    };
    buf.append(&mut (locked_dpm.pf.pixel_bits as u8).to_be_bytes().to_vec()); // Bit per pixel.
    buf.append(&mut (locked_dpm.pf.depth as u8).to_be_bytes().to_vec()); // Depth.
    buf.append(&mut (big_endian as u8).to_be_bytes().to_vec()); // Big-endian flag.
    buf.append(&mut (1_u8).to_be_bytes().to_vec()); // True-color flag.
    buf.append(&mut (locked_dpm.pf.red.max as u16).to_be_bytes().to_vec()); // Red max.
    buf.append(&mut (locked_dpm.pf.green.max as u16).to_be_bytes().to_vec()); // Green max.
    buf.append(&mut (locked_dpm.pf.blue.max as u16).to_be_bytes().to_vec()); // Blue max.
    buf.append(&mut (locked_dpm.pf.red.shift as u8).to_be_bytes().to_vec()); // Red shift.
    buf.append(&mut (locked_dpm.pf.green.shift as u8).to_be_bytes().to_vec()); // Green shift.
    buf.append(&mut (locked_dpm.pf.blue.shift as u8).to_be_bytes().to_vec()); // Blue shift.
    buf.append(&mut [0; 3].to_vec()); // Padding.
    drop(locked_dpm);
}

/// Set Desktop Size.
pub fn desktop_resize(client: &Arc<ClientState>, server: &Arc<VncServer>, buf: &mut Vec<u8>) {
    let locked_state = client.conn_state.lock().unwrap();
    let locked_surface = server.vnc_surface.lock().unwrap();
    if !locked_state.has_feature(VncFeatures::VncFeatureResizeExt)
        && !locked_state.has_feature(VncFeatures::VncFeatureResize)
    {
        return;
    }
    let width = get_image_width(locked_surface.server_image);
    let height = get_image_height(locked_surface.server_image);
    drop(locked_state);
    drop(locked_surface);

    if !(0..=MAX_WINDOW_WIDTH as i32).contains(&width)
        || !(0..=MAX_WINDOW_HEIGHT as i32).contains(&height)
    {
        error!("Invalid Image Size!");
        return;
    }

    buf.append(&mut (ServerMsg::FramebufferUpdate as u8).to_be_bytes().to_vec());
    buf.append(&mut (0_u8).to_be_bytes().to_vec());
    buf.append(&mut (1_u16).to_be_bytes().to_vec());
    framebuffer_upadate(0, 0, width, height, ENCODING_DESKTOPRESIZE, buf);
}

/// Set color depth for client.
pub fn set_color_depth(client: &Arc<ClientState>, server: &Arc<VncServer>, buf: &mut Vec<u8>) {
    let locked_state = client.conn_state.lock().unwrap();
    if locked_state.has_feature(VncFeatures::VncFeatureWmvi) {
        let locked_surface = server.vnc_surface.lock().unwrap();
        let width = get_image_width(locked_surface.server_image);
        let height = get_image_height(locked_surface.server_image);
        drop(locked_surface);

        if !(0..=MAX_WINDOW_WIDTH as i32).contains(&width)
            || !(0..=MAX_WINDOW_HEIGHT as i32).contains(&height)
        {
            error!("Invalid Image Size!");
            return;
        }
        buf.append(&mut (ServerMsg::FramebufferUpdate as u8).to_be_bytes().to_vec());
        buf.append(&mut (0_u8).to_be_bytes().to_vec()); // Padding.
        buf.append(&mut (1_u16).to_be_bytes().to_vec()); // Number of pixel block.
        framebuffer_upadate(0, 0, width, height, ENCODING_WMVI, buf);
        buf.append(&mut (ENCODING_RAW as u32).to_be_bytes().to_vec());
        pixel_format_message(client, buf);
    } else if !client
        .client_dpm
        .lock()
        .unwrap()
        .pf
        .is_default_pixel_format()
    {
        client.client_dpm.lock().unwrap().convert = true;
    }
}

/// Send framebuf of mouse to the client.
pub fn display_cursor_define(
    client: &Arc<ClientState>,
    server: &Arc<VncServer>,
    buf: &mut Vec<u8>,
) {
    let mut cursor: DisplayMouse;
    let mut mask: Vec<u8>;
    let locked_cursor = server.vnc_cursor.lock().unwrap();
    match &locked_cursor.cursor {
        Some(c) => {
            cursor = c.clone();
        }
        None => {
            return;
        }
    }
    match &locked_cursor.mask {
        Some(m) => {
            mask = m.clone();
        }
        None => {
            return;
        }
    }
    drop(locked_cursor);
    if cursor.data.is_empty()
        || cursor.data.len() != ((cursor.width * cursor.height) as usize) * bytes_per_pixel()
    {
        return;
    }
    if client
        .conn_state
        .lock()
        .unwrap()
        .has_feature(VncFeatures::VncFeatureAlphaCursor)
    {
        buf.append(&mut (ServerMsg::FramebufferUpdate as u8).to_be_bytes().to_vec());
        buf.append(&mut (0_u8).to_be_bytes().to_vec()); // padding
        buf.append(&mut (1_u16).to_be_bytes().to_vec()); // number of rects

        framebuffer_upadate(
            cursor.hot_x as i32,
            cursor.hot_y as i32,
            cursor.width as i32,
            cursor.height as i32,
            ENCODING_ALPHA_CURSOR as i32,
            buf,
        );
        buf.append(&mut (ENCODING_RAW as u32).to_be_bytes().to_vec());
        buf.append(&mut cursor.data);
        return;
    }

    if client
        .conn_state
        .lock()
        .unwrap()
        .has_feature(VncFeatures::VncFeatureRichCursor)
    {
        buf.append(&mut (ServerMsg::FramebufferUpdate as u8).to_be_bytes().to_vec());
        buf.append(&mut (0_u8).to_be_bytes().to_vec()); // padding
        buf.append(&mut (1_u16).to_be_bytes().to_vec()); // number of rects

        framebuffer_upadate(
            cursor.hot_x as i32,
            cursor.hot_y as i32,
            cursor.width as i32,
            cursor.height as i32,
            ENCODING_RICH_CURSOR as i32,
            buf,
        );
        let dpm = client.client_dpm.lock().unwrap().clone();
        let data_size = cursor.width * cursor.height * dpm.pf.pixel_bytes as u32;
        let data_ptr = cursor.data.as_ptr() as *mut u8;
        write_pixel(data_ptr, data_size as usize, &dpm, buf);
        buf.append(&mut mask);
    }
}

/// Consume the output buffer.
pub fn vnc_flush_notify(client: &Arc<ClientState>) {
    client.write_fd.lock().unwrap().write(1).unwrap();
}