lavina/src/protos/xmpp/client.rs

use quick_xml::events::Event;
use quick_xml::name::ResolveResult;

use crate::prelude::*;
use crate::util::xml::*;

pub static XMLNS: &'static str = "jabber:client";

#[derive(PartialEq, Eq, Debug)]
pub struct Message {
    pub from: Option<String>,
    pub id: Option<String>,
    pub to: Option<String>,
    // default is Normal
    pub r#type: MessageType,
    pub lang: Option<String>,

    pub subject: Option<String>,
    pub body: String,
}
impl Message {
    pub fn parse() -> impl Parser<Output = Result<Self>> {
        MessageParser::Init
    }
}

#[derive(Default)]
enum MessageParser {
    #[default]
    Init,
    Outer(MessageParserState),
    InSubject(MessageParserState),
    InBody(MessageParserState),
}
#[derive(Default)]
struct MessageParserState {
    from: Option<String>,
    id: Option<String>,
    to: Option<String>,
    r#type: MessageType,
    lang: Option<String>,
    subject: Option<String>,
    body: Option<String>,
}
impl Parser for MessageParser {
    type Output = Result<Message>;

    fn consume<'a>(
        self: Self,
        namespace: ResolveResult,
        event: &Event<'a>,
    ) -> Continuation<Self, Self::Output> {
        // TODO validate tag name and namespace at each stage
        match self {
            MessageParser::Init => {
                if let Event::Start(ref bytes) = event {
                    let mut state: MessageParserState = Default::default();
                    for attr in bytes.attributes() {
                        let attr = fail_fast!(attr);
                        if attr.key.0 == b"from" {
                            let value = fail_fast!(std::str::from_utf8(&*attr.value));
                            state.from = Some(value.to_string())
                        } else if attr.key.0 == b"id" {
                            let value = fail_fast!(std::str::from_utf8(&*attr.value));
                            state.id = Some(value.to_string())
                        } else if attr.key.0 == b"to" {
                            let value = fail_fast!(std::str::from_utf8(&*attr.value));
                            state.to = Some(value.to_string())
                        } else if attr.key.0 == b"type" {
                            let value = fail_fast!(MessageType::from_str(&*attr.value));
                            state.r#type = value;
                        }
                    }
                    Continuation::Continue(MessageParser::Outer(state))
                } else {
                    Continuation::Final(Err(ffail!("Expected start")))
                }
            }
            MessageParser::Outer(state) => match event {
                Event::Start(ref bytes) => {
                    if bytes.name().0 == b"subject" {
                        Continuation::Continue(MessageParser::InSubject(state))
                    } else if bytes.name().0 == b"body" {
                        Continuation::Continue(MessageParser::InBody(state))
                    } else {
                        Continuation::Final(Err(ffail!("Unexpected XML tag")))
                    }
                }
                Event::End(_) => {
                    if let Some(body) = state.body {
                        Continuation::Final(Ok(Message {
                            from: state.from,
                            id: state.id,
                            to: state.to,
                            r#type: state.r#type,
                            lang: state.lang,
                            subject: state.subject,
                            body,
                        }))
                    } else {
                        Continuation::Final(Err(ffail!("Body not found")))
                    }
                }
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            MessageParser::InSubject(mut state) => match event {
                Event::Text(ref bytes) => {
                    let subject = fail_fast!(std::str::from_utf8(&*bytes));
                    state.subject = Some(subject.to_string());
                    Continuation::Continue(MessageParser::InSubject(state))
                }
                Event::End(_) => Continuation::Continue(MessageParser::Outer(state)),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            MessageParser::InBody(mut state) => match event {
                Event::Text(ref bytes) => match std::str::from_utf8(&*bytes) {
                    Ok(subject) => {
                        state.body = Some(subject.to_string());
                        Continuation::Continue(MessageParser::InBody(state))
                    }
                    Err(err) => Continuation::Final(Err(err.into())),
                },
                Event::End(_) => Continuation::Continue(MessageParser::Outer(state)),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
        }
    }
}

#[derive(PartialEq, Eq, Debug)]
pub enum MessageType {
    Chat,
    Error,
    Groupchat,
    Headline,
    Normal,
}

impl Default for MessageType {
    fn default() -> Self {
        MessageType::Normal
    }
}

impl MessageType {
    pub fn from_str(s: &[u8]) -> Result<MessageType> {
        use MessageType::*;
        let s = std::str::from_utf8(s)?;
        match s {
            "chat" => Ok(Chat),
            "error" => Ok(Error),
            "groupchat" => Ok(Groupchat),
            "headline" => Ok(Headline),
            "normal" => Ok(Normal),
            t => Err(ffail!("Unknown message type: {t}")),
        }
    }
}

pub struct Iq<T> {
    pub from: Option<String>,
    pub id: String,
    pub to: Option<String>,
    pub r#type: IqType,
    pub body: T,
}

impl<T: FromXml> FromXml for Iq<T> {
    type P = IqParser<T>;

    fn parse() -> Self::P {
        IqParser(IqParserInner::Init)
    }
}

pub struct IqParser<T: FromXml>(IqParserInner<T>);

enum IqParserInner<T: FromXml> {
    Init,
    ParsingBody(IqParserState<T>, T::P),
    Final(IqParserState<T>),
}
struct IqParserState<T> {
    pub from: Option<String>,
    pub id: Option<String>,
    pub to: Option<String>,
    pub r#type: Option<IqType>,
    pub body: Option<T>,
}

impl<T: FromXml> Parser for IqParser<T> {
    type Output = Result<Iq<T>>;

    fn consume<'a>(
        self: Self,
        namespace: ResolveResult,
        event: &Event<'a>,
    ) -> Continuation<Self, Self::Output> {
        match self.0 {
            IqParserInner::Init => {
                if let Event::Start(ref bytes) = event {
                    let mut state: IqParserState<T> = IqParserState { from: None, id: None, to: None, r#type: None, body: None };
                    for attr in bytes.attributes() {
                        let attr = fail_fast!(attr);
                        if attr.key.0 == b"from" {
                            let value = fail_fast!(std::str::from_utf8(&*attr.value));
                            state.from = Some(value.to_string())
                        } else if attr.key.0 == b"id" {
                            let value = fail_fast!(std::str::from_utf8(&*attr.value));
                            state.id = Some(value.to_string())
                        } else if attr.key.0 == b"to" {
                            let value = fail_fast!(std::str::from_utf8(&*attr.value));
                            state.to = Some(value.to_string())
                        } else if attr.key.0 == b"type" {
                            let value = fail_fast!(IqType::from_str(&*attr.value));
                            state.r#type = Some(value);
                        }
                    }
                    Continuation::Continue(IqParser(IqParserInner::ParsingBody(state, T::parse())))
                } else {
                    Continuation::Final(Err(ffail!("Expected start")))
                }
            }
            IqParserInner::ParsingBody(mut state, parser) => {
                match parser.consume(namespace, event) {
                    Continuation::Final(f) => {
                        let body = fail_fast!(f);
                        state.body = Some(body);
                        Continuation::Continue(IqParser(IqParserInner::Final(state)))
                    }
                    Continuation::Continue(parser) => {
                        Continuation::Continue(IqParser(IqParserInner::ParsingBody(state, parser)))
                    }
                }
            }
            IqParserInner::Final(state) => {
                if let Event::End(ref bytes) = event {
                    let id = fail_fast!(state.id.ok_or_else(|| fail("No id provided")));
                    let r#type = fail_fast!(state.r#type.ok_or_else(|| fail("No type provided")));
                    let body = fail_fast!(state.body.ok_or_else(|| fail("No body provided")));
                    Continuation::Final(Ok(Iq {
                        from: state.from,
                        id,
                        to: state.to,
                        r#type,
                        body,
                    }))
                } else {
                    Continuation::Final(Err(ffail!("Unexpected event: {event:?}")))
                }
            }
        }
    }
}

pub enum IqType {
    Error,
    Get,
    Result,
    Set,
}

impl IqType {
    pub fn from_str(s: &[u8]) -> Result<IqType> {
        use IqType::*;
        let s = std::str::from_utf8(s)?;
        match s {
            "error" => Ok(Error),
            "get" => Ok(Get),
            "result" => Ok(Result),
            "set" => Ok(Set),
            t => Err(ffail!("Unknown iq type: {t}")),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use quick_xml::NsReader;

    #[tokio::test]
    async fn parse_message() {
        let input = r#"<message id="aacea" type="chat" to="nikita@vlnv.dev"><subject>daa</subject><body>bbb</body></message>"#;
        let mut reader = NsReader::from_reader(input.as_bytes());
        let mut buf = vec![];
        let (ns, event) = reader.read_resolved_event_into_async(&mut buf).await.unwrap();
        let mut parser = Message::parse().consume(ns, &event);
        let result = loop {
            match parser {
                Continuation::Final(res) => break res,
                Continuation::Continue(next) => {
                    let (ns, event) = reader.read_resolved_event_into_async(&mut buf).await.unwrap();
                    parser = next.consume(ns, &event);
                }
            }
        }
        .unwrap();
        assert_eq!(
            result,
            Message {
                from: None,
                id: Some("aacea".to_string()),
                to: Some("nikita@vlnv.dev".to_string()),
                r#type: MessageType::Chat,
                lang: None,
                subject: Some("daa".to_string()),
                body: "bbb".to_string(),
            }
        )
    }
}