lavina/src/protos/xmpp/client.rs

use derive_more::From;
use quick_xml::events::attributes::Attribute;
use quick_xml::events::{BytesEnd, BytesStart, BytesText, 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")));
    }
}

#[derive(PartialEq, Eq, Debug)]
pub struct Presence<T> {
    pub to: Option<String>,
    pub from: Option<String>,
    pub priority: Option<PresencePriority>,
    pub show: Option<PresenceShow>,
    pub status: Vec<String>,
    pub custom: Vec<T>,
}

impl<T> Default for Presence<T> {
    fn default() -> Self {
        Self {
            to: Default::default(),
            from: Default::default(),
            priority: Default::default(),
            show: Default::default(),
            status: Default::default(),
            custom: Default::default(),
        }
    }
}

#[derive(PartialEq, Eq, Debug)]
pub enum PresenceShow {
    Away,
    Chat,
    Dnd,
    Xa,
}

/// Presence priority is an integer number in range [-128; 127].
///
/// Presence priority < 0 means that the bound resource will never be chosen unless it was asked for specifically.
#[derive(PartialEq, Eq, Debug)]
pub struct PresencePriority(pub i8);

impl PresenceShow {
    pub fn from_str(s: &[u8]) -> Result<Self> {
        use PresenceShow::*;
        let s = std::str::from_utf8(s)?;
        match s {
            "away" => Ok(Away),
            "chat" => Ok(Chat),
            "dnd" => Ok(Dnd),
            "xa" => Ok(Xa),
            t => Err(ffail!("Unknown presence show type: {t}")),
        }
    }
    pub fn as_str(&self) -> &'static str {
        use PresenceShow::*;
        match self {
            Away => "away",
            Chat => "chat",
            Dnd => "dnd",
            Xa => "xa",
        }
    }
}

#[derive(From)]
pub struct PresenceParser<T: FromXml>(PresenceParserInner<T>);
enum PresenceParserInner<T: FromXml> {
    Initial,
    InPresence(Presence<T>),
    InPriority(Presence<T>),
    InPriorityEnd(Presence<T>),
    InShow(Presence<T>),
    InShowEnd(Presence<T>),
    InStatus(Presence<T>),
    InStatusEnd(Presence<T>),
    InCustom(Presence<T>, T::P),
}

impl<T: FromXml> Parser for PresenceParser<T> {
    type Output = Result<Presence<T>>;

    fn consume<'a>(
        self: Self,
        namespace: ResolveResult,
        event: &Event<'a>,
    ) -> Continuation<Self, Self::Output> {
        match self.0 {
            PresenceParserInner::Initial => match event {
                //TODO validate
                Event::Start(bytes) => Continuation::Continue(
                    PresenceParserInner::InPresence(Presence::default()).into(),
                ),
                Event::Empty(bytes) => Continuation::Final(Ok(Presence::default())),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InPresence(p) => match event {
                Event::Start(bytes) => match bytes.name().0 {
                    b"show" => Continuation::Continue(PresenceParserInner::InShow(p).into()),
                    b"status" => Continuation::Continue(PresenceParserInner::InStatus(p).into()),
                    b"priority" => {
                        Continuation::Continue(PresenceParserInner::InPriority(p).into())
                    }
                    _ => PresenceParser(PresenceParserInner::InCustom(p, T::parse()))
                        .consume(namespace, event),
                },
                Event::Empty(_) => PresenceParser(PresenceParserInner::InCustom(p, T::parse()))
                    .consume(namespace, event),
                Event::End(_) => Continuation::Final(Ok(p)),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InPriority(p) => match event {
                Event::Text(bytes) => {
                    match (|| {
                        let s = std::str::from_utf8(bytes)?;
                        let i = s.parse()?;
                        Ok(Presence {
                            priority: Some(PresencePriority(i)),
                            ..p
                        })
                    })() {
                        Ok(res) => {
                            Continuation::Continue(PresenceParserInner::InPriorityEnd(res).into())
                        }
                        Err(e) => Continuation::Final(Err(e)),
                    }
                }
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InPriorityEnd(p) => match event {
                Event::End(_) => Continuation::Continue(PresenceParserInner::InPresence(p).into()),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InShow(p) => match event {
                Event::Text(bytes) => {
                    match (|| {
                        let i = PresenceShow::from_str(bytes)?;
                        Ok(Presence { show: Some(i), ..p })
                    })() {
                        Ok(res) => {
                            Continuation::Continue(PresenceParserInner::InShowEnd(res).into())
                        }
                        Err(e) => Continuation::Final(Err(e)),
                    }
                }
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InShowEnd(p) => match event {
                Event::End(_) => Continuation::Continue(PresenceParserInner::InPresence(p).into()),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InStatus(mut p) => match event {
                Event::Text(bytes) => {
                    match (|| {
                        let i = std::str::from_utf8(bytes)?;
                        p.status.push(i.to_string());
                        Ok(p)
                    })() {
                        Ok(res) => {
                            Continuation::Continue(PresenceParserInner::InStatusEnd(res).into())
                        }
                        Err(e) => Continuation::Final(Err(e)),
                    }
                }
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InStatusEnd(p) => match event {
                Event::End(_) => Continuation::Continue(PresenceParserInner::InPresence(p).into()),
                _ => Continuation::Final(Err(ffail!("Unexpected XML event: {event:?}"))),
            },
            PresenceParserInner::InCustom(mut p, c) => match c.consume(namespace, event) {
                Continuation::Final(Ok(res)) => {
                    p.custom.push(res);
                    Continuation::Continue(PresenceParserInner::InPresence(p).into())
                }
                Continuation::Final(Err(err)) => Continuation::Final(Err(ffail!(
                    "Failed to parse custom presence element: {err}"
                ))),
                Continuation::Continue(c) => {
                    Continuation::Continue(PresenceParserInner::InCustom(p, c).into())
                }
            },
        }
    }
}

impl<T: FromXml> FromXml for Presence<T> {
    type P = PresenceParser<T>;

    fn parse() -> Self::P {
        PresenceParserInner::Initial.into()
    }
}

impl<T: FromXml> FromXmlTag for Presence<T> {
    const NAME: &'static str = "presence";
    const NS: &'static str = XMLNS;
}

impl<T: ToXml> ToXml for Presence<T> {
    fn serialize(&self, events: &mut Vec<Event<'static>>) {
        let mut start = BytesStart::new("presence");
        if let Some(ref to) = self.to {
            start.extend_attributes([Attribute {
                key: QName(b"to"),
                value: to.as_bytes().into(),
            }]);
        }
        if let Some(ref from) = self.from {
            start.extend_attributes([Attribute {
                key: QName(b"from"),
                value: from.as_bytes().into(),
            }]);
        }
        events.push(Event::Start(start));
        if let Some(ref priority) = self.priority {
            let s = priority.0.to_string();
            events.extend_from_slice(&[
                Event::Start(BytesStart::new(r#"priority"#)),
                Event::Text(BytesText::new(s.as_str()).into_owned()),
                Event::End(BytesEnd::new("priority")),
            ]);
        }
        events.push(Event::End(BytesEnd::new("presence")));
    }
}

#[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()))
            }
        )
    }
}