lavina/crates/proto-xmpp/src/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 anyhow::{Result, anyhow as ffail};


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

use super::bind::Jid;

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

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

    pub subject: Option<Str>,
    pub body: Str,
}

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<Jid>,
    id: Option<String>,
    to: Option<Jid>,
    r#type: MessageType,
    lang: Option<Str>,
    subject: Option<Str>,
    body: Option<Str>,
}
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));
                            let value = fail_fast!(Jid::from_string(value));
                            state.from = Some(value)
                        } 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));
                            let value = fail_fast!(Jid::from_string(value));
                            state.to = Some(value)
                        } 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 event: {event:?}")))
                    }
                }
                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.into());
                    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.into());
                        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:?}"))),
            },
        }
    }
}

impl ToXml for Message {
    fn serialize(&self, events: &mut Vec<Event<'static>>) {
        let mut bytes = BytesStart::new(format!(r#"message xmlns="{}""#, XMLNS));
        if let Some(from) = &self.from {
            bytes.push_attribute(Attribute {
                key: QName(b"from"),
                value: from.to_string().into_bytes().into(),
            });
        }
        if let Some(to) = &self.to {
            bytes.push_attribute(Attribute {
                key: QName(b"to"),
                value: to.to_string().into_bytes().into(),
            });
        }
        if let Some(id) = &self.id {
            bytes.push_attribute(Attribute {
                key: QName(b"id"),
                value: id.clone().into_bytes().into(),
            });
        }
        bytes.push_attribute(Attribute {
            key: QName(b"type"),
            value: self.r#type.as_str().as_bytes().into(),
        });
        events.push(Event::Start(bytes));
        events.push(Event::Start(BytesStart::new("body")));
        events.push(Event::Text(BytesText::new(&self.body).into_owned()));
        events.push(Event::End(BytesEnd::new("body")));
        events.push(Event::End(BytesEnd::new("message")));
    }
}

#[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 fn as_str(&self) -> &'static str {
        match self {
            MessageType::Chat => "chat",
            MessageType::Error => "error",
            MessageType::Groupchat => "groupchat",
            MessageType::Headline => "headline",
            MessageType::Normal => "normal",
        }
    }
}

#[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(|| ffail!("No id provided")));
                    let r#type = fail_fast!(state.r#type.ok_or_else(|| ffail!("No type provided")));
                    let body = fail_fast!(state.body.ok_or_else(|| ffail!("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<Jid>,
    pub from: Option<Jid>,
    pub priority: Option<PresencePriority>,
    pub show: Option<PresenceShow>,
    pub status: Vec<String>,
    pub custom: Vec<T>,
    pub r#type: Option<String>,
}

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(),
            r#type: None,
        }
    }
}

#[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",
        }
    }
}

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

    fn parse() -> Self::P {
        |(namespace, event): (ResolveResult<'static>, &'static Event<'static>)| -> Result<Self> {
            let (bytes, end) = match event {
                Event::Start(bytes) => (bytes, false),
                Event::Empty(bytes) => (bytes, true),
                _ => return Err(ffail!("Unexpected XML event: {event:?}")),
            };
            let mut p = Presence::<T>::default();
            for attr in bytes.attributes() {
                let attr = attr?;
                match attr.key.0 {
                    b"to" => {
                        let s = std::str::from_utf8(&attr.value)?;
                        p.to = Some(Jid::from_string(s)?);
                    }
                    b"from" => {
                        let s = std::str::from_utf8(&attr.value)?;
                        p.to = Some(Jid::from_string(s)?);
                    }
                    b"type" => {
                        let s = std::str::from_utf8(&attr.value)?;
                        p.r#type = Some(s.into());
                    }
                    _ => {}
                }
            }
            if end {
                return Ok(p);
            }
            loop {
                let (namespace, event) = yield;
                match event {
                    Event::Start(bytes) => match bytes.name().0 {
                        b"show" => {
                            let (_, event) = yield;
                            let Event::Text(bytes) = event else {
                                return Err(ffail!("Unexpected XML event: {event:?}"));
                            };
                            let i = PresenceShow::from_str(bytes)?;
                            p.show = Some(i);

                            let (_, event) = yield;
                            let Event::End(_) = event else {
                                return Err(ffail!("Unexpected XML event: {event:?}"));
                            };
                        }
                        b"status" => {
                            let (_, event) = yield;
                            let Event::Text(bytes) = event else {
                                return Err(ffail!("Unexpected XML event: {event:?}"));
                            };
                            let s = std::str::from_utf8(bytes)?;
                            p.status.push(s.to_string());

                            let (_, event) = yield;
                            let Event::End(_) = event else {
                                return Err(ffail!("Unexpected XML event: {event:?}"));
                            };
                        }
                        b"priority" => {
                            let (_, event) = yield;
                            let Event::Text(bytes) = event else {
                                return Err(ffail!("Unexpected XML event: {event:?}"));
                            };
                            let s = std::str::from_utf8(bytes)?;
                            let i = s.parse()?;
                            p.priority = Some(PresencePriority(i));

                            let (_, event) = yield;
                            let Event::End(_) = event else {
                                return Err(ffail!("Unexpected XML event: {event:?}"));
                            };
                        }
                        _ => {
                            let res = delegate_parsing!(T, namespace, event);
                            p.custom.push(res?);
                        }
                    },
                    Event::Empty(_) => {
                        let res = delegate_parsing!(T, namespace, event);
                        p.custom.push(res?);
                    }
                    Event::End(_) => return Ok(p),
                    _ => return Err(ffail!("Unexpected XML event: {event:?}")),
                }
            }
        }
    }
}

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.to_string().as_bytes().into(),
            }]);
        }
        if let Some(ref from) = self.from {
            start.extend_attributes([Attribute {
                key: QName(b"from"),
                value: from.to_string().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::bind::{BindRequest, Name, Resource, Server};

    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(Jid {
                    name: Some(Name("nikita".into())),
                    server: Server("vlnv.dev".into()),
                    resource: None
                }),
                r#type: MessageType::Chat,
                lang: None,
                subject: Some("daa".into()),
                body: "bbb".into(),
            }
        )
    }

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