lavina/src/protos/xmpp/client.rs

use derive_more::From;
use quick_xml::events::attributes::Attribute;
use quick_xml::events::{BytesEnd, BytesStart, Event};
use quick_xml::name::{QName, ResolveResult};

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

pub const 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 FromXmlTag for Message {
    const NS: &'static str = XMLNS;
    const NAME: &'static str = "message";
}

impl FromXml for Message {
    type P = MessageParser;

    fn parse() -> Self::P {
        MessageParserInner::Init.into()
    }
}

#[derive(From)]
pub struct MessageParser(MessageParserInner);

#[derive(Default)]
enum MessageParserInner {
    #[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
        use MessageParserInner::*;
        match self.0 {
            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(Outer(state).into())
                } else {
                    Continuation::Final(Err(ffail!("Expected start")))
                }
            }
            Outer(state) => match event {
                Event::Start(ref bytes) => {
                    if bytes.name().0 == b"subject" {
                        Continuation::Continue(InSubject(state).into())
                    } else if bytes.name().0 == b"body" {
                        Continuation::Continue(InBody(state).into())
                    } 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:?}"))),
            },
            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(InSubject(state).into())
                }
                Event::End(_) => Continuation::Continue(Outer(state).into()),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            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(InBody(state).into())
                    }
                    Err(err) => Continuation::Final(Err(err.into())),
                },
                Event::End(_) => Continuation::Continue(Outer(state).into()),
                _ => 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}")),
        }
    }
}

#[derive(PartialEq, Eq, Debug)]
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> FromXmlTag for Iq<T> {
    const NAME: &'static str = "iq";
    const NS: &'static str = XMLNS;
}

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:?}")))
                }
            }
        }
    }
}

#[derive(PartialEq, Eq, Debug)]
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}")),
        }
    }
    pub fn as_str(&self) -> &'static str {
        match self {
            IqType::Error => "error",
            IqType::Get => "get",
            IqType::Result => "result",
            IqType::Set => "set",
        }
    }
}

impl<T: ToXml> ToXml for Iq<T> {
    fn serialize(&self, events: &mut Vec<Event<'static>>) {
        let start = format!(r#"iq xmlns="{}""#, XMLNS);
        let mut start = BytesStart::new(start);
        let mut attrs = vec![];
        if let Some(ref from) = self.from {
            attrs.push(Attribute {
                key: QName(b"from"),
                value: from.as_bytes().into(),
            });
        };
        if let Some(ref to) = self.to {
            attrs.push(Attribute {
                key: QName(b"to"),
                value: to.as_bytes().into(),
            });
        }
        attrs.push(Attribute {
            key: QName(b"id"),
            value: self.id.as_bytes().into(),
        });
        attrs.push(Attribute {
            key: QName(b"type"),
            value: self.r#type.as_str().as_bytes().into(),
        });
        start.extend_attributes(attrs.into_iter());

        events.push(Event::Start(start));
        self.body.serialize(events);
        events.push(Event::End(BytesEnd::new("iq")));
    }
}

#[cfg(test)]
mod tests {
    use crate::protos::xmpp::bind::{BindRequest, Resource};

    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(),
            }
        )
    }

    #[tokio::test]
    async fn parse_iq() {
        let input = r#"<iq id="bind_1" type="set"><bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"><resource>mobile</resource></bind></iq>"#;
        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 = Iq::<BindRequest>::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,
            Iq {
                from: None,
                id: "bind_1".to_string(),
                to: None,
                r#type: IqType::Set,
                body: BindRequest(Resource("mobile".to_string()))
            }
        )
    }
}