lavina/crates/proto-xmpp/src/bind.rs

use std::fmt::Display;

use anyhow::{anyhow, Result};
use quick_xml::events::{BytesEnd, BytesStart, BytesText, Event};
use quick_xml::name::{Namespace, ResolveResult};

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

pub const XMLNS: &'static str = "urn:ietf:params:xml:ns:xmpp-bind";

// TODO remove `pub` in newtypes, introduce validation

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Name(pub Str);

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Server(pub Str);

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Resource(pub Str);

#[derive(PartialEq, Eq, Debug, Clone)]
pub struct Jid {
    pub name: Option<Name>,
    pub server: Server,
    pub resource: Option<Resource>,
}

impl Display for Jid {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(name) = &self.name {
            write!(f, "{}@", &name.0)?;
        }
        write!(f, "{}", &self.server.0)?;
        if let Some(resource) = &self.resource {
            write!(f, "/{}", &resource.0)?;
        }
        Ok(())
    }
}

impl Jid {
    pub fn from_string(i: &str) -> Result<Jid> {
        use regex::Regex;
        use lazy_static::lazy_static;
        lazy_static! {
            static ref RE: Regex = Regex::new(r"^(([a-zA-Z]+)@)?([a-zA-Z.]+)(/([a-zA-Z\-]+))?$").unwrap();
        }
        let m = RE
            .captures(i)
            .ok_or(anyhow!("Incorrectly format jid: {i}"))?;

        let name = m.get(2).map(|name| Name(name.as_str().into()));
        let server = m.get(3).unwrap();
        let server = Server(server.as_str().into());
        let resource = m
            .get(5)
            .map(|resource| Resource(resource.as_str().into()));

        Ok(Jid {
            name,
            server,
            resource,
        })
    }
}

/// Request to bind to a resource.
///
/// Example:
/// ```xml
/// <bind xmlns="urn:ietf:params:xml:ns:xmpp-bind">
///   <resource>mobile</resource>
/// </bind>
/// ```
///
#[derive(PartialEq, Eq, Debug)]
pub struct BindRequest(pub Resource);

pub struct BindRequestParser(BindRequestParserInner);

enum BindRequestParserInner {
    Initial,
    /// Consumed <bind> start and expects <resource>
    InBind(Option<String>),
    /// Consumed <resource> start
    InBindResourceInitial,
    /// Consumer <resource> start and inner text
    InBindResourceEnd(String),
}

impl FromXmlTag for BindRequest {
    const NS: &'static str = XMLNS;
    const NAME: &'static str = "bind";
}

impl FromXml for BindRequest {
    type P = BindRequestParser;

    fn parse() -> Self::P {
        BindRequestParser(BindRequestParserInner::Initial)
    }
}

// TODO rewrite as a generator
impl Parser for BindRequestParser {
    type Output = Result<BindRequest>;

    fn consume<'a>(
        self: Self,
        namespace: ResolveResult,
        event: &Event<'a>,
    ) -> Continuation<Self, Self::Output> {
        // TODO validate tag names and namespaces
        use BindRequestParserInner::*;
        match self.0 {
            Initial => {
                let Event::Start(bytes) = event else {
                    return Continuation::Final(Err(anyhow!("Unexpected XML event: {event:?}")));
                };
                if bytes.name().0 != BindRequest::NAME.as_bytes() {
                    return Continuation::Final(Err(anyhow!(
                        "Unexpected XML tag: {:?}",
                        bytes.name()
                    )));
                }
                let ResolveResult::Bound(Namespace(ns)) = namespace else {
                    return Continuation::Final(Err(anyhow!("No namespace provided")));
                };
                if ns != XMLNS.as_bytes() {
                    return Continuation::Final(Err(anyhow!("Incorrect namespace")));
                }
                Continuation::Continue(BindRequestParser(InBind(None)))
            }
            InBind(resource) => match event {
                Event::Start(bytes) => {
                    Continuation::Continue(BindRequestParser(InBindResourceInitial))
                }
                Event::End(bytes) => {
                    let Some(resource) = resource else {
                            return Continuation::Final(Err(anyhow!("No resource was provided")));
                        };
                    Continuation::Final(Ok(BindRequest(Resource(resource.into()))))
                }
                _ => Continuation::Final(Err(anyhow!("Unexpected XML event: {event:?}"))),
            },
            InBindResourceInitial => {
                let Event::Text(text) = event else {
                    return Continuation::Final(Err(anyhow!("Unexpected XML event: {event:?}")));
                };
                let resource = match std::str::from_utf8(&*text) {
                    Ok(e) => e.to_string(),
                    Err(err) => return Continuation::Final(Err(err.into())),
                };
                Continuation::Continue(BindRequestParser(InBindResourceEnd(resource)))
            }
            InBindResourceEnd(resource) => {
                let Event::End(bytes) = event else {
                    return Continuation::Final(Err(anyhow!("Unexpected XML event: {event:?}")));
                };
                Continuation::Continue(BindRequestParser(InBind(Some(resource))))
            }
        }
    }
}

pub struct BindResponse(pub Jid);

impl ToXml for BindResponse {
    fn serialize(&self, events: &mut Vec<Event<'static>>) {
        events.extend_from_slice(&[
            Event::Start(BytesStart::new(
                r#"bind xmlns="urn:ietf:params:xml:ns:xmpp-bind""#,
            )),
            Event::Start(BytesStart::new(r#"jid"#)),
            Event::Text(BytesText::new(self.0.to_string().as_str()).into_owned()),
            Event::End(BytesEnd::new("jid")),
            Event::End(BytesEnd::new("bind")),
        ]);
    }
}

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

    use super::*;

    #[tokio::test]
    async fn parse_message() {
        let input =
            r#"<bind xmlns="urn:ietf:params:xml:ns:xmpp-bind"><resource>mobile</resource></bind>"#;
        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 = 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, BindRequest(Resource("mobile".into())),)
    }

    #[test]
    fn jid_parse_full() {
        let input = "chelik@server.example/kek";
        let expected = Jid {
            name: Some(Name("chelik".into())),
            server: Server("server.example".into()),
            resource: Some(Resource("kek".into())),
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }

    #[test]
    fn jid_parse_user() {
        let input = "chelik@server.example";
        let expected = Jid {
            name: Some(Name("chelik".into())),
            server: Server("server.example".into()),
            resource: None,
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }

    #[test]
    fn jid_parse_server() {
        let input = "server.example";
        let expected = Jid {
            name: None,
            server: Server("server.example".into()),
            resource: None,
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }

    #[test]
    fn jid_parse_server_resource() {
        let input = "server.example/kek";
        let expected = Jid {
            name: None,
            server: Server("server.example".into()),
            resource: Some(Resource("kek".into())),
        };
        let res = Jid::from_string(input).unwrap();
        assert_eq!(res, expected);
    }
}