MetaCrate/crates/libremetaverse/src/protocol_manager.rs

use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::{BufRead, BufReader, BufWriter, Read, Write};

use libremetaverse_types::compat::Object;

use crate::{Error, FieldType, GridClient, PacketFrequency};

#[derive(Clone, Debug)]
pub struct MapField {
    pub count: i32,
    pub keyword_position: i32,
    pub name: String,
    pub type_: FieldType,
}

impl Default for MapField {
    fn default() -> Self {
        Self {
            count: 0,
            keyword_position: 0,
            name: String::new(),
            type_: FieldType::U8,
        }
    }
}

impl MapField {
    pub fn new() -> Result<Self, Error> {
        Ok(Self::default())
    }

    pub fn compare_to(&self, obj: Option<Object>) -> Result<i32, Error> {
        let Some(obj) = obj else {
            return Ok(1);
        };
        let other = obj.downcast_ref::<Self>().ok_or(Error::Argument)?;
        Ok(self.keyword_position.cmp(&other.keyword_position) as i32)
    }
}

#[derive(Clone, Debug, Default)]
pub struct MapBlock {
    pub count: i32,
    pub fields: Vec<MapField>,
    pub keyword_position: i32,
    pub name: String,
}

impl MapBlock {
    pub fn new() -> Result<Self, Error> {
        Ok(Self::default())
    }

    pub fn compare_to(&self, obj: Option<Object>) -> Result<i32, Error> {
        let Some(obj) = obj else {
            return Ok(1);
        };
        let other = obj.downcast_ref::<Self>().ok_or(Error::Argument)?;
        Ok(self.keyword_position.cmp(&other.keyword_position) as i32)
    }
}

#[derive(Clone, Debug)]
pub struct MapPacket {
    pub blocks: Vec<MapBlock>,
    pub encoded: bool,
    pub frequency: PacketFrequency,
    pub id: u16,
    pub name: String,
    pub trusted: bool,
}

impl Default for MapPacket {
    fn default() -> Self {
        Self {
            blocks: Vec::new(),
            encoded: false,
            frequency: PacketFrequency::Low,
            id: 0,
            name: String::new(),
            trusted: false,
        }
    }
}

impl MapPacket {
    pub fn new() -> Result<Self, Error> {
        Ok(Self::default())
    }
}

pub struct ProtocolManager {
    pub high_maps: Vec<MapPacket>,
    pub keyword_positions: HashMap<String, i32>,
    pub low_maps: Vec<MapPacket>,
    pub medium_maps: Vec<MapPacket>,
    pub type_sizes: HashMap<FieldType, i32>,
}

impl ProtocolManager {
    pub fn new(map_file: String, client: GridClient) -> Result<Self, Error> {
        // The C# instance retains this reference only to enrich parse-error logging.
        // Rust reports parsing failures to the caller, so no client ownership is needed.
        drop(client);
        let mut manager = Self::empty();
        manager.load_map_file(&map_file)?;
        Ok(manager)
    }

    fn empty() -> Self {
        let type_sizes = HashMap::from([
            (FieldType::U8, 1),
            (FieldType::U16, 2),
            (FieldType::U32, 4),
            (FieldType::U64, 8),
            (FieldType::S8, 1),
            (FieldType::S16, 2),
            (FieldType::S32, 4),
            (FieldType::F32, 4),
            (FieldType::F64, 8),
            (FieldType::UUID, 16),
            (FieldType::BOOL, 1),
            (FieldType::Vector3, 12),
            (FieldType::Vector3d, 24),
            (FieldType::Vector4, 16),
            (FieldType::Quaternion, 16),
            (FieldType::IPADDR, 4),
            (FieldType::IPPORT, 2),
            (FieldType::Variable, -1),
            (FieldType::Fixed, -2),
        ]);
        Self {
            high_maps: vec![MapPacket::default(); 256],
            keyword_positions: HashMap::new(),
            low_maps: vec![MapPacket::default(); 65_536],
            medium_maps: vec![MapPacket::default(); 256],
            type_sizes,
        }
    }

    pub fn command_with_bytes(&self, data: Vec<u8>) -> Result<Option<MapPacket>, Error> {
        let Some(&frequency_byte) = data.get(4) else {
            return Ok(None);
        };
        let result = if frequency_byte != 0xff {
            self.command_with_u_int16_packet_frequency(
                u16::from(frequency_byte),
                PacketFrequency::High,
            )
        } else if data.get(5) != Some(&0xff) {
            let Some(&command) = data.get(5) else {
                return Ok(None);
            };
            self.command_with_u_int16_packet_frequency(u16::from(command), PacketFrequency::Medium)
        } else {
            let (Some(&high), Some(&low)) = (data.get(6), data.get(7)) else {
                return Ok(None);
            };
            self.command_with_u_int16_packet_frequency(
                u16::from_be_bytes([high, low]),
                PacketFrequency::Low,
            )
        };
        match result {
            Ok(packet) => Ok(Some(packet)),
            Err(Error::IndexOutOfRange | Error::InvalidOperation) => Ok(None),
            Err(error) => Err(error),
        }
    }

    pub fn command_with_string(&self, command: String) -> Result<MapPacket, Error> {
        self.high_maps
            .iter()
            .chain(&self.medium_maps)
            .chain(&self.low_maps)
            .find(|packet| !packet.name.is_empty() && packet.name == command)
            .cloned()
            .ok_or(Error::InvalidOperation)
    }

    pub fn command_with_u_int16_packet_frequency(
        &self,
        command: u16,
        frequency: PacketFrequency,
    ) -> Result<MapPacket, Error> {
        let maps = match frequency {
            PacketFrequency::High => &self.high_maps,
            PacketFrequency::Medium => &self.medium_maps,
            PacketFrequency::Low => &self.low_maps,
        };
        let packet = maps
            .get(usize::from(command))
            .ok_or(Error::IndexOutOfRange)?;
        if packet.name.is_empty() {
            return Err(Error::InvalidOperation);
        }
        Ok(packet.clone())
    }

    pub fn decode_map_file(map_file: String, output_file: String) -> Result<(), Error> {
        let input = File::open(map_file).map_err(|_| Error::InvalidOperation)?;
        let output = OpenOptions::new()
            .write(true)
            .create_new(true)
            .open(output_file)
            .map_err(|_| Error::InvalidOperation)?;
        Self::decode(BufReader::new(input), BufWriter::new(output))
    }

    fn decode(mut input: impl Read, mut output: impl Write) -> Result<(), Error> {
        let mut key = 0_u8;
        let mut buffer = [0_u8; 2048];
        loop {
            let count = input
                .read(&mut buffer)
                .map_err(|_| Error::InvalidOperation)?;
            if count == 0 {
                break;
            }
            for byte in &mut buffer[..count] {
                *byte ^= key;
                key = key.wrapping_add(43);
            }
            output
                .write_all(&buffer[..count])
                .map_err(|_| Error::InvalidOperation)?;
        }
        output.flush().map_err(|_| Error::InvalidOperation)
    }

    pub fn print_map(&self) -> Result<(), Error> {
        self.print_one_map(&self.low_maps, "Low   ");
        self.print_one_map(&self.medium_maps, "Medium");
        self.print_one_map(&self.high_maps, "High  ");
        Ok(())
    }

    fn print_one_map(&self, maps: &[MapPacket], frequency: &str) {
        for (index, packet) in maps
            .iter()
            .enumerate()
            .filter(|(_, packet)| !packet.name.is_empty())
        {
            println!(
                "{frequency} {index:5} - {} - {} - {}",
                packet.name,
                if packet.trusted {
                    "Trusted"
                } else {
                    "Untrusted"
                },
                if packet.encoded {
                    "Unencoded"
                } else {
                    "Zerocoded"
                }
            );
            for block in &packet.blocks {
                if block.count == -1 {
                    println!("\t{:4} {} (Variable)", block.keyword_position, block.name);
                } else {
                    println!(
                        "\t{:4} {} ({})",
                        block.keyword_position, block.name, block.count
                    );
                }
                for field in &block.fields {
                    println!(
                        "\t\t{:4} {} ({:?} / {})",
                        field.keyword_position, field.name, field.type_, field.count
                    );
                }
            }
        }
    }

    fn load_map_file(&mut self, map_file: &str) -> Result<(), Error> {
        let input = File::open(map_file).map_err(|_| Error::InvalidOperation)?;
        self.parse(BufReader::new(input))
    }

    fn parse(&mut self, input: impl BufRead) -> Result<(), Error> {
        let mut low = 1_u16;
        let mut medium = 1_u16;
        let mut high = 1_u16;
        let mut in_packet = false;
        let mut in_block = false;
        let mut current_packet: Option<(PacketFrequency, usize)> = None;
        let mut current_block: Option<MapBlock> = None;

        for (line_number, line) in input.lines().enumerate() {
            let line = line.map_err(|_| Error::Parse {
                position: line_number,
                context: "could not read protocol map line",
            })?;
            let trimmed = line.split_whitespace().collect::<Vec<_>>().join(" ");
            if !in_packet {
                if trimmed == "{" {
                    in_packet = true;
                }
                continue;
            }

            if !in_block {
                if trimmed == "{" {
                    if current_packet.is_none() {
                        return Err(Self::parse_error(line_number, "block has no packet header"));
                    }
                    in_block = true;
                } else if trimmed == "}" {
                    if let Some((frequency, index)) = current_packet.take() {
                        self.packet_mut(frequency, index)?
                            .blocks
                            .sort_by_key(|block| block.keyword_position);
                    }
                    in_packet = false;
                } else if !trimmed.is_empty() && !trimmed.starts_with("//") {
                    current_packet = Some(self.parse_packet_header(
                        &trimmed,
                        line_number,
                        &mut low,
                        &mut medium,
                        &mut high,
                    )?);
                }
                continue;
            }

            if trimmed.starts_with('{') {
                let tokens: Vec<_> = trimmed.split_whitespace().collect();
                if tokens.len() < 4 {
                    return Err(Self::parse_error(
                        line_number,
                        "malformed field declaration",
                    ));
                }
                let type_ = Self::parse_field_type(tokens[2])
                    .ok_or_else(|| Self::parse_error(line_number, "unknown field type"))?;
                let count = if tokens[3] == "}" {
                    1
                } else {
                    tokens[3]
                        .parse()
                        .map_err(|_| Self::parse_error(line_number, "invalid field count"))?
                };
                let keyword_position = self.keyword_position(tokens[1])?;
                let block = current_block
                    .as_mut()
                    .ok_or_else(|| Self::parse_error(line_number, "field has no block header"))?;
                block.fields.push(MapField {
                    count,
                    keyword_position,
                    name: tokens[1].to_owned(),
                    type_,
                });
            } else if trimmed == "}" {
                let mut block = current_block.take().ok_or_else(|| {
                    Self::parse_error(line_number, "block terminator has no header")
                })?;
                block.fields.sort_by_key(|field| field.keyword_position);
                let (frequency, index) = current_packet
                    .ok_or_else(|| Self::parse_error(line_number, "block has no packet"))?;
                self.packet_mut(frequency, index)?.blocks.push(block);
                in_block = false;
            } else if !trimmed.is_empty() && !trimmed.starts_with("//") {
                if current_block.is_some() {
                    return Err(Self::parse_error(line_number, "block header is incomplete"));
                }
                let tokens: Vec<_> = trimmed.split_whitespace().collect();
                if tokens.len() < 2 {
                    return Err(Self::parse_error(line_number, "malformed block header"));
                }
                let count = match tokens[1] {
                    "Single" => 1,
                    "Variable" => -1,
                    "Multiple" => tokens
                        .get(2)
                        .ok_or_else(|| Self::parse_error(line_number, "missing block count"))?
                        .parse()
                        .map_err(|_| Self::parse_error(line_number, "invalid block count"))?,
                    _ => return Err(Self::parse_error(line_number, "unknown block frequency")),
                };
                current_block = Some(MapBlock {
                    count,
                    fields: Vec::new(),
                    keyword_position: self.keyword_position(tokens[0])?,
                    name: tokens[0].to_owned(),
                });
            }
        }

        if in_packet || in_block || current_block.is_some() {
            return Err(Self::parse_error(
                0,
                "unterminated protocol map declaration",
            ));
        }
        Ok(())
    }

    fn parse_packet_header(
        &mut self,
        line: &str,
        position: usize,
        low: &mut u16,
        medium: &mut u16,
        high: &mut u16,
    ) -> Result<(PacketFrequency, usize), Error> {
        let tokens: Vec<_> = line.split_whitespace().collect();
        if tokens.len() < 4 {
            return Err(Self::parse_error(position, "malformed packet header"));
        }
        self.keyword_position(tokens[0])?;
        let (frequency, id, trusted_index, encoded_index) = match tokens[1] {
            "Fixed" => {
                if tokens.len() < 5 {
                    return Err(Self::parse_error(position, "malformed fixed packet header"));
                }
                let raw = tokens[2].strip_prefix("0x").unwrap_or(tokens[2]);
                let raw = u32::from_str_radix(raw, 16)
                    .map_err(|_| Self::parse_error(position, "invalid fixed packet id"))?;
                (PacketFrequency::Low, (raw ^ 0xffff_0000) as u16, 3, 4)
            }
            "Low" => {
                let id = *low;
                *low = low.checked_add(1).ok_or(Error::IndexOutOfRange)?;
                (PacketFrequency::Low, id, 2, 3)
            }
            "Medium" => {
                let id = *medium;
                *medium = medium.checked_add(1).ok_or(Error::IndexOutOfRange)?;
                (PacketFrequency::Medium, id, 2, 3)
            }
            "High" => {
                let id = *high;
                *high = high.checked_add(1).ok_or(Error::IndexOutOfRange)?;
                (PacketFrequency::High, id, 2, 3)
            }
            _ => return Err(Self::parse_error(position, "unknown packet frequency")),
        };
        let packet = MapPacket {
            blocks: Vec::new(),
            encoded: tokens.get(encoded_index) == Some(&"Zerocoded"),
            frequency,
            id,
            name: tokens[0].to_owned(),
            trusted: tokens.get(trusted_index) == Some(&"Trusted"),
        };
        let index = usize::from(id);
        *self.packet_mut(frequency, index)? = packet;
        Ok((frequency, index))
    }

    fn packet_mut(
        &mut self,
        frequency: PacketFrequency,
        index: usize,
    ) -> Result<&mut MapPacket, Error> {
        match frequency {
            PacketFrequency::High => self.high_maps.get_mut(index),
            PacketFrequency::Medium => self.medium_maps.get_mut(index),
            PacketFrequency::Low => self.low_maps.get_mut(index),
        }
        .ok_or(Error::IndexOutOfRange)
    }

    fn keyword_position(&mut self, keyword: &str) -> Result<i32, Error> {
        if let Some(position) = self.keyword_positions.get(keyword) {
            return Ok(*position);
        }
        let mut hash = keyword
            .chars()
            .skip(1)
            .fold(0_i32, |hash, character| {
                hash.wrapping_add(character as i32).wrapping_mul(2)
            })
            .wrapping_mul(2)
            & 0x1fff;
        let start = hash;
        while self
            .keyword_positions
            .values()
            .any(|position| *position == hash)
        {
            hash = (hash + 1) & 0x1fff;
            if hash == start {
                return Err(Error::InvalidOperation);
            }
        }
        self.keyword_positions.insert(keyword.to_owned(), hash);
        Ok(hash)
    }

    fn parse_field_type(value: &str) -> Option<FieldType> {
        Some(match value.to_ascii_lowercase().as_str() {
            "u8" => FieldType::U8,
            "u16" => FieldType::U16,
            "u32" => FieldType::U32,
            "u64" => FieldType::U64,
            "s8" => FieldType::S8,
            "s16" => FieldType::S16,
            "s32" => FieldType::S32,
            "f32" => FieldType::F32,
            "f64" => FieldType::F64,
            "uuid" => FieldType::UUID,
            "bool" => FieldType::BOOL,
            "vector3" => FieldType::Vector3,
            "vector3d" => FieldType::Vector3d,
            "vector4" => FieldType::Vector4,
            "quaternion" => FieldType::Quaternion,
            "ipaddr" => FieldType::IPADDR,
            "ipport" => FieldType::IPPORT,
            "variable" => FieldType::Variable,
            "fixed" => FieldType::Fixed,
            "single" => FieldType::Single,
            "multiple" => FieldType::Multiple,
            _ => return None,
        })
    }

    const fn parse_error(position: usize, context: &'static str) -> Error {
        Error::Parse { position, context }
    }
}

#[cfg(test)]
mod tests {
    use std::io::Cursor;

    use super::*;

    const MAP: &str = r#"
// version header
{
TestHigh High Trusted Zerocoded
{
AgentData Single
{ AgentID UUID }
{ SessionID UUID }
}
}
{
TestLow Fixed 0xFFFF00F2 Untrusted Unencoded
{
Data Variable
{ Value U32 }
}
}
"#;

    #[test]
    fn parses_and_resolves_protocol_commands() {
        let mut manager = ProtocolManager::empty();
        manager.parse(Cursor::new(MAP)).unwrap();

        let high = manager.command_with_string("TestHigh".to_owned()).unwrap();
        assert_eq!(high.frequency, PacketFrequency::High);
        assert_eq!(high.id, 1);
        assert!(high.trusted);
        assert!(high.encoded);
        assert_eq!(high.blocks[0].fields.len(), 2);

        let low = manager
            .command_with_bytes(vec![0, 0, 0, 0, 0xff, 0xff, 0, 0xf2])
            .unwrap()
            .unwrap();
        assert_eq!(low.name, "TestLow");
        assert_eq!(low.blocks[0].count, -1);
        assert_eq!(manager.type_sizes[&FieldType::Vector3d], 24);
    }

    #[test]
    fn xor_decoder_tracks_key_across_chunks() {
        let input: Vec<_> = (0_u16..4097).map(|value| value as u8).collect();
        let mut encoded = Vec::new();
        ProtocolManager::decode(Cursor::new(&input), &mut encoded).unwrap();
        let mut decoded = Vec::new();
        ProtocolManager::decode(Cursor::new(encoded), &mut decoded).unwrap();
        assert_eq!(decoded, input);
    }

    #[test]
    fn malformed_short_packet_headers_do_not_panic() {
        let manager = ProtocolManager::empty();
        assert!(manager.command_with_bytes(vec![0; 5]).unwrap().is_none());
        assert!(
            manager
                .command_with_bytes(vec![0, 0, 0, 0, 0xff])
                .unwrap()
                .is_none()
        );
        assert!(
            manager
                .command_with_bytes(vec![0, 0, 0, 0, 0xff, 0xff, 0])
                .unwrap()
                .is_none()
        );
    }
}