pxl8/pxl8d/src/transport.rs

extern crate alloc;

use alloc::vec;
use alloc::vec::Vec;
use crate::protocol::*;
use crate::protocol::pxl8_msg_type::*;
use crate::voxel::VoxelChunk;

pub const DEFAULT_PORT: u16 = 7777;
pub const CHUNK_MAX_PAYLOAD: usize = 1400;
pub const CHUNK_FLAG_RLE: u8 = 0x01;
pub const CHUNK_FLAG_FINAL: u8 = 0x04;

pub const CHUNK_TYPE_VXL: u8 = 0;
pub const CHUNK_TYPE_BSP: u8 = 1;

pub struct ChunkMessage {
    pub chunk_type: u8,
    pub id: u32,
    pub cx: i32,
    pub cy: i32,
    pub cz: i32,
    pub version: u32,
    pub flags: u8,
    pub fragment_idx: u8,
    pub fragment_count: u8,
    pub payload: Vec<u8>,
}

impl ChunkMessage {
    pub fn from_voxel(chunk: &VoxelChunk, version: u32) -> Vec<ChunkMessage> {
        let rle_data = chunk.rle_encode();
        let total_size = rle_data.len();

        if total_size <= CHUNK_MAX_PAYLOAD {
            return vec![ChunkMessage {
                chunk_type: CHUNK_TYPE_VXL,
                id: 0,
                cx: chunk.cx,
                cy: chunk.cy,
                cz: chunk.cz,
                version,
                flags: CHUNK_FLAG_RLE | CHUNK_FLAG_FINAL,
                fragment_idx: 0,
                fragment_count: 1,
                payload: rle_data,
            }];
        }

        let fragment_count = (total_size + CHUNK_MAX_PAYLOAD - 1) / CHUNK_MAX_PAYLOAD;
        let mut messages = Vec::new();

        for i in 0..fragment_count {
            let start = i * CHUNK_MAX_PAYLOAD;
            let end = ((i + 1) * CHUNK_MAX_PAYLOAD).min(total_size);
            let is_final = i == fragment_count - 1;

            messages.push(ChunkMessage {
                chunk_type: CHUNK_TYPE_VXL,
                id: 0,
                cx: chunk.cx,
                cy: chunk.cy,
                cz: chunk.cz,
                version,
                flags: CHUNK_FLAG_RLE | if is_final { CHUNK_FLAG_FINAL } else { 0 },
                fragment_idx: i as u8,
                fragment_count: fragment_count as u8,
                payload: rle_data[start..end].to_vec(),
            });
        }

        messages
    }

    pub fn from_bsp(data: Vec<u8>, chunk_id: u32, version: u32) -> Vec<ChunkMessage> {
        let total_size = data.len();

        if total_size <= CHUNK_MAX_PAYLOAD {
            return vec![ChunkMessage {
                chunk_type: CHUNK_TYPE_BSP,
                id: chunk_id,
                cx: 0,
                cy: 0,
                cz: 0,
                version,
                flags: CHUNK_FLAG_FINAL,
                fragment_idx: 0,
                fragment_count: 1,
                payload: data,
            }];
        }

        let fragment_count = (total_size + CHUNK_MAX_PAYLOAD - 1) / CHUNK_MAX_PAYLOAD;
        let mut messages = Vec::new();

        for i in 0..fragment_count {
            let start = i * CHUNK_MAX_PAYLOAD;
            let end = ((i + 1) * CHUNK_MAX_PAYLOAD).min(total_size);
            let is_final = i == fragment_count - 1;

            messages.push(ChunkMessage {
                chunk_type: CHUNK_TYPE_BSP,
                id: chunk_id,
                cx: 0,
                cy: 0,
                cz: 0,
                version,
                flags: if is_final { CHUNK_FLAG_FINAL } else { 0 },
                fragment_idx: i as u8,
                fragment_count: fragment_count as u8,
                payload: data[start..end].to_vec(),
            });
        }

        messages
    }
}

#[cfg(unix)]
mod sys {
    use libc::{c_int, c_void, sockaddr, sockaddr_in, socklen_t};

    pub type RawSocket = c_int;
    pub const INVALID_SOCKET: RawSocket = -1;

    pub fn socket() -> RawSocket {
        unsafe { libc::socket(libc::AF_INET, libc::SOCK_DGRAM, 0) }
    }

    pub fn bind(sock: RawSocket, port: u16) -> c_int {
        let addr = sockaddr_in {
            sin_family: libc::AF_INET as u16,
            sin_port: port.to_be(),
            sin_addr: libc::in_addr { s_addr: u32::from_be_bytes([127, 0, 0, 1]).to_be() },
            sin_zero: [0; 8],
        };
        unsafe { libc::bind(sock, &addr as *const _ as *const sockaddr, core::mem::size_of::<sockaddr_in>() as socklen_t) }
    }

    pub fn set_nonblocking(sock: RawSocket) -> c_int {
        unsafe {
            let flags = libc::fcntl(sock, libc::F_GETFL, 0);
            libc::fcntl(sock, libc::F_SETFL, flags | libc::O_NONBLOCK)
        }
    }

    pub fn recvfrom(sock: RawSocket, buf: &mut [u8], addr: &mut sockaddr_in) -> isize {
        let mut addr_len = core::mem::size_of::<sockaddr_in>() as socklen_t;
        unsafe {
            libc::recvfrom(
                sock,
                buf.as_mut_ptr() as *mut c_void,
                buf.len(),
                0,
                addr as *mut _ as *mut sockaddr,
                &mut addr_len,
            )
        }
    }

    pub fn sendto(sock: RawSocket, buf: &[u8], addr: &sockaddr_in) -> isize {
        unsafe {
            libc::sendto(
                sock,
                buf.as_ptr() as *const c_void,
                buf.len(),
                0,
                addr as *const _ as *const sockaddr,
                core::mem::size_of::<sockaddr_in>() as socklen_t,
            )
        }
    }

    pub fn close(sock: RawSocket) {
        unsafe { libc::close(sock) };
    }
}

#[cfg(windows)]
mod sys {
    use windows_sys::Win32::Networking::WinSock::*;

    pub type RawSocket = SOCKET;
    pub const INVALID_SOCKET_VAL: RawSocket = INVALID_SOCKET;

    pub fn socket() -> RawSocket {
        unsafe { socket(AF_INET as i32, SOCK_DGRAM as i32, 0) }
    }

    pub fn bind(sock: RawSocket, port: u16) -> i32 {
        let addr = SOCKADDR_IN {
            sin_family: AF_INET,
            sin_port: port.to_be(),
            sin_addr: IN_ADDR { S_un: IN_ADDR_0 { S_addr: u32::from_be_bytes([127, 0, 0, 1]).to_be() } },
            sin_zero: [0; 8],
        };
        unsafe { bind(sock, &addr as *const _ as *const SOCKADDR, core::mem::size_of::<SOCKADDR_IN>() as i32) }
    }

    pub fn set_nonblocking(sock: RawSocket) -> i32 {
        let mut nonblocking: u32 = 1;
        unsafe { ioctlsocket(sock, FIONBIO as i32, &mut nonblocking) }
    }

    pub fn recvfrom(sock: RawSocket, buf: &mut [u8], addr: &mut SOCKADDR_IN) -> i32 {
        let mut addr_len = core::mem::size_of::<SOCKADDR_IN>() as i32;
        unsafe {
            recvfrom(
                sock,
                buf.as_mut_ptr(),
                buf.len() as i32,
                0,
                addr as *mut _ as *mut SOCKADDR,
                &mut addr_len,
            )
        }
    }

    pub fn sendto(sock: RawSocket, buf: &[u8], addr: &SOCKADDR_IN) -> i32 {
        unsafe {
            sendto(
                sock,
                buf.as_ptr(),
                buf.len() as i32,
                0,
                addr as *const _ as *const SOCKADDR,
                core::mem::size_of::<SOCKADDR_IN>() as i32,
            )
        }
    }

    pub fn close(sock: RawSocket) {
        unsafe { closesocket(sock) };
    }
}

#[cfg(unix)]
type SockAddr = libc::sockaddr_in;

#[cfg(windows)]
type SockAddr = windows_sys::Win32::Networking::WinSock::SOCKADDR_IN;

pub struct Transport {
    client_addr: SockAddr,
    has_client: bool,
    recv_buf: [u8; 512],
    send_buf: [u8; 2048],
    socket: sys::RawSocket,
}

impl Transport {
    pub fn bind(port: u16) -> Option<Self> {
        let sock = sys::socket();
        if sock == sys::INVALID_SOCKET {
            return None;
        }

        if sys::bind(sock, port) < 0 {
            sys::close(sock);
            return None;
        }

        if sys::set_nonblocking(sock) < 0 {
            sys::close(sock);
            return None;
        }

        Some(Self {
            client_addr: unsafe { core::mem::zeroed() },
            has_client: false,
            recv_buf: [0u8; 512],
            send_buf: [0u8; 2048],
            socket: sock,
        })
    }

    pub fn recv(&mut self) -> Option<u8> {
        let mut addr: SockAddr = unsafe { core::mem::zeroed() };
        let len = sys::recvfrom(self.socket, &mut self.recv_buf, &mut addr);

        if len <= 0 || (len as usize) < size_of::<pxl8_msg_header>() {
            return None;
        }

        self.client_addr = addr;
        self.has_client = true;

        let header = self.deserialize_header();
        Some(header.type_)
    }

    pub fn get_input(&self) -> pxl8_input_msg {
        self.deserialize_input()
    }

    pub fn get_command(&self) -> pxl8_command_msg {
        self.deserialize_command()
    }

    pub fn send_snapshot(
        &mut self,
        header: &pxl8_snapshot_header,
        entities: &[pxl8_entity_state],
        sequence: u32,
    ) {
        if !self.has_client {
            return;
        }

        let mut offset = 0;

        let msg_header = pxl8_msg_header {
            sequence,
            size: 0,
            type_: PXL8_MSG_SNAPSHOT as u8,
            version: PXL8_PROTOCOL_VERSION as u8,
        };
        offset += self.serialize_header(&msg_header, offset);
        offset += self.serialize_snapshot_header(header, offset);

        for entity in entities {
            offset += self.serialize_entity_state(entity, offset);
        }

        sys::sendto(self.socket, &self.send_buf[..offset], &self.client_addr);
    }

    pub fn send_chunk(&mut self, msg: &ChunkMessage, sequence: u32) {
        if !self.has_client {
            return;
        }

        let mut offset = 0;

        let msg_header = pxl8_msg_header {
            sequence,
            size: 0,
            type_: PXL8_MSG_CHUNK as u8,
            version: PXL8_PROTOCOL_VERSION as u8,
        };
        offset += self.serialize_header(&msg_header, offset);
        offset += self.serialize_chunk_msg_header(msg, offset);

        let payload_len = msg.payload.len().min(self.send_buf.len() - offset);
        self.send_buf[offset..offset + payload_len].copy_from_slice(&msg.payload[..payload_len]);
        offset += payload_len;

        sys::sendto(self.socket, &self.send_buf[..offset], &self.client_addr);
    }

    pub fn send_chunk_enter(&mut self, chunk_id: u32, chunk_type: u8, sequence: u32) {
        if !self.has_client {
            return;
        }

        let mut offset = 0;

        let msg_header = pxl8_msg_header {
            sequence,
            size: 0,
            type_: PXL8_MSG_CHUNK_ENTER as u8,
            version: PXL8_PROTOCOL_VERSION as u8,
        };
        offset += self.serialize_header(&msg_header, offset);

        let buf = &mut self.send_buf[offset..];
        buf[0..4].copy_from_slice(&chunk_id.to_be_bytes());
        buf[4] = chunk_type;
        buf[5] = 0;
        buf[6] = 0;
        buf[7] = 0;
        offset += 8;

        sys::sendto(self.socket, &self.send_buf[..offset], &self.client_addr);
    }

    fn serialize_chunk_msg_header(&mut self, msg: &ChunkMessage, offset: usize) -> usize {
        let buf = &mut self.send_buf[offset..];
        buf[0] = msg.chunk_type;
        buf[1] = msg.flags;
        buf[2] = msg.fragment_idx;
        buf[3] = msg.fragment_count;
        buf[4..8].copy_from_slice(&msg.id.to_be_bytes());
        buf[8..12].copy_from_slice(&(msg.cx as u32).to_be_bytes());
        buf[12..16].copy_from_slice(&(msg.cy as u32).to_be_bytes());
        buf[16..20].copy_from_slice(&(msg.cz as u32).to_be_bytes());
        buf[20..24].copy_from_slice(&msg.version.to_be_bytes());
        let payload_size = msg.payload.len() as u16;
        buf[24..26].copy_from_slice(&payload_size.to_be_bytes());
        buf[26..28].copy_from_slice(&[0u8; 2]);
        28
    }

    fn serialize_header(&mut self, h: &pxl8_msg_header, offset: usize) -> usize {
        let buf = &mut self.send_buf[offset..];
        buf[0..4].copy_from_slice(&h.sequence.to_be_bytes());
        buf[4..6].copy_from_slice(&h.size.to_be_bytes());
        buf[6] = h.type_;
        buf[7] = h.version;
        8
    }

    fn serialize_snapshot_header(&mut self, h: &pxl8_snapshot_header, offset: usize) -> usize {
        let buf = &mut self.send_buf[offset..];
        buf[0..2].copy_from_slice(&h.entity_count.to_be_bytes());
        buf[2..4].copy_from_slice(&h.event_count.to_be_bytes());
        buf[4..12].copy_from_slice(&h.player_id.to_be_bytes());
        buf[12..20].copy_from_slice(&h.tick.to_be_bytes());
        buf[20..24].copy_from_slice(&h.time.to_be_bytes());
        24
    }

    fn serialize_entity_state(&mut self, e: &pxl8_entity_state, offset: usize) -> usize {
        let buf = &mut self.send_buf[offset..];
        buf[0..8].copy_from_slice(&e.entity_id.to_be_bytes());
        buf[8..64].copy_from_slice(&e.userdata);
        64
    }

    fn deserialize_header(&self) -> pxl8_msg_header {
        let buf = &self.recv_buf;
        pxl8_msg_header {
            sequence: u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]),
            size: u16::from_be_bytes([buf[4], buf[5]]),
            type_: buf[6],
            version: buf[7],
        }
    }

    fn deserialize_input(&self) -> pxl8_input_msg {
        let buf = &self.recv_buf[8..];
        pxl8_input_msg {
            buttons: u32::from_be_bytes([buf[0], buf[1], buf[2], buf[3]]),
            look_dx: f32::from_be_bytes([buf[4], buf[5], buf[6], buf[7]]),
            look_dy: f32::from_be_bytes([buf[8], buf[9], buf[10], buf[11]]),
            move_x: f32::from_be_bytes([buf[12], buf[13], buf[14], buf[15]]),
            move_y: f32::from_be_bytes([buf[16], buf[17], buf[18], buf[19]]),
            yaw: f32::from_be_bytes([buf[20], buf[21], buf[22], buf[23]]),
            tick: u64::from_be_bytes([buf[24], buf[25], buf[26], buf[27], buf[28], buf[29], buf[30], buf[31]]),
            timestamp: u64::from_be_bytes([buf[32], buf[33], buf[34], buf[35], buf[36], buf[37], buf[38], buf[39]]),
        }
    }

    fn deserialize_command(&self) -> pxl8_command_msg {
        let buf = &self.recv_buf[8..];
        let mut cmd = pxl8_command_msg {
            cmd_type: u16::from_be_bytes([buf[0], buf[1]]),
            payload: [0u8; 64],
            payload_size: u16::from_be_bytes([buf[66], buf[67]]),
            tick: u64::from_be_bytes([buf[68], buf[69], buf[70], buf[71], buf[72], buf[73], buf[74], buf[75]]),
        };
        cmd.payload.copy_from_slice(&buf[2..66]);
        cmd
    }
}

impl Drop for Transport {
    fn drop(&mut self) {
        sys::close(self.socket);
    }
}