herodb/crates/herodb/protocol.rs

use core::fmt;

use crate::error::DBError;

#[derive(Debug, Clone)]
pub enum Protocol {
    SimpleString(String),
    BulkString(String),
    Null,
    Array(Vec<Protocol>),
    Error(String), // NEW
}

impl fmt::Display for Protocol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.decode().as_str())
    }
}

impl Protocol {
    pub fn from(protocol: &str) -> Result<(Self, &str), DBError> {
        let ret = match protocol.chars().nth(0) {
            Some('+') => Self::parse_simple_string_sfx(&protocol[1..]),
            Some('$') => Self::parse_bulk_string_sfx(&protocol[1..]),
            Some('*') => Self::parse_array_sfx(&protocol[1..]),
            _ => Err(DBError(format!(
                "[from] unsupported protocol: {:?}",
                protocol
            ))),
        };
        ret
    }

    pub fn from_vec(array: Vec<&str>) -> Self {
        let array = array
            .into_iter()
            .map(|x| Protocol::BulkString(x.to_string()))
            .collect();
        Protocol::Array(array)
    }

    #[inline]
    pub fn ok() -> Self {
        Protocol::SimpleString("ok".to_string())
    }

    #[inline]
    pub fn err(msg: &str) -> Self {
        Protocol::Error(msg.to_string())
    }

    #[inline]
    pub fn write_on_slave_err() -> Self {
        Self::err("DISALLOW WRITE ON SLAVE")
    }

    #[inline]
    pub fn psync_on_slave_err() -> Self {
        Self::err("PSYNC ON SLAVE IS NOT ALLOWED")
    }

    #[inline]
    pub fn none() -> Self {
        Self::SimpleString("none".to_string())
    }

    pub fn decode(&self) -> String {
        match self {
            Protocol::SimpleString(s) => s.to_string(),
            Protocol::BulkString(s) => s.to_string(),
            Protocol::Null => "".to_string(),
            Protocol::Array(s) => s.iter().map(|x| x.decode()).collect::<Vec<_>>().join(" "),
            Protocol::Error(s) => s.to_string(),
        }
    }

    pub fn encode(&self) -> String {
        match self {
            Protocol::SimpleString(s) => format!("+{}\r\n", s),
            Protocol::BulkString(s)   => format!("${}\r\n{}\r\n", s.len(), s),
            Protocol::Array(ss)       => {
                format!("*{}\r\n", ss.len()) + &ss.iter().map(|x| x.encode()).collect::<String>()
            }
            Protocol::Null            => "$-1\r\n".to_string(),
            Protocol::Error(s)        => format!("-{}\r\n", s), // proper RESP error
        }
    }

    fn parse_simple_string_sfx(protocol: &str) -> Result<(Self, &str), DBError> {
        match protocol.find("\r\n") {
            Some(x) => Ok((Self::SimpleString(protocol[..x].to_string()), &protocol[x + 2..])),
            _ => Err(DBError(format!(
                "[new simple string] unsupported protocol: {:?}",
                protocol
            ))),
        }
    }

    fn parse_bulk_string_sfx(protocol: &str) -> Result<(Self, &str), DBError> {
        if let Some(len_end) = protocol.find("\r\n") {
            let size = Self::parse_usize(&protocol[..len_end])?;
            let data_start = len_end + 2;
            let data_end = data_start + size;
            let s = Self::parse_string(&protocol[data_start..data_end])?;

            if protocol.len() < data_end + 2 || &protocol[data_end..data_end+2] != "\r\n" {
                 Err(DBError(format!(
                    "[new bulk string] unmatched string length in prototocl {:?}",
                    protocol,
                )))
            } else {
                Ok((Protocol::BulkString(s), &protocol[data_end + 2..]))
            }
        } else {
            Err(DBError(format!(
                "[new bulk string] unsupported protocol: {:?}",
                protocol
            )))
        }
    }

    fn parse_array_sfx(s: &str) -> Result<(Self, &str), DBError> {
        if let Some(len_end) = s.find("\r\n") {
            let array_len = s[..len_end].parse::<usize>()?;
            let mut remaining = &s[len_end + 2..];
            let mut vec = vec![];
            for _ in 0..array_len {
                let (p, rem) = Protocol::from(remaining)?;
                vec.push(p);
                remaining = rem;
            }
            Ok((Protocol::Array(vec), remaining))
        } else {
            Err(DBError(format!(
                "[new array] unsupported protocol: {:?}",
                s
            )))
        }
    }

    fn parse_usize(protocol: &str) -> Result<usize, DBError> {
        if protocol.is_empty() {
            Err(DBError("Cannot parse usize from empty string".to_string()))
        } else {
            protocol
                .parse::<usize>()
                .map_err(|_| DBError(format!("Failed to parse usize from: {}", protocol)))
        }
    }

    fn parse_string(protocol: &str) -> Result<String, DBError> {
        if protocol.is_empty() {
            // Allow empty strings, but handle appropriately
            Ok("".to_string())
        } else {
            Ok(protocol.to_string())
        }
    }
}